#########################################################################
# MacSyLib - Python library to detect macromolecular systems            #
#            in prokaryotes protein dataset using systems modelling     #
#            and similarity search.                                     #
#                                                                       #
# Authors: Sophie Abby, Bertrand Neron                                  #
# Copyright (c) 2014-2025  Institut Pasteur (Paris) and CNRS.           #
# See the COPYRIGHT file for details                                    #
#                                                                       #
# This file is part of MacSyLib package.                                #
#                                                                       #
# MacSyLib is free software: you can redistribute it and/or modify      #
# it under the terms of the GNU General Public License as published by  #
# the Free Software Foundation, either version 3 of the License, or     #
# (at your option) any later version.                                   #
#                                                                       #
# MacSyLib is distributed in the hope that it will be useful,           #
# but WITHOUT ANY WARRANTY; without even the implied warranty of        #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          #
# GNU General Public License for more details .                         #
#                                                                       #
# You should have received a copy of the GNU General Public License     #
# along with MacSyLib (COPYING).                                        #
# If not, see <https://www.gnu.org/licenses/>.                          #
#########################################################################

import os
import tempfile
import unittest
import urllib.request
import urllib.error
import json
import io
import shutil
import tarfile
import glob

try:
    import lxml
except ModuleNotFoundError:
    lxml = None

import yaml
import colorlog
from unittest.mock import patch

import macsylib
from macsylib import model_package
from macsylib.metadata import Maintainer
from macsylib import model_conf_parser
from macsylib.error import MacsydataError, MacsyDataLimitError

from tests import MacsyTest


class TestPackageFunc(MacsyTest):

    def test_parse_arch_path(self):
        self.assertTupleEqual(model_package.parse_arch_path("pack-3.0.tar.gz"),
                              ('pack', '3.0'))
        self.assertTupleEqual(model_package.parse_arch_path("pack-3.0.tgz"),
                              ('pack', '3.0'))

        pack = "pack-3.0.foo"
        with self.assertRaises(ValueError) as ctx:
            model_package.parse_arch_path(pack)
        self.assertEqual(str(ctx.exception),
                         f"{pack} does not seem to be a package (a tarball).")
        pack = "pack.tar.gz"
        with self.assertRaises(ValueError) as ctx:
            model_package.parse_arch_path(pack)
        self.assertEqual(str(ctx.exception),
                         f"{pack} does not seem to not be versioned.")


class TestLocalModelIndex(MacsyTest):

    def test_init(self):
        lmi = model_package.LocalModelIndex()
        self.assertEqual(lmi.org_name, 'local')
        expected_cache = os.path.join(tempfile.gettempdir(), 'tmp-macsy-cache')
        self.assertEqual(lmi.cache, expected_cache)

    def test_repos_url(self):
        lmi = model_package.LocalModelIndex()
        self.assertEqual(lmi.repos_url, 'local')


class TestRemoteModelIndex(MacsyTest):

    def setUp(self) -> None:
        self._tmp_dir = tempfile.TemporaryDirectory(prefix='test_msl_package_')
        self.tmpdir = self._tmp_dir.name


    def tearDown(self) -> None:
        self._tmp_dir.cleanup()


    def mocked_requests_get(url: str, context:None=None):
        # cannot type the return value the class is defined inside de method

        class MockResponse:
            def __init__(self, data, status_code):
                self.data = io.BytesIO(bytes(data.encode("utf-8")))
                self.status_code = status_code

            def read(self, length=-1):
                return self.data.read(length)

            def __enter__(self):
                return self

            def __exit__(self, type, value, traceback):
                return False

        if url == 'https://test_url_json/':
            resp = {'fake': ['json', 'response']}
            return MockResponse(json.dumps(resp), 200)
        elif url == 'https://test_url_json/limit':
            raise urllib.error.HTTPError(url, 403, 'forbidden', None, None)
        elif url == 'https://api.github.com/orgs/remote_exists_true':
            resp = {'type': 'Organization'}
            return MockResponse(json.dumps(resp), 200)
        elif url == 'https://api.github.com/orgs/remote_exists_false':
            raise urllib.error.HTTPError(url, 404, 'not found', None, None)
        elif url == 'https://api.github.com/orgs/remote_exists_server_error':
            raise urllib.error.HTTPError(url, 500, 'Server Error', None, None)
        elif url == 'https://api.github.com/orgs/remote_exists_unexpected_error':
            raise urllib.error.HTTPError(url, 204, 'No Content', None, None)
        elif url == 'https://api.github.com/orgs/list_packages/repos':
            resp = [{'name': 'model_1'}, {'name': 'model_2'}, {'name':'.github'}]
            return MockResponse(json.dumps(resp), 200)
        elif url == 'https://api.github.com/repos/list_package_vers/model_1/tags':
            resp = [{'name': 'v_1'}, {'name': 'v_2'}]
            return MockResponse(json.dumps(resp), 200)
        elif url == 'https://api.github.com/repos/list_package_vers/model_2/tags':
            raise urllib.error.HTTPError(url, 404, 'not found', None, None)
        elif url == 'https://api.github.com/repos/list_package_vers/model_3/tags':
            raise urllib.error.HTTPError(url, 500, 'Server Error', None, None)
        elif 'https://api.github.com/repos/package_download/fake/tarball/1.0' in url:
            return MockResponse('fake data ' * 2, 200)
        elif url == 'https://api.github.com/repos/package_download/bad_pack/tarball/0.2':
            raise urllib.error.HTTPError(url, 404, 'not found', None, None)
        elif url == 'https://raw.githubusercontent.com/get_metadata/foo/0.0/metadata.yml':
            data = yaml.dump({"maintainer": {"name": "moi"}})
            return MockResponse(data, 200)
        else:
            raise RuntimeError("test non prevu", url)


    def test_init(self):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        model_package.RemoteModelIndex.remote_exists = lambda x: True
        try:
            remote = model_package.RemoteModelIndex()
            remote.cache = self.tmpdir
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists
        self.assertEqual(remote.org_name, 'macsy-models')
        self.assertEqual(remote.base_url, 'https://api.github.com')
        self.assertEqual(remote.cache, self.tmpdir)

        model_package.RemoteModelIndex.remote_exists = lambda x: True
        try:
            remote = model_package.RemoteModelIndex(org='foo')
            remote.cache = self.tmpdir
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists
        self.assertEqual(remote.org_name, 'foo')

        model_package.RemoteModelIndex.remote_exists = lambda x: False
        try:
            with self.assertRaises(ValueError) as ctx:
                model_package.RemoteModelIndex(org='foo')
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists
        self.assertEqual(str(ctx.exception), "the 'foo' organization does not exist.")


    def test_repos_url(self):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        model_package.RemoteModelIndex.remote_exists = lambda x: True
        try:
            org = "package_repos_url"
            remote = model_package.RemoteModelIndex(org=org)
            self.assertEqual(remote.repos_url, f"https://github.com/{org}")
        finally:
            rem_exists

    @patch('urllib.request.urlopen', side_effect=mocked_requests_get)
    def test_url_json(self, mock_urlopen):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        model_package.RemoteModelIndex.remote_exists = lambda x: True
        remote = model_package.RemoteModelIndex(org="nimportnaoik")
        remote.cache = self.tmpdir
        try:
            j = remote._url_json("https://test_url_json/")
            self.assertDictEqual(j, {'fake': ['json', 'response']})
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists


    @patch('urllib.request.urlopen', side_effect=mocked_requests_get)
    def test_url_json_reach_limit(self, mock_urlopen):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        model_package.RemoteModelIndex.remote_exists = lambda x: True
        remote = model_package.RemoteModelIndex(org="nimportnaoik")
        remote.cache = self.tmpdir
        try:
            with self.assertRaises(MacsyDataLimitError) as ctx:
                remote._url_json("https://test_url_json/limit")
            self.assertEqual(str(ctx.exception),
                             """You reach the maximum number of request per hour to github.
Please wait before to try again.""")
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists


    @patch('urllib.request.urlopen', side_effect=mocked_requests_get)
    def test_remote_exists(self, mock_urlopen):
        remote = model_package.RemoteModelIndex(org="remote_exists_true")
        remote.cache = self.tmpdir
        exists = remote.remote_exists()
        self.assertTrue(exists)

        remote.org_name = "remote_exists_false"
        exists = remote.remote_exists()
        self.assertFalse(exists)

        remote.org_name = "remote_exists_server_error"
        with self.assertRaises(urllib.error.HTTPError) as ctx:
            remote.remote_exists()
        ctx.exception.close()
        self.assertEqual(str(ctx.exception),
                         "HTTP Error 500: Server Error")

        remote.org_name = "remote_exists_unexpected_error"
        with self.assertRaises(urllib.error.HTTPError) as ctx:
            remote.remote_exists()
        ctx.exception.close()
        self.assertEqual(str(ctx.exception),
                         "HTTP Error 204: No Content")

    @patch('urllib.request.urlopen', side_effect=mocked_requests_get)
    def test_list_packages(self, mock_urlopen):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        try:
            model_package.RemoteModelIndex.remote_exists = lambda x: True
            remote = model_package.RemoteModelIndex(org="list_packages")
            remote.cache = self.tmpdir
            self.assertListEqual(remote.list_packages(), ['model_1', 'model_2'])
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists


    @patch('urllib.request.urlopen', side_effect=mocked_requests_get)
    def test_list_package_vers(self, mock_urlopen):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        try:
            model_package.RemoteModelIndex.remote_exists = lambda x: True
            remote = model_package.RemoteModelIndex(org="list_package_vers")
            remote.cache = self.tmpdir
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists

        self.assertListEqual(remote.list_package_vers('model_1'), ['v_1', 'v_2'])

        with self.assertRaises(ValueError) as ctx:
            _ = remote.list_package_vers('model_2')
        self.assertEqual(str(ctx.exception), "package 'model_2' does not exists on repos 'list_package_vers'")

        with self.assertRaises(urllib.error.HTTPError) as ctx:
            _ = remote.list_package_vers('model_3')
        ctx.exception.close()
        self.assertEqual(str(ctx.exception), "HTTP Error 500: Server Error")


    @patch('urllib.request.urlopen', side_effect=mocked_requests_get)
    def test_download(self, mock_urlopen):
        rem_exists = model_package.RemoteModelIndex.remote_exists
        try:
            model_package.RemoteModelIndex.remote_exists = lambda x: True
            remote = model_package.RemoteModelIndex(org="package_download")
            remote.cache = self.tmpdir
            pack_name = "fake"
            pack_vers = "1.0"
            # ensure that remote.cache does not exists
            if os.path.exists(remote.cache):
                if os.path.isdir(remote.cache):
                    shutil.rmtree(remote.cache)
                elif os.path.isfile(remote.cache) or os.path.islink(remote.cache):
                    os.unlink(remote.cache)

            arch_path = remote.download(pack_name, pack_vers)
            self.assertEqual(os.path.join(remote.cache, remote.org_name, f"{pack_name}-{pack_vers}.tar.gz"),
                             arch_path)
            self.assertFileEqual(arch_path, io.StringIO('fake data ' * 2))

            # download again with existing remote.cache and replace old archive
            os.unlink(arch_path)
            arch_path = remote.download(pack_name, pack_vers)
            self.assertFileEqual(arch_path, io.StringIO('fake data ' * 2))

            # download again with existing remote.cache and replace old archive
            os.unlink(arch_path)
            dest = os.path.join(self.tmpdir, 'dest')
            os.makedirs(dest)
            arch_path = remote.download(pack_name, pack_vers, dest=dest)
            self.assertEqual(os.path.join(dest, f'{pack_name}-{pack_vers}.tar.gz'), arch_path)

            # remote cache exist and is a file
            shutil.rmtree(remote.cache)
            open(remote.cache, 'w').close()
            try:
                with self.assertRaises(NotADirectoryError) as ctx:
                    remote.download(pack_name, pack_vers)
                self.assertEqual(str(ctx.exception),
                                 f"The tmp cache '{remote.cache}' already exists")
            finally:
                os.unlink(remote.cache)

            with self.assertRaises(ValueError) as ctx:
                _ = remote.download("bad_pack", "0.2")
            self.assertEqual(str(ctx.exception),
                             "package 'bad_pack-0.2' does not exists on repos 'package_download'")

        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists


    def test_unarchive(self):

        def create_pack(dir_, repo, name, vers, key):
            pack_name = f"{name}-{vers}"
            tar_path = os.path.join(dir_, f"{pack_name}.tar.gz")
            with tarfile.open(tar_path, "w:gz") as tar:
                tmp_pack = os.path.join(dir_, f"{repo}-{name}-{key}")
                os.mkdir(tmp_pack)
                for i in range(3):
                    name = f"file_{i}"
                    tmp_file = os.path.join(tmp_pack, name)
                    with open(tmp_file, 'w') as f:
                        f.write(f"Content of file {i}\n")
                tar.add(tmp_pack, arcname=os.path.basename(tmp_pack))
            shutil.rmtree(tmp_pack)
            return tar_path

        pack_name = 'model-toto'
        pack_vers = '2.0'

        rem_exists = model_package.RemoteModelIndex.remote_exists
        model_package.RemoteModelIndex.remote_exists = lambda x: True
        try:
            remote = model_package.RemoteModelIndex(org="package_unarchive")
            arch = create_pack(self.tmpdir, remote.org_name, pack_name, pack_vers, 'e020300')
            remote.cache = self.tmpdir

            model_path = remote.unarchive_package(arch)
            unpacked_path = os.path.join(self.tmpdir, remote.org_name, pack_name, pack_vers, pack_name)
            self.assertEqual(model_path, unpacked_path)
            self.assertTrue(os.path.exists(unpacked_path))
            self.assertTrue(os.path.isdir(unpacked_path))
            self.assertListEqual(sorted(glob.glob(f"{unpacked_path}/*")),
                                 sorted([os.path.join(unpacked_path, f"file_{i}") for i in range(3)])
                                 )
            # test package is remove before to unarchive a new one
            open(os.path.join(unpacked_path, "file_must_be_removed"), 'w').close()
            model_path = remote.unarchive_package(arch)
            self.assertListEqual(sorted(glob.glob(f"{unpacked_path}/*")),
                                 sorted([os.path.join(unpacked_path, f"file_{i}") for i in range(3)])
                                 )
        finally:
            model_package.RemoteModelIndex.remote_exists = rem_exists


class TestModelPackage(MacsyTest):

    def setUp(self) -> None:
        self._tmp_dir = tempfile.TemporaryDirectory(prefix='test_msl_package_')
        self.tmpdir = self._tmp_dir.name

        macsylib.init_logger()
        macsylib.logger_set_level(level=30)
        logger = colorlog.getLogger('macsylib.package')
        model_package._log = logger
        logger = colorlog.getLogger('macsylib.model_conf_parser')
        model_conf_parser._log = logger
        maintainer = Maintainer("auth_name", "auth_name@mondomain.fr")
        self.metadata = model_package.Metadata(maintainer,
                                         "this is a short description of the repos")
        self.metadata.vers = "0.0b2"
        self.metadata.cite = ["bla bla",
                                  "link to publication",
                                  """ligne 1
ligne 2
ligne 3 et bbbbb
"""]
        self.metadata.doc = "http://link/to/the/documentation"
        self.metadata.license = "CC BY-NC-SA 4.0 (https://creativecommons.org/licenses/by-nc-sa/4.0/)"
        self.metadata.copyright_date = "2019"
        self.metadata.copyright_holder = "Institut Pasteur, CNRS"


    def tearDown(self) -> None:
        self._tmp_dir.cleanup()


    def create_fake_package(self, model,
                            definitions=True,
                            bad_definitions=False,
                            profiles=True,
                            skip_hmm=None,
                            metadata='good_metadata.yml',
                            readme=True,
                            license=True,
                            conf=True,
                            vers=True,
                            grammar="2.1",
                            bad_conf=False):
        pack_path = os.path.join(self.tmpdir, model)
        os.mkdir(pack_path)
        if definitions:
            def_dir = os.path.join(pack_path, 'definitions')
            os.mkdir(def_dir)
            with open(os.path.join(def_dir, "model_1.xml"), 'w') as f:
                f.write(f"""<model inter_gene_max_space="20" min_mandatory_genes_required="1" min_genes_required="2" vers="{grammar}">
    <gene name="flgB" presence="mandatory"/>
    <gene name="flgC" presence="mandatory" inter_gene_max_space="2"/>
</model>""")
            with open(os.path.join(def_dir, "model_2.xml"), 'w') as f:
                f.write(f"""<model inter_gene_max_space="20" min_mandatory_genes_required="1" min_genes_required="2" vers="{grammar}">
    <gene name="fliE" presence="mandatory" multi_system="true"/>
    <gene name="tadZ" presence="accessory" loner="true"/>
    <gene name="sctC" presence="forbidden"/>
</model>""")
        if bad_definitions:
            with open(os.path.join(def_dir, "model_3.xml"), 'w') as f:
                f.write(f"""<model inter_gene_max_space="20" min_mandatory_genes_required="2" min_genes_required="1" vers="{grammar}">
    <gene name="flgB" presence="mandatory"/>
    <gene name="flgC" presence="mandatory" inter_gene_max_space="2"/>
</model>""")
        if profiles:
            profile_dir = os.path.join(pack_path, 'profiles')
            os.mkdir(profile_dir)
            for name in ('flgB', 'flgC', 'fliE', 'tadZ', 'sctC'):
                if skip_hmm and name in skip_hmm:
                    continue
                shutil.copyfile(self.find_data('models', 'foo', 'profiles', f'{name}.hmm'),
                                os.path.join(profile_dir, f"{name}.hmm")
                                )
        if metadata:
            meta_path = self.find_data('pack_metadata', metadata)
            meta_dest = os.path.join(pack_path, model_package.Metadata.name)
            with open(meta_path) as meta_file:
                meta = yaml.safe_load(meta_file)
            if not vers:
                meta['vers'] = None
            with open(meta_dest, 'w') as meta_file:
                yaml.dump(meta, meta_file,allow_unicode=True, indent=2)
        if readme:
            with open(os.path.join(pack_path, "README"), 'w') as f:
                f.write("# This a README\n")
        if license:
            with open(os.path.join(pack_path, "LICENSE"), 'w') as f:
                f.write("# This the License\n")
        if conf:
            with open(os.path.join(pack_path, "model_conf.xml"), 'w') as f:
                conf = """<model_config>
    <weights>
        <itself>11</itself>
        <exchangeable>12</exchangeable>
        <mandatory>13</mandatory>
        <accessory>14</accessory>
        <neutral>0</neutral>
        <out_of_cluster>10</out_of_cluster>
    </weights>
    <filtering>
        <e_value_search>0.12</e_value_search>
        <i_evalue_sel>0.012</i_evalue_sel>
        <coverage_profile>0.55</coverage_profile>
        <cut_ga>false</cut_ga>
    </filtering>
</model_config>
"""
                f.write(conf)
        elif bad_conf:
            with open(os.path.join(pack_path, "model_conf.xml"), 'w') as f:
                conf = """<model_config>
    <weights>
        <itself>FOO</itself>
        <exchangeable>BAR</exchangeable>
    </weights>
</model_config>
"""
                f.write(conf)

        return pack_path

    def test_check_structure_bad_path(self):
        foobar = os.path.join(self.tmpdir, "foobar")
        pack = model_package.ModelPackage(foobar)
        errors, warnings = pack._check_structure()
        self.assertListEqual(errors, ["The model package 'foobar' does not exists."])
        self.assertListEqual(warnings, [])

        open(foobar, 'w').close()
        errors, warnings = pack._check_structure()
        self.assertListEqual(errors, ["The model package 'foobar' is not a directory "])
        self.assertListEqual(warnings, [])