#!/usr/bin/env python3

import os
import unittest
import subprocess
import shutil
import tarfile


class TestFormatAndPackage(unittest.TestCase):
    """ Tests the format_and_package_l2 script and the scripts it calls.
        Verification is different between the TIFF and HDF options.

        For TIFF it executes the script and verifies the output tarball is
        valid by extracting it and verifying each file. By verifying each file
        individually, the scripts that are called by format_and_package_l2 are
        checked to make sure they are working as expected.

        For HDF it executes the script and verifies the output files. It also
        verifies that create_l2_mtl and package_product weren't executed.
    """
    def setUp(self):
        """ Setup up shared values and files for both tests. """
        # Get the test case name
        self.test_case_dir = self.id().split('.')[-1]

        # Create a new directory for the test and switch to that directory.
        # This is done so that the tests can be run in parallel.
        self.cwd = os.getcwd()
        os.makedirs(self.test_case_dir, exist_ok=True)
        os.chdir(self.test_case_dir)

        self.test_file_dir = os.path.join(
            os.environ['ESPA_UNIT_TEST_DATA_DIR'],
            'espa-product-formatter/format_and_package_l2')
        self.output_dir = 'test_output'
        self.symlink_test_data(self.test_file_dir)

    def tearDown(self):
        """ Make sure the current working directory is changed back """
        # Make sure that we change the current working directory back to the
        # original directory regardless of test outcome
        os.chdir(self.cwd)

    def symlink_test_data(self, data_directory):
        """ Helper function to create symlinks to the unit test data files. """
        for test_file in os.listdir(data_directory):
            test_file_path = os.path.join(data_directory, test_file)
            if os.path.isfile(test_file_path):
                # If the test file is already symlinked, remove the link and
                # create a new one.
                if os.path.exists(test_file):
                    os.unlink(test_file)
                os.symlink(test_file_path, test_file)

    def get_product_id(self):
        """ Parses the product_data.txt file and returns product ID from the
            file.
        """
        with open('product_data.txt') as product_data:
            for line in product_data:
                try:
                    (key, value) = line.split('=')
                except ValueError:
                    # Just want to skip to the next one as some lines don't
                    # have a value
                    pass
                if key == 'LANDSAT_L2_PRODUCT_ID':
                    output_product_id = value.strip()
                    return output_product_id

    def verify_files(self, output_name, expected_dir, output_dir):
        """ Checks that the number of output files matches the number of
            expected files. After that, each output file is individually
            compared against its expected file.

            Note:
                The diff command that is used to verify the individual files
                takes into account values that change from each test run so
                that they won't cause failures.

            Args:
                output_name (str): The portion of the output file names that
                    is the same for all output files. For TIFF this is the
                    product ID and HDF is what was passed into the hdf_filename
                    parameter.
                expected_dir (str): The path to the directory containing the
                    expected files.
                output_dir (str): The path to the directory contain the output
                    to compare against the expected files.
        """
        # Verify that the correct number of files are in the output dir
        expected_num_files = len(os.listdir(expected_dir))
        output_num_files = len(os.listdir(output_dir))
        fail_msg = ('Incorrect number of files found in ' + output_dir + '. '
                    'Expected: ' + str(expected_num_files) + ' Found: ' +
                    str(output_num_files))
        self.assertEqual(expected_num_files, output_num_files, msg=fail_msg)

        # Verify each extracted file
        for file_name in os.listdir(expected_dir):
            expected_file = os.path.join(expected_dir, file_name)
            # Get the part of the file name that indicates the files contents
            file_type = file_name.split('expected')[-1]
            # This check is for the hdf file. Diff won't work correctly due to
            # variations in the hdf file, so this skips it.
            if file_type == '':
                continue

            output_file = os.path.join(output_dir, output_name + file_type)

            # Sed command to ignore any fields that can change from run to run.
            # (For metadata files only)
            if 'MTL.txt' in file_type:
                sed_cmd = ['sed', '-i',
                           '-e', '/LANDSAT_PRODUCT_ID/ s/ ".*"//',
                           '-e', '/FILE_NAME/ s/ ".*"//',
                           '-e', '/PROCESSING_SOFTWARE_VERSION/ s/= .*/=/',
                           '-e', '/DATE_PRODUCT_GENERATED/ s/= .*/=/',
                           output_file]
                status = subprocess.call(sed_cmd)
                self.assertEqual(status, 0, msg='sed command failed (MTL.txt)')

            elif 'MTL.xml' in file_type:
                sed_cmd = ['sed', '-i',
                           '-e', '/LANDSAT_PRODUCT_ID/ s/>.*</></',
                           '-e', '/FILE_NAME/ s/>.*</></',
                           '-e', '/DATE_PRODUCT_GENERATED/ s/>.*</></',
                           '-e', '/PROCESSING_SOFTWARE_VERSION/ s/>.*</></',
                           output_file]
                status = subprocess.call(sed_cmd)
                self.assertEqual(status, 0, msg='sed command failed (MTL.xml')

            # Metadata files went through the sed command already
            # so plain diff command for all output files
            diff_cmd = ['diff', output_file, expected_file]

            status = subprocess.call(diff_cmd)
            self.assertEqual(
                status, 0, msg=output_file + ' should match ' + expected_file)

            # Remove the file that passed. This way if the test fails, only the
            # files that weren't checked or failed will remain
            os.unlink(output_file)

    def execute_format_and_package(self, xml, output_type,
                                   l2_product_type=None, disable_cog=True,
                                   include_angle_bands=False):
        """ Executes the format_and_package_l2 script and verifies that it
            succeeded in executing.

            Note:
                The fail assertion inside the except will cause the stack trace
                upon a test failure to show two exceptions. However, this is
                the usual way to handle test failure on an exception while
                using Python's unittest. Without it, the test will fail but
                will fail with an error rather than show a test failure.

            Args:
                xml (str): The filename of the level 2 xml file to use as
                           input.
                output_type (str): The type of output to use. This is either
                    HDF or TIFF
                l2_product_type (None/str): The defined product type. This is
                    either None(L2 SR/ST product) or 'sr'(L2 SR-Only product)
                    or 'sr_st_toa_bt'(L2 SR/ST/TOA/BT product)
                    or 'sr_toa_bt'(L2 SR-Only Albers Product).
                disable_cog (True/False): Flag to indicate if the GeoTIFFs
                    files should be converted to COG format. A true value
                    indicates they should not be converted.
        """
        # Depends on the desired product type, call the different command
        if (l2_product_type is None):
            cmd = ['format_and_package_l2.py', '--xml', xml, '--output_dir',
                   self.output_dir]
        else:
            cmd = ['format_and_package_l2.py', '--xml', xml, '--output_dir',
                   self.output_dir, '--l2_product_type', l2_product_type]

        if output_type == 'hdf':
            cmd.append('--hdf_filename')
            cmd.append('hdf_test')
        else:
            if disable_cog:
                cmd.append('--disable_cog')

        if include_angle_bands:
            cmd.append('--include_angle_bands')

        try:
            subprocess.check_call(cmd)
        except Exception:
            # This will cause a double exception message in the stack trace if
            # it fails. Please see note in doc string for more information.
            self.fail(msg='Should be able to execute ' + str(cmd))

    def run_format_and_package_tiff_test(self, xml, expected_dir,
                                         l2_product_type=None,
                                         disable_cog=True,
                                         include_angle_bands=False):
        """ Tests the format_and_package_l2 script with TIFF output. The script
            is executed and its output tarball is extracted. The number of
            files that were extracted is compared against the expected amount.
            Each file is then verified individually against an expected
            version.

            Note: On failure the file that caused the failure will not be
                  deleted.
                  The product ID that is generated will change depending on the
                  time so the test figures out what product ID was generated so
                  that it can get the extracted files.
                  A few of the files have fields that will change run to run
                  due to date, algorithm and product ID variations. These are
                  accounted for in the diff command to avoid false test
                  failures.

            Args:
                xml (str): The xml file to use in the call to
                    format_and_package_l2.
                expected_dir (str): The directory containing the expected files
                    to compare the output too.
                l2_product_type (None/str): The defined product type. This is
                    either None(L2 SR/ST product) or 'sr'(L2 SR-Only product)
                    or 'sr_st_toa_bt'(L2 SR/ST/TOA/BT product)
                    or 'sr_toa_bt'(L2 SR-Only Albers Product).
                disable_cog (True/False): Flag to indicate if the GeoTIFFs
                    files should be converted to COG format. A true value
                    indicates they should not be converted.
        """
        # Setup specific data for the test
        test_directory = os.path.join(self.test_file_dir, expected_dir)
        self.symlink_test_data(test_directory)

        # If the previous run was a failure, the output dir may still exist
        if os.path.isdir(self.output_dir):
            shutil.rmtree(self.output_dir)

        # Execute format_and_package_l2 with TIFF output with l2_product_type
        self.execute_format_and_package(xml, 'tiff', l2_product_type,
                                        disable_cog, include_angle_bands)

        # Get the output product ID
        output_product_id = self.get_product_id()
        # If there is no LANDSAT_L2_PRODUCT_ID then this is not a valid
        # product_data.txt file and package_product failed.
        fail_msg = 'product_data.txt should have LANDSAT_L2_PRODUCT_ID'
        self.assertIsNotNone(output_product_id, msg=fail_msg)

        # Get the tarball file name
        tarball_file_name = output_product_id + '.tar.gz'

        # Extract the files from the tarball for verification
        extracted_dir = 'extracted_files'
        # If the previous run was a failure, the extracted dir may still exist
        if os.path.isdir(extracted_dir):
            shutil.rmtree(extracted_dir)
        tarball = tarfile.open(tarball_file_name)
        tarball.extractall(path=extracted_dir)

        # Verify the extracted files
        expected_dir = os.path.join(test_directory, 'expected_output')
        self.verify_files(output_product_id, expected_dir, extracted_dir)

        # Return to the original dir and remove the test case dir since the
        # test passed
        os.chdir(self.cwd)
        shutil.rmtree(self.test_case_dir)

    def test_format_and_package_tiff_oli_tirs(self):
        """ Runs the format_and_package test with an OLITIRS scene and TIFF
            output.
        """
        xml = 'LC08_L1TP_043028_20160401_20181108_02_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'tiff_oli_tirs')

    def test_format_and_package_tiff_etm(self):
        """ Runs the format_and_package test with an ETM scene and TIFF output.
        """
        xml = 'LE07_L1TP_043028_20020419_20181108_01_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'tiff_etm')

    def test_format_and_package_tiff_oli_tirs_sr_only(self):
        """ Runs the format_and_package test with an OLITIRS scene and TIFF
            output (SR-Only Product).
        """
        xml = 'LC08_L1TP_043028_20160401_20181108_02_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'tiff_oli_tirs_sr_only',
                                              'sr')

    def test_format_and_package_tiff_etm_sr_only(self):
        """ Runs the format_and_package test with an ETM scene and TIFF output
            (SR-Only Product).
        """
        xml = 'LE07_L1TP_043028_20020419_20181108_01_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'tiff_etm_sr_only', 'sr')

    def test_TM_sr_st_toa_bt(self):
        """ Runs the format_and_package test with a TM scene and TIFF output
            (SR-ST-TOA-BT Product).
        """
        xml = 'LT05_L1GS_022042_20090614_20190923_02_T2.xml'
        self.run_format_and_package_tiff_test(xml, 'test5', 'sr_st_toa_bt')

    def test_ETM_sr_st_toa_bt(self):
        """ Runs the format_and_package test with an ETM scene and TIFF output
            (SR-ST-TOA-BT Product).
        """
        xml = 'LE07_L1TP_162045_20130825_20190909_02_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'test6', 'sr_st_toa_bt')

    def test_OLITIRS_sr_st_toa_bt(self):
        """ Runs the format_and_package_l2 test with an OLI-TIRS scene and
            TIFF output (SR-ST-TOA-BT Product).
        """
        xml = 'LC08_L1TP_157070_20180321_20190830_02_T1.xml'
        self.run_format_and_package_tiff_test(
            xml, 'test7', 'SR_ST_TOA_BT', include_angle_bands=True)

    def test_format_and_package_cog_oli_tirs(self):
        """ Runs the format_and_package test with an OLITIRS scene and COG
            output.
        """
        xml = 'LC08_L1TP_043028_20160401_20181108_02_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'cog_oli_tirs', None, False)

    def test_format_and_package_cog_etm(self):
        """ Runs the format_and_package test with an ETM scene and COG output.
        """
        xml = 'LE07_L1TP_043028_20020419_20181108_01_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'cog_etm', None, False)

    def test_format_and_package_cog_oli_tirs_sr_only(self):
        """ Runs the format_and_package test with an OLITIRS scene and COG
            output (SR-Only Product).
        """
        xml = 'LC08_L1TP_043028_20160401_20181108_02_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'cog_oli_tirs_sr_only',
                                              'sr', False)

    def test_format_and_package_cog_etm_sr_only(self):
        """ Runs the format_and_package test with an ETM scene and COG output
            (SR-Only Product).
        """
        xml = 'LE07_L1TP_043028_20020419_20181108_01_T1.xml'
        self.run_format_and_package_tiff_test(xml, 'cog_etm_sr_only', 'sr',
                                              False)

    def test_TM_sr_toa_bt(self):
        """ Runs the format_and_package_l2 test with a TM scene and TIFF output
            (SR-Only Albers Product).
        """
        xml = 'LT05_L1TP_011028_20110927_20201006_02_A1.xml'
        self.run_format_and_package_tiff_test(xml, 'test8', 'sr_toa_bt')

    def test_ETM_sr_toa_bt(self):
        """ Runs the format_and_package_l2 test with an ETM scene and TIFF
            output (SR-Only Albers Product).
        """
        xml = 'LE07_L1TP_015029_20021024_20201006_02_A1.xml'
        self.run_format_and_package_tiff_test(
            xml, 'test9', 'sr_toa_bt', include_angle_bands=True)

    def test_OLITIRS_sr_toa_bt(self):
        """ Runs the format_and_package_l2 test with an OLITIRS scene and
            TIFF output (SR-Only Albers Product).
        """
        xml = 'LC08_L1TP_040033_20130709_20201005_02_A1.xml'
        self.run_format_and_package_tiff_test(xml, 'test10', 'SR_TOA_BT')

    def test_format_and_package_hdf(self):
        """ Runs the format_and_package test with an ETM scene and HDF output.
            The number of files that are in the output directory is compared
            against the expected amount. Each output file is then verified
            individually against an expected version. Verifies that the HDF
            option causes format_and_package_l2 to stop after running
            product_formatter.
        """
        xml = 'LE07_L1TP_043028_20020419_20181108_01_T1.xml'
        output_type = 'hdf'

        # Setup specific data for the test
        test_directory = os.path.join(self.test_file_dir, 'hdf')
        self.symlink_test_data(test_directory)

        # If the previous run was a failure, the output dir may still exist
        if os.path.isdir(self.output_dir):
            shutil.rmtree(self.output_dir)

        # Execute format_and_package_l2 with optional hdf_filename parameter
        # No hdf unit test for SR-Only product. L2 product type is set to None
        self.execute_format_and_package(xml, output_type)

        # Check to see if product_data.txt was edited. If LANDSAT_L2_PRODUCT_ID
        # exists then package_product was executed when it shouldn't have been.
        output_product_id = self.get_product_id()
        fail_msg = 'product_data.txt should not be edited with HDF'
        self.assertIsNone(output_product_id, msg=fail_msg)

        # Verify that create_l2_mtl didn't run
        for file_name in os.listdir(self.output_dir):
            fail_msg = ('Should not have any mtl files in the output '
                        'dir since create_l2_mtl should not run.')
            self.assertNotIn('_MTL', file_name, msg=fail_msg)

        # Verify the output files
        expected_dir = os.path.join(test_directory, 'expected_output')
        self.verify_files('hdf_test', expected_dir, self.output_dir)

        # Return to the original dir and remove the test case dir since the
        # test passed
        os.chdir(self.cwd)
        shutil.rmtree(self.test_case_dir)

if __name__ == '__main__':
    unittest.main(verbosity=2)
