# ----------------------------------------------------------------------------
# Copyright (c) 2016-2026, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------
from datetime import timedelta
import os
import networkx as nx
from unittest.mock import MagicMock
import pandas as pd
import unittest
import zipfile
import copy

import pytest

from .._checksum_validator import ChecksumDiff, ValidationCode
from .testing_utilities import DummyArtifacts
from ..archive_parser import (
    ProvNode, Config, _Action, _Citations, _ResultMetadata, ParserResults,
    ArchiveParser, ParserV0, ParserV1, ParserV2, ParserV3, ParserV4, ParserV5,
    ParserV6, ParserV7
)
from ...provenance import MetadataInfo

from rachis.core.testing.util import ReallyEqualMixin


class ParserVxTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.das = DummyArtifacts()
        cls.tempdir = cls.das.tempdir

    @classmethod
    def tearDownClass(cls):
        cls.das.free()

    @pytest.mark.filterwarnings('ignore::UserWarning')
    def test_populate_archive(self):
        for artifact in self.das.all_artifact_versions:
            parser = ArchiveParser.get_parser(artifact.archiver)
            uuid = artifact.uuid
            version = artifact.archive_version

            if version == 0:
                with self.assertWarnsRegex(
                    UserWarning,
                    'Artifact .*prior to provenance'
                ):
                    res = parser.parse_prov(
                        Config(), artifact.archiver)

            else:
                res = parser.parse_prov(Config(), artifact.archiver)
                self.assertIsInstance(res, ParserResults)
                pa_uuids = res.parsed_artifact_uuids
                self.assertIsInstance(pa_uuids, set)
                self.assertIsInstance(next(iter(pa_uuids)), str)
                self.assertIsInstance(res.prov_digraph,
                                      (type(None), nx.DiGraph))
                self.assertIsInstance(res.provenance_is_valid, ValidationCode)

                if version < 5:
                    self.assertIsInstance(res.checksum_diff, type(None))
                else:
                    self.assertIsInstance(res.checksum_diff, ChecksumDiff)

                self.assertIn(uuid, res.prov_digraph)
                self.assertIsInstance(
                    res.prov_digraph.nodes[uuid]['node_data'], ProvNode)

    def test_validate_checksums(self):
        for artifact in self.das.all_artifact_versions:
            parser = ArchiveParser.get_parser(artifact.archiver)
            is_valid, diff = \
                parser._validate_checksums(artifact.archiver)
            if artifact.archive_version < 5:
                self.assertEqual(is_valid, ValidationCode.PREDATES_CHECKSUMS)
                self.assertEqual(diff, None)
            else:
                self.assertEqual(is_valid, ValidationCode.VALID)
                self.assertEqual(diff, ChecksumDiff({}, {}, {}))

    @pytest.mark.filterwarnings('ignore::UserWarning')
    def test_correct_validate_checksums_method_called(self):
        '''
        We want to confirm that parse_prov uses the local _validate_checksums
        even when it calls super().parse_prov() internally
        '''
        for artifact in self.das.all_artifact_versions:
            parser = ArchiveParser.get_parser(artifact.archiver)
            if artifact.archive_version < 5:
                parser._validate_checksums = MagicMock(
                    # return values only here to facilitate normal execution
                    return_value=(ValidationCode.PREDATES_CHECKSUMS, None)
                )
                parser.parse_prov(Config(), artifact.archiver)
                parser._validate_checksums.assert_called_once()
            else:
                parser._validate_checksums = MagicMock(
                    return_value=(
                        ValidationCode.VALID,
                        ChecksumDiff({}, {}, {})
                    )
                )
                parser.parse_prov(Config(), artifact.archiver)
                parser._validate_checksums.assert_called_once()


class ArchiveParserTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.das = DummyArtifacts()
        cls.tempdir = cls.das.tempdir

    @classmethod
    def tearDownClass(cls):
        cls.das.free()

    def test_get_parser(self):
        parsers = [
            ParserV0, ParserV1, ParserV2, ParserV3, ParserV4, ParserV5,
            ParserV6, ParserV7
        ]
        for artifact, parser_version in zip(
            self.das.all_artifact_versions, parsers
        ):
            parser = ArchiveParser.get_parser(artifact.archiver)
            self.assertEqual(type(parser), parser_version)

    def test_artifact_parser_parse_prov(self):
        with self.assertRaisesRegex(NotImplementedError, "Use a subclass"):
            ArchiveParser().parse_prov(Config(), 'doesnotmatter.txt')


class ResultMetadataTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.das = DummyArtifacts()
        cls.tempdir = cls.das.tempdir

        cls.uuid = cls.das.concated_ints.uuid
        md_fp = cls.das.concated_ints.artifact._archiver.path / 'metadata.yaml'
        cls.root_md = _ResultMetadata(md_fp)

    @classmethod
    def tearDownClass(cls):
        cls.das.free()

    def test_smoke(self):
        self.assertEqual(self.root_md.uuid, self.uuid)
        self.assertEqual(self.root_md.type, 'IntSequence1')
        self.assertEqual(self.root_md.format, 'IntSequenceDirectoryFormat')

    def test_repr(self):
        exp = (f'UUID:\t\t{self.uuid}\n'
               'Type:\t\tIntSequence1\n'
               'Data Format:\tIntSequenceDirectoryFormat')
        self.assertEqual(repr(self.root_md), exp)


class ActionTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.das = DummyArtifacts()
        cls.tempdir = cls.das.tempdir

        action_path = os.path.join(
            cls.das.concated_ints_v6.artifact._archiver.path, 'provenance',
            'action', 'action.yaml'
        )
        cls.concat_action = _Action(action_path)

        action_path = os.path.join(
            cls.das.single_int.artifact._archiver.path, 'provenance', 'action',
            'action.yaml'
        )
        cls.import_action = _Action(action_path)

        action_path = os.path.join(
            cls.das.pipeline_viz.artifact._archiver.path, 'provenance',
            'action', 'action.yaml'
        )
        cls.pipeline_action = _Action(action_path)

    @classmethod
    def tearDownClass(cls):
        cls.das.free()

    def test_action_id(self):
        exp = '5035a60e-6f9a-40d4-b412-48ae52255bb5'
        self.assertEqual(self.concat_action.action_id, exp)

    def test_action_type(self):
        self.assertEqual(self.concat_action.action_type, 'method')
        self.assertEqual(self.import_action.action_type, 'import')
        self.assertEqual(self.pipeline_action.action_type, 'pipeline')

    def test_runtime(self):
        exp_t = timedelta
        exp = timedelta(microseconds=6840)
        self.assertIsInstance(self.concat_action.runtime, exp_t)
        self.assertEqual(self.concat_action.runtime, exp)

    def test_runtime_str(self):
        exp = '6840 microseconds'
        self.assertEqual(self.concat_action.runtime_str, exp)

    def test_action(self):
        exp = 'concatenate_ints'
        self.assertEqual(self.concat_action.action_name, exp)

    def test_plugin(self):
        exp = 'dummy_plugin'
        self.assertEqual(self.concat_action.plugin, exp)

    '''
    Import is not handled by a plugin, so the parser provides values
    for the action_name and plugin properties not present in action.yaml
    '''

    def test_action_for_import_node(self):
        exp = 'import'
        self.assertEqual(self.import_action.action_name, exp)

    def test_plugin_for_import_node(self):
        exp = 'framework'
        self.assertEqual(self.import_action.plugin, exp)

    def test_inputs(self):
        exp = {
            'ints1': '8dea2f1a-2164-4a85-9f7d-e0641b1db22b',
            'ints2': '8dea2f1a-2164-4a85-9f7d-e0641b1db22b',
            'ints3': '7727c060-5384-445d-b007-b64b41a090ee'
        }
        self.assertEqual(self.concat_action.inputs, exp)

        exp = {}
        self.assertEqual(self.import_action.inputs, exp)

    def test_parameters(self):
        exp = {
            'int1': 7,
            'int2': 100,
        }
        self.assertEqual(self.concat_action.parameters, exp)

        exp = {}
        self.assertEqual(self.import_action.parameters, exp)

    def test_output_name(self):
        exp = 'concatenated_ints'
        self.assertEqual(self.concat_action.output_name, exp)

        exp = None
        self.assertEqual(self.import_action.output_name, exp)

    def test_format(self):
        exp = None
        self.assertEqual(self.concat_action.format, exp)
        self.assertEqual(self.import_action.format, exp)

    def test_transformers(self):
        int_seq_dir_citation = (
            'view|dummy-plugin:0.0.0-dev|IntSequenceDirectoryFormat|0'
        )
        transformer_citation = (
            'transformer|dummy-plugin:0.0.0-dev|builtins:list'
            '->IntSequenceDirectoryFormat|0'
        )
        output_citations = [transformer_citation, int_seq_dir_citation]

        exp = {
            'inputs': {
                'ints1':
                    [{
                        'from': 'IntSequenceDirectoryFormat',
                        'to': 'builtins:list',
                        'plugin': 'dummy-plugin',
                        'citations': [int_seq_dir_citation]
                    }],
                'ints2':
                    [{
                        'from': 'IntSequenceDirectoryFormat',
                        'to': 'builtins:list',
                        'plugin': 'dummy-plugin',
                        'citations': [int_seq_dir_citation]
                    }],
                'ints3':
                    [{
                        'from': 'IntSequenceV2DirectoryFormat',
                        'to': 'builtins:list',
                        'plugin': 'dummy-plugin',
                    }],
            },
            'output': [{
                'from': 'builtins:list',
                'to': 'IntSequenceDirectoryFormat',
                'plugin': 'dummy-plugin',
                'citations': output_citations
            }]
        }
        self.assertEqual(self.concat_action.transformers, exp)

    def test_repr(self):
        exp = (
            '_Action(action_id=5035a60e-6f9a-40d4-b412-48ae52255bb5, '
            'type=method, plugin=dummy_plugin, '
            'action=concatenate_ints)'
        )
        self.assertEqual(repr(self.concat_action), exp)


class CitationsTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.das = DummyArtifacts()
        cls.tempdir = cls.das.tempdir

        cls.cite_none = os.path.join(cls.das.datadir, 'cite_none.bib')
        cls.cite_one = os.path.join(cls.das.datadir, 'cite_one.bib')
        cls.cite_many = os.path.join(cls.das.datadir, 'cite_many.bib')

    @classmethod
    def tearDownClass(cls):
        cls.das.free()

    def test_empty_bib(self):
        citations = _Citations(self.cite_none)
        self.assertEqual(len(citations.citations), 0)

    def test_citation(self):
        exp = 'framework'
        citations = _Citations(self.cite_one)
        for key in citations.citations:
            self.assertRegex(key, exp)

    def test_many_citations(self):
        exp = ['2020.6.0.dev0', 'unweighted_unifrac.+0',
               'unweighted_unifrac.+1', 'unweighted_unifrac.+2',
               'unweighted_unifrac.+3', 'unweighted_unifrac.+4',
               'BIOMV210DirFmt', 'BIOMV210Format']
        citations = _Citations(self.cite_many)
        for i, key in enumerate(citations.citations):
            self.assertRegex(key, exp[i])

    def test_repr(self):
        exp = ("Citations(['framework|qiime2:2020.6.0.dev0|0'])")
        citations = _Citations(self.cite_one)
        self.assertEqual(repr(citations), exp)


class ProvNodeTests(unittest.TestCase, ReallyEqualMixin):
    @classmethod
    def setUpClass(cls):
        cls.das = DummyArtifacts()
        cls.tempdir = cls.das.tempdir

        # build root nodes for all archive format versions
        cfg = Config(parse_study_metadata=True)
        cls.nodes = {}
        for artifact in cls.das.all_artifact_versions:
            cls.nodes[str(artifact.archive_version)] = \
                ProvNode(cfg, artifact.archiver)

        with zipfile.ZipFile(cls.das.concated_ints_with_md.filepath) as zf:
            root_node_id = cls.das.concated_ints_with_md.uuid
            all_filenames = zf.namelist()

        dag = cls.das.concated_ints_with_md.dag
        for node in dag.nodes:
            md_path = os.path.join(
                root_node_id, 'provenance', 'artifacts', node, 'action',
                'metadata.tsv'
            )
            if md_path in all_filenames:
                md_node_id = node
            else:
                non_md_node_id = node

        # build a nonroot node without study metadata
        cls.nonroot_non_md_node = ProvNode(
            cfg, cls.das.concated_ints_with_md.artifact._archiver,
            uuid=non_md_node_id
        )

        # build a nonroot node with study metadata
        cls.nonroot_md_node = ProvNode(
            cfg, cls.das.concated_ints_with_md.artifact._archiver,
            uuid=md_node_id
        )

        # build a root node and parse study metadata files
        cfg = Config(parse_study_metadata=True)
        cls.root_node_parse_md = ProvNode(
            cfg, cls.das.concated_ints_with_md.artifact._archiver
        )

        # build a root node and don't parse study metadata files
        cfg = Config(parse_study_metadata=False)
        cls.root_node_dont_parse_md = ProvNode(
            cfg, cls.das.concated_ints_with_md.artifact._archiver
        )

        # build a node with a collection as input
        cls.input_collection_node = ProvNode(
            cfg, cls.das.int_from_collection.artifact._archiver
        )

        # build a node with an optional input that defaults to None
        cls.optional_input_node = ProvNode(
            cfg, cls.das.int_seq_optional_input.artifact._archiver
        )

    @classmethod
    def tearDownClass(cls):
        cls.das.free()

    def test_smoke(self):
        self.assertTrue(True)
        for node_vzn in self.nodes:
            self.assertIsInstance(self.nodes[node_vzn], ProvNode)

    def test_node_properties(self):
        # hardcoded from test data
        framework_versions = {
            '0': '2.0.5',
            '1': '2017.2.0',
            '2': '2017.9.0',
            '3': '2018.2.0',
            '4': '2018.6.0',
            '5': '2018.11.0',
            '6': '2023.5.1',
        }

        for node, archive_version in zip(
            self.nodes, [str(i) for i in range(7)]
        ):
            if archive_version == '0':
                self.assertEqual(self.nodes[node].format, 'BIOMV210DirFmt')
                self.assertEqual(self.nodes[node].type,
                                 'FeatureTable[Frequency]')
            else:
                self.assertEqual(
                    self.nodes[node].format, 'IntSequenceDirectoryFormat'
                )
                self.assertEqual(self.nodes[node].type, 'IntSequence1')

            if archive_version == '0' or archive_version == '1':
                self.assertEqual(self.nodes[node].has_provenance, False)
            else:
                self.assertEqual(self.nodes[node].has_provenance, True)

            self.assertEqual(self.nodes[node].archive_version, archive_version)
            self.assertEqual(
                self.nodes[node].framework_version,
                framework_versions[archive_version]
            )

    def test_self_eq(self):
        self.assertReallyEqual(self.nodes['5'], self.nodes['5'])

    def test_eq(self):
        # Mock has no matching UUID
        mock_node = MagicMock()
        self.assertNotEqual(self.nodes['5'], mock_node)

        # Mock has bad UUID
        mock_node._uuid = 'gerbil'
        self.assertReallyNotEqual(self.nodes['5'], mock_node)

        # Matching UUIDs insufficient if classes differ
        mock_node._uuid = self.das.concated_ints_v5.uuid
        self.assertReallyNotEqual(self.nodes['5'], mock_node)
        mock_node.__class__ = ProvNode
        self.assertReallyEqual(self.nodes['5'], mock_node)

    def test_is_hashable(self):
        exp_hash = hash(self.das.concated_ints_v5.uuid)
        self.assertReallyEqual(hash(self.nodes['5']), exp_hash)

    def test_str(self):
        for node_vzn, artifact in zip(
            self.nodes, self.das.all_artifact_versions
        ):
            uuid = artifact.uuid
            self.assertRegex(str(self.nodes[node_vzn]),
                             f'(?s)UUID:\t\t{uuid}.*Type.*Data Format')

    def test_repr(self):
        for node_vzn, artifact in zip(
            self.nodes, self.das.all_artifact_versions
        ):
            uuid = artifact.uuid
            self.assertRegex(repr(self.nodes[node_vzn]),
                             f'(?s)UUID:\t\t{uuid}.*Type.*Data Format')

    def test_get_metadata_from_action(self):
        new_node = copy.deepcopy(self.root_node_parse_md)
        find_md = new_node._get_metadata_from_Action

        # create dummy hash '0', not relevant here
        md = MetadataInfo(
            ['d5b4cf78-f5e2-44e0-aa24-66b02564e9f1'], 'metadata.tsv', '0'
        )
        action_details = {
            'parameters': [{'metadata': md}]
        }
        new_node.action._action_details = action_details

        all_md, artifacts_as_md = find_md()
        all_exp = {'metadata': 'metadata.tsv'}
        a_as_md_exp = [{
            'artifact_passed_as_metadata':
            'd5b4cf78-f5e2-44e0-aa24-66b02564e9f1'
        }]
        self.assertEqual(all_md, all_exp)
        self.assertEqual(artifacts_as_md, a_as_md_exp)

    def test_get_metadata_from_action_with_no_params(self):
        new_node = copy.deepcopy(self.nodes['5'])
        find_md = new_node._get_metadata_from_Action
        action_details = {'parameters': []}
        new_node.action._action_details = action_details

        all_md, artifacts_as_md = find_md()
        self.assertEqual(all_md, {})
        self.assertEqual(artifacts_as_md, [])

        action_details = {'non-parameters-key': 'here is a thing'}
        new_node.action._action_details = action_details

        all_md, artifacts_as_md = find_md()
        self.assertEqual(all_md, {})
        self.assertEqual(artifacts_as_md, [])

    def test_metadata_available_in_property(self):
        self.assertEqual(type(self.nonroot_md_node.metadata), dict)
        self.assertIn('metadata', self.nonroot_md_node.metadata)
        self.assertEqual(type(self.nonroot_md_node.metadata['metadata']),
                         pd.DataFrame)

    def test_metadata_not_available_in_property_w_opt_out(self):
        self.assertEqual(self.root_node_dont_parse_md.metadata, None)

    def test_metadata_is_correct(self):
        self.assertIn('metadata', self.nonroot_md_node.metadata)

        md_data = {
            'id': ['#q2:types', '0'],
            'a': ['categorical', '42'],
        }
        md_exp = pd.DataFrame(md_data, columns=md_data.keys())
        pd.testing.assert_frame_equal(
            md_exp, self.nonroot_md_node.metadata['metadata']
        )

    def test_has_no_provenance_so_no_metadata(self):
        self.assertEqual(self.nodes['0'].has_provenance, False)
        self.assertEqual(self.nodes['0'].metadata, None)

    def test_node_has_provenance_but_no_metadata(self):
        self.assertEqual(self.nonroot_non_md_node.has_provenance, True)
        self.assertEqual(self.nonroot_non_md_node.metadata, {})

    def test_parse_metadata_for_nonroot_node(self):
        self.assertEqual(self.nonroot_md_node.has_provenance, True)
        self.assertIn('metadata', self.nonroot_md_node.metadata)

    def test_parents(self):
        actual_parent_names = []
        for parent in self.nodes['5']._parents:
            actual_parent_names += parent.keys()

        self.assertIn('ints1', actual_parent_names)
        self.assertIn('ints2', actual_parent_names)
        self.assertIn('ints3', actual_parent_names)

        self.assertEqual(len(self.nodes['5']._parents), 3)
        self.assertEqual(len(actual_parent_names), 3)

    def test_parents_no_prov(self):
        no_prov_node = self.nodes['0']
        self.assertFalse(no_prov_node.has_provenance)
        self.assertEqual(no_prov_node._parents, None)

    def test_parents_with_artifact_passed_as_md(self):
        actual_parent_names = []
        for parent in self.nonroot_md_node._parents:
            actual_parent_names += parent.keys()

        self.assertIn('ints', actual_parent_names)
        self.assertIn('artifact_passed_as_metadata', actual_parent_names)

        self.assertEqual(len(self.nonroot_md_node._parents), 2)
        self.assertEqual(len(actual_parent_names), 2)

    def test_parents_for_import_node(self):
        import_node = ProvNode(
            Config(), self.das.single_int.artifact._archiver
        )
        self.assertEqual(import_node._parents, [])

    def test_parents_collection_of_inputs(self):
        parents = self.input_collection_node._parents
        self.assertIn('int1', parents[0].keys())
        self.assertIn('int2', parents[1].keys())
        self.assertEqual(len(parents), 2)

    def test_parents_optional_input(self):
        # NOTE: The None-type input is not captured
        parents = self.optional_input_node._parents
        self.assertIn('ints', parents[0].keys())
        self.assertEqual(len(parents), 1)
