import unittest.mock

import pytest
from db_schema.factories.apps import ContentStatusFactory, DatabricksExecutionFactory
from db_schema.schemas.apps import DatabricksExecution, DBXJobStatusEnum

from apps_etl_databricks import exceptions
from apps_etl_databricks.commands import cancel, start, sync
from apps_etl_databricks.entities import CancelInput, Config, StartInput, SyncInput
from apps_etl_databricks.objects.run import RunState


@pytest.mark.integration
@unittest.mock.patch('time.sleep', return_value=None)
@unittest.mock.patch('requests.get')
@unittest.mock.patch('requests.post')
def test_start(mocked_post, mocked_get, _, session):
    sf_name, job_id, run_id = 'sf-execution-name', 14, 99
    mocked_post.return_value.json.return_value = {'run_id': run_id}

    mocked_get.return_value.json.return_value = {'state': RunState(life_cycle_state='IN_PROGRESS')}
    content_status = ContentStatusFactory()

    config = Config(
        environment='dev',
        dbx_url='dbx-url.com',
        dbx_token='dbx-token',
        sqs_arns='sqs_arns',
        rds_secret_key='rds_secret_key',
    )
    start_input = StartInput(
        dsp='apple',
        report_date='2020-02-02',
        unit_of_work_id=content_status.unit_of_work.unit_of_work_id,
        licensor_ids=[],
        content_status_ids=[content_status.content_status_id],
        sf_token='test-sf-token',
        sf_execution_name=sf_name,
        dbx_job_id=job_id,
        dbx_job_max_lifetime=3600,
        dbx_spark_params=[],
        conversion_params={},
    )

    start.execute(unittest.mock.Mock(), session, config, start_input)

    dbx_execution = session.query(DatabricksExecution).one_or_none()

    assert dbx_execution
    assert dbx_execution.status == DBXJobStatusEnum.IN_PROGRESS
    assert dbx_execution.sf_execution_name == sf_name
    assert dbx_execution.unit_of_work_id == content_status.unit_of_work_id
    assert content_status.latest_run_id == run_id


@pytest.mark.integration
@unittest.mock.patch('boto3.client')
@unittest.mock.patch('requests.get')
def test_sync(mocked_get, _, session):
    mocked_get.return_value.json.return_value = {
        'state': {
            'life_cycle_state': 'TERMINATED',
            'result_state': 'SUCCESS',
        },
    }
    dbx_execution = DatabricksExecutionFactory()

    config = Config(
        environment='dev',
        dbx_url='dbx-url.com',
        dbx_token='dbx-token',
        sqs_arns='sqs_arns',
        rds_secret_key='rds_secret_key',
    )

    sync_input = SyncInput(
        sf_execution_name=dbx_execution.sf_execution_name,
        conversion_params={},
    )

    sync.execute(unittest.mock.Mock(), session, config, sync_input)

    assert dbx_execution.completed_at
    assert dbx_execution.status == DBXJobStatusEnum.COMPLETE


@pytest.mark.integration
@unittest.mock.patch('requests.post')
def test_cancel(_, session):
    dbx_execution = DatabricksExecutionFactory()

    config = Config(
        environment='dev',
        dbx_url='dbx-url.com',
        dbx_token='dbx-token',
        sqs_arns='sqs_arns',
        rds_secret_key='rds_secret_key',
    )

    cancel_input = CancelInput(sf_execution_name=dbx_execution.sf_execution_name)
    cancel.execute(unittest.mock.Mock(), session, config, cancel_input)

    assert dbx_execution.completed_at
    assert dbx_execution.status == DBXJobStatusEnum.CANCELLED


@pytest.mark.integration
@pytest.mark.parametrize('sqs_push_result', (True, False))
@unittest.mock.patch('boto3.client')
@unittest.mock.patch('apps_etl_utils.utils.get_s3_keys')
@unittest.mock.patch('slz_notification.sqs_service.SQSNotificationService.push')
@unittest.mock.patch('requests.get')
def test_sync_with_conversion_params_send_message_to_sqs(
    mocked_get, mocked_sqs_push, _1, _2, sqs_push_result, session
):
    mocked_get.return_value.json.return_value = {
        'state': {
            'life_cycle_state': 'TERMINATED',
            'result_state': 'SUCCESS',
        },
    }
    mocked_sqs_push.return_value = sqs_push_result, None

    dbx_execution = DatabricksExecutionFactory()

    config = Config(
        environment='dev',
        dbx_url='dbx-url.com',
        dbx_token='dbx-token',
        sqs_arns='sqs_arns',
        rds_secret_key='rds_secret_key',
    )

    conversion_params = {
        'contexts': [{
            'report_type': 'report_type',
            'uow_id': 'uow_id',
            'compressed_path': 'compressed_path',
            'decompressed_bucket': 'decompressed_bucket',
            'decompressed_folder': 'decompressed_folder',
            'content_name': 'content_name',
            'context': 'context',
            'unit_of_work_id': 'unit_of_work_id'
        }]
    }

    sync_input = SyncInput(
        sf_execution_name=dbx_execution.sf_execution_name,
        conversion_params=conversion_params,
    )

    if sqs_push_result:
        sync.execute(unittest.mock.Mock(), session, config, sync_input)
        assert dbx_execution.completed_at
        assert dbx_execution.status == DBXJobStatusEnum.COMPLETE

    else:
        # exception should be raised
        with pytest.raises(exceptions.SQSSendMessageError):
            sync.execute(unittest.mock.Mock(), session, config, sync_input)