import unittest
from unittest import mock

import pytest
from airflow.exceptions import AirflowException
from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook

from twitch_airflow_components.operators.sagemaker_training import SageMakerTrainingOperator

role = "arn:aws:iam:role/test-role"

bucket = "test-bucket"

key = "test/data"
data_url = f"s3://{bucket}/{key}"

job_name = "test-job-name"

image = "test-image"

output_url = f"s3://{bucket}/test/output"
create_training_params = {
    "AlgorithmSpecification": {"TrainingImage": image, "TrainingInputMode": "File"},
    "RoleArn": role,
    "OutputDataConfig": {"S3OutputPath": output_url},
    "ResourceConfig": {
        "InstanceCount": "2",
        "InstanceType": "ml.c4.8xlarge",
        "VolumeSizeInGB": "50",
    },
    "TrainingJobName": job_name,
    "HyperParameters": {
        "k": "10",
        "feature_dim": "784",
        "mini_batch_size": "500",
        "force_dense": "True",
    },
    "StoppingCondition": {"MaxRuntimeInSeconds": "3600"},
    "InputDataConfig": [
        {
            "ChannelName": "train",
            "DataSource": {
                "S3DataSource": {
                    "S3DataType": "S3Prefix",
                    "S3Uri": data_url,
                    "S3DataDistributionType": "FullyReplicated",
                }
            },
            "CompressionType": "None",
            "RecordWrapperType": "None",
        }
    ],
}


# pylint: disable=unused-argument
class TestSageMakerTrainingOperator(unittest.TestCase):
    def setUp(self):
        self.sagemaker = SageMakerTrainingOperator(
            task_id="test_sagemaker_operator",
            aws_conn_id="sagemaker_test_id",
            config=create_training_params,
            wait_for_completion=False,
            check_interval=5,
        )

    def test_parse_config_integers(self):
        self.sagemaker.parse_config_integers()
        assert self.sagemaker.config["ResourceConfig"]["InstanceCount"] == int(
            self.sagemaker.config["ResourceConfig"]["InstanceCount"]
        )
        assert self.sagemaker.config["ResourceConfig"]["VolumeSizeInGB"] == int(
            self.sagemaker.config["ResourceConfig"]["VolumeSizeInGB"]
        )
        assert self.sagemaker.config["StoppingCondition"]["MaxRuntimeInSeconds"] == int(
            self.sagemaker.config["StoppingCondition"]["MaxRuntimeInSeconds"]
        )

    @mock.patch.object(SageMakerHook, "get_conn")
    @mock.patch.object(SageMakerHook, "create_training_job")
    def test_execute(self, mock_training, mock_client):
        mock_training.return_value = {
            "TrainingJobArn": "testarn",
            "ResponseMetadata": {"HTTPStatusCode": 200},
        }
        self.sagemaker.execute(None)
        mock_training.assert_called_once_with(
            create_training_params,
            wait_for_completion=False,
            print_log=True,
            check_interval=5,
            max_ingestion_time=None,
        )

    @mock.patch.object(SageMakerHook, "get_conn")
    @mock.patch.object(SageMakerHook, "create_training_job")
    def test_execute_with_failure(self, mock_training, mock_client):
        mock_training.return_value = {
            "TrainingJobArn": "testarn",
            "ResponseMetadata": {"HTTPStatusCode": 404},
        }
        with pytest.raises(AirflowException):
            self.sagemaker.execute(None)
