import unittest
from unittest.mock import MagicMock, Mock, patch

from db_schema.schemas.exploration import DisassembleStatusEnum
from exp_sf_storage.models.inbound_file_status import InboundFileStatusType

from exp_results_process_lambda import STATUS_FAILURE, handler
from exp_results_process_lambda.const import STATUS_OK
from exp_results_process_lambda.entities import Input


class HandlerTestCase(unittest.TestCase):
    def setUp(self):
        self.input = Input(
            archive_status='OK',
            chunks_count=10,
            content_name='foo_content_name',
            context='foo',
            dbx_state='OK',
            dest_bucket='foo_bucket',
            disassemble_content_status_id=11,
            file_lines_count=12,
            source_path='foo',
            start_dbx_status='OK',
            sync_dbx_status='OK',
            uow_id='bar',
            use_databricks=False,
        )

        self.s3_mock_list_objects = Mock(return_value=[
            'foo.gz',
            'bar.gz',
        ])
        self.slz_mock_fetch_db_data = Mock(
            return_value={
                'dsp': 'foo',
                'report_name': 'bar',
                'report_date': 'baz',
                'context': 'foobar',
                'licensor': 'bazbar',
                'replaced_disassemble_content_status_id': None,
                'retried_disassemble_content_status_id': None,
            }
        )
        self.slz_mock_update_disassemble_content_status = Mock(return_value=True)
        self.snowflake_mock_create_file_metadata = Mock()

        self.patchers = [
            patch(
                'exp_s3_storage.service.S3Service.list_objects',
                self.s3_mock_list_objects,
            ),
            patch(
                'exp_results_process_lambda.services.slz_db.SlzDbService.fetch_db_data',
                self.slz_mock_fetch_db_data,
            ),
            patch(
                'exp_results_process_lambda.services.slz_db.SlzDbService.'
                'update_disassemble_content_status',
                self.slz_mock_update_disassemble_content_status,
            ),
            patch(
                'exp_results_process_lambda.services.snowflake.SnowflakeService.'
                'create_file_metadata',
                self.snowflake_mock_create_file_metadata,
            ),
        ]

        for patcher in self.patchers:
            patcher.start()

    def tearDown(self):
        for patcher in self.patchers:
            patcher.stop()

    def test_inbound_archive_status_failure(self):
        self.input.archive_status = STATUS_FAILURE
        expected = {'meta': {'status': STATUS_FAILURE, 'message': 'Invalid dep status'}}

        actual = handler(Mock(), Mock(), Mock(), self.input, Mock())

        self.assertEqual(expected, actual)

    def test_inbound_start_dbx_status_failure(self):
        self.input.start_dbx_status = STATUS_FAILURE
        expected = {'meta': {'status': STATUS_FAILURE, 'message': 'Invalid dep status'}}

        actual = handler(Mock(), Mock(), Mock(), self.input, Mock())

        self.assertEqual(expected, actual)

    def test_inbound_sync_dbx_status_failure(self):
        self.input.sync_dbx_status = STATUS_FAILURE
        expected = {'meta': {'status': STATUS_FAILURE, 'message': 'Invalid dep status'}}

        actual = handler(Mock(), Mock(), Mock(), self.input, Mock())

        self.assertEqual(expected, actual)

    def test_arhive_flow__ok(self):
        self.input.chunks_count = 2  # use 2 as an example
        expected = {'meta': {'status': STATUS_OK, 'sf_status': 'LOADED', 'slz_status': 'COMPLETE'}}

        pg_repo = Mock()
        pg_repo.advisory_locked_transaction = MagicMock(MockContextManager())

        actual = handler(Mock(), pg_repo, Mock(), self.input, Mock())

        self.assertEqual(expected, actual)
        self.s3_mock_list_objects.assert_called_once_with()
        self.slz_mock_fetch_db_data.assert_called_once_with()
        self.slz_mock_update_disassemble_content_status.assert_called_once_with(
            DisassembleStatusEnum.COMPLETE,
            self.input.chunks_count,
        )
        self.snowflake_mock_create_file_metadata.assert_called_once_with(
            {
                'dsp': 'foo',
                'report_name': 'bar',
                'report_date': 'baz',
                'context': 'foobar',
                'licensor': 'bazbar',
                'replaced_disassemble_content_status_id': None,
                'retried_disassemble_content_status_id': None,
            },
            InboundFileStatusType.LOADED,
            ['foo.gz', 'bar.gz'],
        )

    def test_dbx_flow__ok(self):
        self.input.use_databricks = True
        self.input.chunks_count = 2
        expected = {'meta': {'status': STATUS_OK, 'sf_status': 'LOADED', 'slz_status': 'COMPLETE'}}

        pg_repo = Mock()
        pg_repo.advisory_locked_transaction = MagicMock(MockContextManager())

        actual = handler(Mock(), pg_repo, Mock(), self.input, Mock())

        self.assertEqual(expected, actual)
        self.s3_mock_list_objects.assert_called_once_with()
        self.slz_mock_fetch_db_data.assert_called_once_with()
        self.slz_mock_update_disassemble_content_status.assert_called_once_with(
            DisassembleStatusEnum.COMPLETE,
            self.input.chunks_count,
        )
        self.snowflake_mock_create_file_metadata.assert_called_once_with(
            {
                'dsp': 'foo',
                'report_name': 'bar',
                'report_date': 'baz',
                'context': 'foobar',
                'licensor': 'bazbar',
                'replaced_disassemble_content_status_id': None,
                'retried_disassemble_content_status_id': None,
            },
            InboundFileStatusType.LOADED,
            ['foo.gz', 'bar.gz'],
        )

    def test_dbx_flow__count_mismatch(self):
        self.input.use_databricks = True
        self.input.chunks_count = 3
        expected = {'meta': {'status': STATUS_OK, 'sf_status': 'FAILED', 'slz_status': 'FAILED'}}

        pg_repo = Mock()
        pg_repo.advisory_locked_transaction = MagicMock(MockContextManager())

        actual = handler(Mock(), pg_repo, Mock(), self.input, Mock())

        self.assertEqual(expected, actual)
        self.s3_mock_list_objects.assert_called_once_with()
        self.slz_mock_fetch_db_data.assert_called_once_with()
        self.slz_mock_update_disassemble_content_status.assert_called_once_with(
            DisassembleStatusEnum.FAILED,
            self.input.chunks_count,
        )
        self.snowflake_mock_create_file_metadata.assert_called_once_with(
            {
                'dsp': 'foo',
                'report_name': 'bar',
                'report_date': 'baz',
                'context': 'foobar',
                'licensor': 'bazbar',
                'replaced_disassemble_content_status_id': None,
                'retried_disassemble_content_status_id': None,
            },
            InboundFileStatusType.FAILED,
            ['foo.gz', 'bar.gz'],
        )


class MockContextManager:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        pass