# -----------------------------------------------------------------------------
# Copyright (c) 2014--, The Qiita Development Team.
#
# Distributed under the terms of the BSD 3-clause License.
#
# The full license is in the file LICENSE, distributed with this software.
# -----------------------------------------------------------------------------

from os import remove, close, mkdir
from os.path import exists, join
from tempfile import mkstemp, mkdtemp
from shutil import rmtree
from unittest import TestCase, main
from six import StringIO
from functools import partial

import pandas as pd

from qiita_core.util import qiita_test_checker

import qiita_db as qdb

import configparser


@qiita_test_checker()
class TestMakeStudyFromCmd(TestCase):
    def setUp(self):
        qdb.study.StudyPerson.create(
            'SomeDude', 'somedude@foo.bar', 'some',
            '111 fake street', '111-121-1313')
        qdb.user.User.create('test@test.com', 'password')
        self.config1 = CONFIG_1
        self.config2 = CONFIG_2

    def test_make_study_from_cmd(self):
        fh = StringIO(self.config1)
        qdb.commands.load_study_from_cmd('test@test.com', 'newstudy', fh)

        with qdb.sql_connection.TRN:
            sql = """SELECT study_id
                     FROM qiita.study
                     WHERE email = %s AND study_title = %s"""
            qdb.sql_connection.TRN.add(sql, ['test@test.com', 'newstudy'])
            study_id = qdb.sql_connection.TRN.execute_fetchflatten()
        self.assertEqual(study_id, [2])

        fh2 = StringIO(self.config2)
        with self.assertRaises(configparser.NoOptionError):
            qdb.commands.load_study_from_cmd('test@test.com', 'newstudy2', fh2)


@qiita_test_checker()
class TestLoadArtifactFromCmd(TestCase):
    def setUp(self):
        self.artifact_count = qdb.util.get_count('qiita.artifact')
        self.fp_count = qdb.util.get_count('qiita.filepath')
        self.files_to_remove = []

    def tearDown(self):
        for fp in self.files_to_remove:
            if exists(fp):
                remove(fp)

    def test_load_artifact_from_cmd_error(self):
        with self.assertRaises(ValueError):
            qdb.commands.load_artifact_from_cmd(
                ["fp1", "fp2"], ["preprocessed_fasta"], "Demultiplexed",
                parents=[1], dflt_params_id=10,
                required_params='{"input_data": 1}')

        with self.assertRaises(ValueError):
            qdb.commands.load_artifact_from_cmd(
                ["fp1"], ["preprocessed_fasta"], "Demultiplexed",
                parents=[1, 2], dflt_params_id=10)

    def test_load_artifact_from_cmd_root(self):
        fd, forward_fp = mkstemp(suffix='_forward.fastq.gz')
        close(fd)
        self.files_to_remove.append(forward_fp)
        fd, reverse_fp = mkstemp(suffix='_reverse.fastq.gz')
        close(fd)
        self.files_to_remove.append(reverse_fp)
        fd, barcodes_fp = mkstemp(suffix='_barcodes.fastq.gz')
        close(fd)
        self.files_to_remove.append(barcodes_fp)
        fps = [forward_fp, reverse_fp, barcodes_fp]
        for fp in fps:
            with open(fp, 'w') as f:
                f.write('\n')
        ftypes = ['raw_forward_seqs', 'raw_reverse_seqs', 'raw_barcodes']
        metadata = pd.DataFrame.from_dict(
            {'SKB8.640193': {'center_name': 'ANL',
                             'primer': 'GTGCCAGCMGCCGCGGTAA',
                             'barcode': 'GTCCGCAAGTTA',
                             'run_prefix': "s_G1_L001_sequences",
                             'platform': 'Illumina',
                             'instrument_model': 'Illumina MiSeq',
                             'library_construction_protocol': 'AAAA',
                             'experiment_design_description': 'BBBB'}},
            orient='index', dtype=str)
        pt = qdb.metadata_template.prep_template.PrepTemplate.create(
            metadata, qdb.study.Study(1), "16S")
        obs = qdb.commands.load_artifact_from_cmd(
            fps, ftypes, 'FASTQ', prep_template=pt.id)
        self.files_to_remove.extend([x['fp'] for x in obs.filepaths])
        self.assertEqual(obs.id, self.artifact_count + 1)
        self.assertTrue(
            qdb.util.check_count('qiita.filepath', self.fp_count + 4))

    def test_load_artifact_from_cmd_processed(self):
        fd, file1 = mkstemp()
        close(fd)
        self.files_to_remove.append(file1)
        fd, file2 = mkstemp()
        close(fd)
        self.files_to_remove.append(file2)
        fps = [file1, file2]
        ftypes = ['preprocessed_fasta', 'preprocessed_fastq']
        for fp in fps:
            with open(fp, 'w') as f:
                f.write("\n")
        obs = qdb.commands.load_artifact_from_cmd(
            fps, ftypes, 'Demultiplexed', parents=[1], dflt_params_id=1,
            required_params='{"input_data": 1}',
            optional_params='{"min_per_read_length_fraction": 0.80}')
        self.files_to_remove.extend([x['fp'] for x in obs.filepaths])
        self.assertEqual(obs.id, self.artifact_count + 1)
        self.assertTrue(
            qdb.util.check_count('qiita.filepath', self.fp_count + 2))

    def test_load_artifact_from_cmd_biom(self):
        fd, otu_table_fp = mkstemp(suffix='_otu_table.biom')
        close(fd)
        self.files_to_remove.append(otu_table_fp)
        fps = [otu_table_fp]
        ftypes = ['biom']
        for fp in fps:
            with open(fp, 'w') as f:
                f.write("\n")
        obs = qdb.commands.load_artifact_from_cmd(
            fps, ftypes, 'BIOM', parents=[3], dflt_params_id=10,
            required_params='{"input_data": 3}')
        self.files_to_remove.extend([x['fp'] for x in obs.filepaths])
        self.assertEqual(obs.id, self.artifact_count + 1)
        self.assertTrue(
            qdb.util.check_count('qiita.filepath', self.fp_count + 1))


@qiita_test_checker()
class TestLoadSampleTemplateFromCmd(TestCase):
    def setUp(self):
        # Create a sample template file
        self.st_contents = SAMPLE_TEMPLATE

        # create a new study to attach the sample template
        info = {
            "timeseries_type_id": 1,
            "metadata_complete": True,
            "mixs_compliant": True,
            "study_alias": "TestStudy",
            "study_description": "Description of a test study",
            "study_abstract": "No abstract right now...",
            "principal_investigator_id": qdb.study.StudyPerson(3),
            "lab_person_id": qdb.study.StudyPerson(1)
        }
        self.study = qdb.study.Study.create(
            qdb.user.User('test@foo.bar'), "Test study", info)

    def test_load_sample_template_from_cmd(self):
        """Correctly adds a sample template to the DB"""
        fh = StringIO(self.st_contents)
        st = qdb.commands.load_sample_template_from_cmd(fh, self.study.id)
        self.assertEqual(st.id, self.study.id)


@qiita_test_checker()
class TestLoadPrepTemplateFromCmd(TestCase):
    def setUp(self):
        self.pt_contents = PREP_TEMPLATE

    def test_load_prep_template_from_cmd(self):
        """Correctly adds a prep template to the DB"""
        fh = StringIO(self.pt_contents)
        st = qdb.commands.load_prep_template_from_cmd(fh, 1, '18S')
        self.assertEqual(st.id, 3)


@qiita_test_checker()
class TestLoadParametersFromCmd(TestCase):
    def setUp(self):
        fd, self.fp = mkstemp(suffix='_params.txt')
        close(fd)

        fd, self.fp_wrong = mkstemp(suffix='_params.txt')
        close(fd)

        with open(self.fp, 'w') as f:
            f.write(PARAMETERS)

        with open(self.fp_wrong, 'w') as f:
            f.write(PARAMETERS_ERROR)

        self.files_to_remove = [self.fp, self.fp_wrong]

    def tearDown(self):
        for fp in self.files_to_remove:
            if exists(fp):
                remove(fp)


@qiita_test_checker()
class TestPatch(TestCase):
    def setUp(self):
        self.patches_dir = mkdtemp()
        self.py_patches_dir = join(self.patches_dir, 'python_patches')
        mkdir(self.py_patches_dir)
        patch2_fp = join(self.patches_dir, '2.sql')
        patch10_fp = join(self.patches_dir, '10.sql')

        with open(patch2_fp, 'w') as f:
            f.write("CREATE TABLE qiita.patchtest2 (testing integer);\n")
            f.write("INSERT INTO qiita.patchtest2 VALUES (1);\n")
            f.write("INSERT INTO qiita.patchtest2 VALUES (9);\n")

        with open(patch10_fp, 'w') as f:
            f.write("CREATE TABLE qiita.patchtest10 (testing integer);\n")

    def tearDown(self):
        rmtree(self.patches_dir)
        # The tests on this class are really tied up to the status of the
        # database, so we do an exception and reset the DB in each test
        qdb.environment_manager.drop_and_rebuild_tst_database()

    def _check_patchtest2(self, exists=True):
        if exists:
            assertion_fn = self.assertTrue
        else:
            assertion_fn = self.assertFalse

        with qdb.sql_connection.TRN:
            qdb.sql_connection.TRN.add("""SELECT EXISTS(SELECT *
                                          FROM information_schema.tables
                                          WHERE table_name = 'patchtest2')""")
            obs = qdb.sql_connection.TRN.execute_fetchflatten()[0]
        assertion_fn(obs)

        if exists:
            exp = [[1], [9]]
            with qdb.sql_connection.TRN:
                qdb.sql_connection.TRN.add(
                    """SELECT * FROM qiita.patchtest2 ORDER BY testing""")
                obs = qdb.sql_connection.TRN.execute_fetchindex()
            self.assertEqual(obs, exp)

    def _check_patchtest10(self):
        with qdb.sql_connection.TRN:
            qdb.sql_connection.TRN.add(
                """SELECT EXISTS(SELECT * FROM information_schema.tables
                   WHERE table_name = 'patchtest10')""")
            obs = qdb.sql_connection.TRN.execute_fetchflatten()[0]
        self.assertTrue(obs)

        exp = []
        with qdb.sql_connection.TRN:
            qdb.sql_connection.TRN.add("SELECT * FROM qiita.patchtest10")
            obs = qdb.sql_connection.TRN.execute_fetchindex()
        self.assertEqual(obs, exp)

    def _assert_current_patch(self, patch_to_check):
        with qdb.sql_connection.TRN:
            qdb.sql_connection.TRN.add("SELECT current_patch FROM settings")
            current_patch = qdb.sql_connection.TRN.execute_fetchflatten()[0]
        self.assertEqual(current_patch, patch_to_check)

    def test_unpatched(self):
        """Test patching from unpatched state"""
        # Reset the settings table to the unpatched state
        qdb.sql_connection.perform_as_transaction(
            "UPDATE settings SET current_patch = 'unpatched'")

        self._assert_current_patch('unpatched')
        qdb.environment_manager.patch(self.patches_dir)
        self._check_patchtest2()
        self._check_patchtest10()
        self._assert_current_patch('10.sql')

    def test_skip_patch(self):
        """Test patching from a patched state"""
        qdb.sql_connection.perform_as_transaction(
            "UPDATE settings SET current_patch = '2.sql'")
        self._assert_current_patch('2.sql')

        # If it tried to apply patch 2.sql again, this will error
        qdb.environment_manager.patch(self.patches_dir)

        self._assert_current_patch('10.sql')
        self._check_patchtest10()

        # Since we "tricked" the system, patchtest2 should not exist
        self._check_patchtest2(exists=False)

    def test_nonexistent_patch(self):
        """Test case where current patch does not exist"""
        qdb.sql_connection.perform_as_transaction(
            "UPDATE settings SET current_patch = 'nope.sql'")
        self._assert_current_patch('nope.sql')

        with self.assertRaises(RuntimeError):
            qdb.environment_manager.patch(self.patches_dir)

    def test_python_patch(self):
        # Write a test python patch
        patch10_py_fp = join(self.py_patches_dir, '10.py')
        with open(patch10_py_fp, 'w') as f:
            f.write(PY_PATCH)

        # Reset the settings table to the unpatched state
        qdb.sql_connection.perform_as_transaction(
            "UPDATE settings SET current_patch = 'unpatched'")

        self._assert_current_patch('unpatched')

        qdb.environment_manager.patch(self.patches_dir)

        with qdb.sql_connection.TRN:
            qdb.sql_connection.TRN.add("SELECT testing FROM qiita.patchtest10")
            obs = qdb.sql_connection.TRN.execute_fetchindex()
        exp = [[1], [100]]
        self.assertEqual(obs, exp)

        self._assert_current_patch('10.sql')


@qiita_test_checker()
class TestUpdateArtifactFromCmd(TestCase):
    def setUp(self):
        fd, seqs_fp = mkstemp(suffix='_seqs.fastq')
        close(fd)
        fd, barcodes_fp = mkstemp(suffix='_barcodes.fastq')
        close(fd)
        self.filepaths = [seqs_fp, barcodes_fp]
        self.checksums = []
        for fp in sorted(self.filepaths):
            with open(fp, 'w') as f:
                f.write("%s\n" % fp)
            self.checksums.append(qdb.util.compute_checksum(fp))
        self.filepaths_types = ["raw_forward_seqs", "raw_barcodes"]
        self._clean_up_files = [seqs_fp, barcodes_fp]
        self.uploaded_files = qdb.util.get_files_from_uploads_folders("1")

    def tearDown(self):
        new_uploaded_files = qdb.util.get_files_from_uploads_folders("1")
        new_files = set(new_uploaded_files).difference(self.uploaded_files)
        path_builder = partial(
            join, qdb.util.get_mountpoint("uploads")[0][1], '1')
        self._clean_up_files.extend(
            [path_builder(fp) for _, fp, _ in new_files])
        for f in self._clean_up_files:
            if exists(f):
                remove(f)

    def test_update_artifact_from_cmd_error(self):
        with self.assertRaises(ValueError):
            qdb.commands.update_artifact_from_cmd(
                self.filepaths[1:], self.filepaths_types, 1)

        with self.assertRaises(ValueError):
            qdb.commands.update_artifact_from_cmd(
                self.filepaths, self.filepaths_types[1:], 1)

    def test_update_artifact_from_cmd(self):
        # Generate some files for an artifact
        fd, fp1 = mkstemp(suffix='_seqs.fastq')
        close(fd)
        with open(fp1, 'w') as f:
            f.write("@HWI-ST753:189:D1385ACXX:1:1101:1214:1906 1:N:0:\n"
                    "NACGTAGGGTGCAAGCGTTGTCCGGAATNA\n"
                    "+\n"
                    "#1=DDFFFHHHHHJJJJJJJJJJJJGII#0\n")

        fd, fp2 = mkstemp(suffix='_barcodes.fastq')
        close(fd)
        with open(fp2, 'w') as f:
            f.write("@HWI-ST753:189:D1385ACXX:1:1101:1214:1906 2:N:0:\n"
                    "NNNCNNNNNNNNN\n"
                    "+\n"
                    "#############\n")
        filepaths = [(fp1, 1), (fp2, 3)]
        # Create a new prep template
        metadata_dict = {
            'SKB8.640193': {'center_name': 'ANL',
                            'primer': 'GTGCCAGCMGCCGCGGTAA',
                            'barcode': 'GTCCGCAAGTTA',
                            'run_prefix': "s_G1_L001_sequences",
                            'platform': 'Illumina',
                            'instrument_model': 'Illumina MiSeq',
                            'library_construction_protocol': 'AAAA',
                            'experiment_design_description': 'BBBB'}}
        metadata = pd.DataFrame.from_dict(metadata_dict, orient='index',
                                          dtype=str)
        self.prep_template = \
            qdb.metadata_template.prep_template.PrepTemplate.create(
                metadata, qdb.study.Study(1), "16S")
        artifact = qdb.artifact.Artifact.create(
            filepaths, "FASTQ", prep_template=self.prep_template)
        for x in artifact.filepaths:
            self._clean_up_files.append(x['fp'])

        new_artifact = qdb.commands.update_artifact_from_cmd(
            self.filepaths, self.filepaths_types, artifact.id)
        for x in new_artifact.filepaths:
            self._clean_up_files.append(x['fp'])

        for obs, exp in zip(sorted(artifact.filepaths, key=lambda x: x['fp']),
                            self.checksums):
            self.assertEqual(qdb.util.compute_checksum(obs['fp']), exp)


CONFIG_1 = """[required]
timeseries_type_id = 1
metadata_complete = True
mixs_compliant = True
principal_investigator = SomeDude, somedude@foo.bar, some
reprocess = False
study_alias = 'test study'
study_description = 'test study description'
study_abstract = 'study abstract'
efo_ids = 1,2,3,4
[optional]
lab_person = SomeDude, somedude@foo.bar, some
funding = 'funding source'
vamps_id = vamps_id
"""

CONFIG_2 = """[required]
timeseries_type_id = 1
metadata_complete = True
principal_investigator = SomeDude, somedude@foo.bar, some
reprocess = False
study_alias = 'test study'
study_description = 'test study description'
study_abstract = 'study abstract'
efo_ids = 1,2,3,4
[optional]
lab_person = SomeDude, somedude@foo.bar, some
funding = 'funding source'
vamps_id = vamps_id
"""

SAMPLE_TEMPLATE = (
    "sample_name\trequired_sample_info_status\tcollection_timestamp\t"
    "sample_type\tphysical_specimen_remaining\tphysical_specimen_location\t"
    "dna_extracted\thost_subject_id\tTreatment\tDOB\tlatitude\tlongitude"
    "\ttaxon_id\tscientific_name\tDescription\n"
    "PC.354\treceived\t06/18/14 16:44:00\ttype_1\tTrue\tLocation_1\tTrue\t"
    "HS_ID_PC.354\tControl\t20061218\t1.88401499993\t56.0003871552\t"
    "9606\thomo sapiens\tControl_mouse_I.D._354\n"
    "PC.593\treceived\t06/18/14 16:44:00\ttype_1\tTrue\tLocation_1\tTrue\t"
    "HS_ID_PC.593\tControl\t20071210\t35.4079458313\t83.2595338611\t"
    "9606\thomo sapiens\tControl_mouse_I.D._593\n"
    "PC.607\treceived\t06/18/14 16:44:00\ttype_1\tTrue\tLocation_1\tTrue\t"
    "HS_ID_PC.607\tFast\t20071112\t18.3175615444\t91.3713989729\t"
    "9606\thomo sapiens\tFasting_mouse_I.D._607\n"
    "PC.636\treceived\t06/18/14 16:44:00\ttype_1\tTrue\tLocation_1\tTrue\t"
    "HS_ID_PC.636\tFast\t20080116\t31.0856060708\t4.16781143893\t"
    "9606\thomo sapiens\tFasting_mouse_I.D._636")

PREP_TEMPLATE = (
    'sample_name\tbarcode\tcenter_name\tcenter_project_name\t'
    'description_prep\tebi_submission_accession\temp_status\tprimer\t'
    'run_prefix\tstr_column\tplatform\tlibrary_construction_protocol\t'
    'experiment_design_description\tinstrument_model\n'
    'SKB7.640196\tCCTCTGAGAGCT\tANL\tTest Project\tskb7\tNone\tEMP\t'
    'GTGCCAGCMGCCGCGGTAA\tts_G1_L001_sequences\tValue for sample 3\tA\tB\tC\t'
    'Illumina MiSeq\n'
    'SKB8.640193\tGTCCGCAAGTTA\tANL\tTest Project\tskb8\tNone\tEMP\t'
    'GTGCCAGCMGCCGCGGTAA\tts_G1_L001_sequences\tValue for sample 1\tA\tB\tC\t'
    'Illumina MiSeq\n'
    'SKD8.640184\tCGTAGAGCTCTC\tANL\tTest Project\tskd8\tNone\tEMP\t'
    'GTGCCAGCMGCCGCGGTAA\tts_G1_L001_sequences\tValue for sample 2\tA\tB\tC\t'
    'Illumina MiSeq\n')

PY_PATCH = """
from qiita_db.study import Study
from qiita_db.sql_connection import TRN
study = Study(1)

with TRN:
    sql = "INSERT INTO qiita.patchtest10 (testing) VALUES (%s)"
    TRN.add(sql, [[study.id], [study.id*100]], many=True)
    TRN.execute()
"""

PARAMETERS = """max_bad_run_length\t3
min_per_read_length_fraction\t0.75
sequence_max_n\t0
rev_comp_barcode\tFalse
rev_comp_mapping_barcodes\tFalse
rev_comp\tFalse
phred_quality_threshold\t3
barcode_type\thamming_8
max_barcode_errors\t1.5
"""

PARAMETERS_ERROR = """max_bad_run_length\t3\tmin_per_read_length_fraction\t0.75
sequence_max_n\t0
rev_comp_barcode\tFalse
rev_comp_mapping_barcodes\tFalse
rev_comp\tFalse
phred_quality_threshold\t3
barcode_type\thamming_8
max_barcode_errors\t1.5
"""


if __name__ == "__main__":
    main()