import os
import time
from functools import wraps
from typing import NamedTuple
from unittest import mock

import boto3
import botocore
import numpy as np
import pytest
from click.testing import CliRunner
from moto.core import DEFAULT_ACCOUNT_ID
from sklearn.linear_model import LogisticRegression

import mlflow
import mlflow.pyfunc
import mlflow.sagemaker as mfs
import mlflow.sagemaker.cli as mfscli
import mlflow.sklearn
from mlflow.exceptions import MlflowException
from mlflow.models import Model
from mlflow.protos.databricks_pb2 import (
    INTERNAL_ERROR,
    INVALID_PARAMETER_VALUE,
    RESOURCE_DOES_NOT_EXIST,
    ErrorCode,
)
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
from mlflow.tracking.artifact_utils import _download_artifact_from_uri

from tests.helper_functions import set_boto_credentials  # noqa: F401
from tests.sagemaker.mock import TransformJob, TransformJobOperation, mock_sagemaker


class TrainedModel(NamedTuple):
    model_path: str
    run_id: str
    model_uri: str


@pytest.fixture
def pretrained_model():
    model_path = "model"
    with mlflow.start_run():
        X = np.array([-2, -1, 0, 1, 2, 1]).reshape(-1, 1)
        y = np.array([0, 0, 1, 1, 1, 0])
        lr = LogisticRegression(solver="lbfgs")
        lr.fit(X, y)
        mlflow.sklearn.log_model(lr, name=model_path)
        run_id = mlflow.active_run().info.run_id
        model_uri = "runs:/" + run_id + "/" + model_path
        return TrainedModel(model_path, run_id, model_uri)


@pytest.fixture
def sagemaker_client():
    return boto3.client("sagemaker", region_name="us-west-2")


def get_sagemaker_backend(region_name):
    return mock_sagemaker.backends[DEFAULT_ACCOUNT_ID][region_name]


def mock_sagemaker_aws_services(fn):
    from moto import mock_ecr, mock_iam, mock_s3, mock_sts

    @mock_ecr
    @mock_iam
    @mock_s3
    @mock_sagemaker
    @mock_sts
    @wraps(fn)
    def mock_wrapper(*args, **kwargs):
        # Create an ECR repository for the `mlflow-pyfunc` SageMaker docker image
        ecr_client = boto3.client("ecr", region_name="us-west-2")
        ecr_client.create_repository(repositoryName=mfs.DEFAULT_IMAGE_NAME)

        # Create the moto IAM role
        role_policy = """
        {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": "*",
                    "Resource": "*"
                }
            ]
        }
        """
        iam_client = boto3.client("iam", region_name="us-west-2")
        iam_client.create_role(RoleName="moto", AssumeRolePolicyDocument=role_policy)

        return fn(*args, **kwargs)

    return mock_wrapper


def test_batch_deployment_with_unsupported_flavor_raises_exception(pretrained_model):
    unsupported_flavor = "this is not a valid flavor"
    match = "The specified flavor: `this is not a valid flavor` is not supported for deployment"
    with pytest.raises(MlflowException, match=match) as exc:
        mfs.deploy_transform_job(
            job_name="bad_flavor",
            model_uri=pretrained_model.model_uri,
            s3_input_data_type="Some Data Type",
            s3_input_uri="Some Input Uri",
            content_type="Some Content Type",
            s3_output_path="Some Output Path",
            flavor=unsupported_flavor,
        )

    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_batch_deployment_of_model_with_no_supported_flavors_raises_exception(pretrained_model):
    logged_model_path = _download_artifact_from_uri(pretrained_model.model_uri)
    model_config_path = os.path.join(logged_model_path, "MLmodel")
    model_config = Model.load(model_config_path)
    del model_config.flavors[mlflow.pyfunc.FLAVOR_NAME]
    model_config.save(path=model_config_path)

    match = "The specified model does not contain any of the supported flavors for deployment"
    with pytest.raises(MlflowException, match=match) as exc:
        mfs.deploy_transform_job(
            job_name="missing-flavor",
            model_uri=logged_model_path,
            s3_input_data_type="Some Data Type",
            s3_input_uri="Some Input Uri",
            content_type="Some Content Type",
            s3_output_path="Some Output Path",
            flavor=None,
        )

    assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)


def test_deploy_sagemaker_transform_job_in_asynchronous_mode_without_archiving_throws_exception(
    pretrained_model,
):
    with pytest.raises(MlflowException, match="Resources must be archived") as exc:
        mfs.deploy_transform_job(
            job_name="test-job",
            model_uri=pretrained_model.model_uri,
            s3_input_data_type="Some Data Type",
            s3_input_uri="Some Input Uri",
            content_type="Some Content Type",
            s3_output_path="Some Output Path",
            archive=False,
            synchronous=False,
        )

    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


@mock_sagemaker_aws_services
def test_deploy_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_local(
    pretrained_model, sagemaker_client
):
    job_name = "test-job"
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
        archive=True,
    )

    region_name = sagemaker_client.meta.region_name
    s3_client = boto3.client("s3", region_name=region_name)
    default_bucket = mfs._get_default_s3_bucket(region_name)
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    model_name = transform_job_description["ModelName"]
    assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
    object_names = [
        entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
    ]
    assert any(model_name in object_name for object_name in object_names)
    assert job_name in [
        transform_job["TransformJobName"]
        for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
    ]


@mock_sagemaker_aws_services
def test_deploy_cli_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_local(
    pretrained_model, sagemaker_client
):
    job_name = "test-job"
    result = CliRunner(env={"LC_ALL": "en_US.UTF-8", "LANG": "en_US.UTF-8"}).invoke(
        mfscli.commands,
        [
            "deploy-transform-job",
            "--job-name",
            job_name,
            "--model-uri",
            pretrained_model.model_uri,
            "--input-data-type",
            "Some Data Type",
            "--input-uri",
            "Some Input Uri",
            "--content-type",
            "Some Content Type",
            "--output-path",
            "Some Output Path",
            "--archive",
        ],
    )
    assert result.exit_code == 0

    region_name = sagemaker_client.meta.region_name
    s3_client = boto3.client("s3", region_name=region_name)
    default_bucket = mfs._get_default_s3_bucket(region_name)
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    model_name = transform_job_description["ModelName"]
    assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
    object_names = [
        entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
    ]
    assert any(model_name in object_name for object_name in object_names)
    assert job_name in [
        transform_job["TransformJobName"]
        for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
    ]


@mock_sagemaker_aws_services
def test_deploy_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_s3(
    pretrained_model, sagemaker_client
):
    local_model_path = _download_artifact_from_uri(pretrained_model.model_uri)
    artifact_path = "model"
    region_name = sagemaker_client.meta.region_name
    default_bucket = mfs._get_default_s3_bucket(region_name)
    s3_artifact_repo = S3ArtifactRepository(f"s3://{default_bucket}")
    s3_artifact_repo.log_artifacts(local_model_path, artifact_path=artifact_path)
    model_s3_uri = f"s3://{default_bucket}/{pretrained_model.model_path}"

    job_name = "test-job"
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=model_s3_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
        archive=True,
    )

    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    model_name = transform_job_description["ModelName"]
    assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]

    s3_client = boto3.client("s3", region_name=region_name)
    object_names = [
        entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
    ]
    assert any(model_name in object_name for object_name in object_names)
    assert job_name in [
        transform_job["TransformJobName"]
        for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
    ]


@mock_sagemaker_aws_services
def test_deploy_cli_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_s3(
    pretrained_model, sagemaker_client
):
    local_model_path = _download_artifact_from_uri(pretrained_model.model_uri)
    artifact_path = "model"
    region_name = sagemaker_client.meta.region_name
    default_bucket = mfs._get_default_s3_bucket(region_name)
    s3_artifact_repo = S3ArtifactRepository(f"s3://{default_bucket}")
    s3_artifact_repo.log_artifacts(local_model_path, artifact_path=artifact_path)
    model_s3_uri = f"s3://{default_bucket}/{pretrained_model.model_path}"

    job_name = "test-job"
    result = CliRunner(env={"LC_ALL": "en_US.UTF-8", "LANG": "en_US.UTF-8"}).invoke(
        mfscli.commands,
        [
            "deploy-transform-job",
            "--job-name",
            job_name,
            "--model-uri",
            model_s3_uri,
            "--input-data-type",
            "Some Data Type",
            "--input-uri",
            "Some Input Uri",
            "--content-type",
            "Some Content Type",
            "--output-path",
            "Some Output Path",
            "--archive",
        ],
    )
    assert result.exit_code == 0

    region_name = sagemaker_client.meta.region_name
    s3_client = boto3.client("s3", region_name=region_name)
    default_bucket = mfs._get_default_s3_bucket(region_name)
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    model_name = transform_job_description["ModelName"]
    assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
    object_names = [
        entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
    ]
    assert any(model_name in object_name for object_name in object_names)
    assert job_name in [
        transform_job["TransformJobName"]
        for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
    ]


@mock_sagemaker_aws_services
def test_deploying_sagemaker_transform_job_with_preexisting_name_in_create_mode_throws_exception(
    pretrained_model,
):
    job_name = "test-job"
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
    )

    with pytest.raises(
        MlflowException, match="a batch transform job with the same name already exists"
    ) as exc:
        mfs.deploy_transform_job(
            job_name=job_name,
            model_uri=pretrained_model.model_uri,
            s3_input_data_type="Some Data Type",
            s3_input_uri="Some Input Uri",
            content_type="Some Content Type",
            s3_output_path="Some Output Path",
        )

    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


@mock_sagemaker_aws_services
def test_deploy_in_synchronous_mode_waits_for_transform_job_creation_to_complete_before_returning(
    pretrained_model, sagemaker_client
):
    transform_job_creation_latency = 10
    get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
        transform_job_creation_latency
    )

    job_name = "test-job"
    deployment_start_time = time.time()
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
        synchronous=True,
    )
    deployment_end_time = time.time()

    assert (deployment_end_time - deployment_start_time) >= transform_job_creation_latency
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_COMPLETED


@mock_sagemaker_aws_services
def test_deploy_create_in_asynchronous_mode_returns_before_transform_job_creation_completes(
    pretrained_model, sagemaker_client
):
    transform_job_creation_latency = 10
    get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
        transform_job_creation_latency
    )

    job_name = "test-job"
    deployment_start_time = time.time()
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
        archive=True,
        synchronous=False,
    )
    deployment_end_time = time.time()

    assert (deployment_end_time - deployment_start_time) < transform_job_creation_latency
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_IN_PROGRESS


@mock_sagemaker_aws_services
def test_deploy_in_throw_exception_after_transform_job_creation_fails(
    pretrained_model, sagemaker_client
):
    transform_job_creation_latency = 10
    sagemaker_backend = get_sagemaker_backend(sagemaker_client.meta.region_name)
    sagemaker_backend.set_transform_job_update_latency(transform_job_creation_latency)

    boto_caller = botocore.client.BaseClient._make_api_call

    def fail_transform_job_creations(self, operation_name, operation_kwargs):
        """
        Processes all boto3 client operations according to the following rules:
        - If the operation is a transform job creation, create the transform job and
          set its status to ``TransformJob.STATUS_FAILED``.
        - Else, execute the client operation as normal
        """
        result = boto_caller(self, operation_name, operation_kwargs)
        if operation_name == "CreateTransformJob":
            transform_job_name = operation_kwargs["TransformJobName"]
            sagemaker_backend.set_transform_job_latest_operation(
                transform_job_name=transform_job_name,
                operation=TransformJobOperation.create_unsuccessful(
                    latency_seconds=transform_job_creation_latency
                ),
            )
        return result

    with (
        mock.patch("botocore.client.BaseClient._make_api_call", new=fail_transform_job_creations),
        pytest.raises(MlflowException, match="batch transform job failed") as exc,
    ):
        mfs.deploy_transform_job(
            job_name="test-job",
            model_uri=pretrained_model.model_uri,
            s3_input_data_type="Some Data Type",
            s3_input_uri="Some Input Uri",
            content_type="Some Content Type",
            s3_output_path="Some Output Path",
        )

    assert exc.value.error_code == ErrorCode.Name(INTERNAL_ERROR)


@mock_sagemaker_aws_services
def test_attempting_to_terminate_in_asynchronous_mode_without_archiving_throws_exception(
    pretrained_model,
):
    job_name = "test-job"
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
    )

    with pytest.raises(MlflowException, match="Resources must be archived") as exc:
        mfs.terminate_transform_job(
            job_name=job_name,
            archive=False,
            synchronous=False,
        )

    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


@mock_sagemaker_aws_services
def test_terminate_in_sync_mode_waits_for_transform_job_termination_to_complete_before_returning(
    pretrained_model, sagemaker_client
):
    transform_job_termination_latency = 10
    get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
        transform_job_termination_latency
    )

    job_name = "test-job"
    termination_start_time = time.time()
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
        archive=True,
        synchronous=True,
    )

    mfs.terminate_transform_job(job_name=job_name, synchronous=True)
    termination_end_time = time.time()

    assert (termination_end_time - termination_start_time) >= transform_job_termination_latency
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_STOPPED


@mock_sagemaker_aws_services
def test_terminate_in_asynchronous_mode_returns_before_transform_job_termination_completes(
    pretrained_model, sagemaker_client
):
    transform_job_termination_latency = 10
    get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
        transform_job_termination_latency
    )

    job_name = "test-job"
    termination_start_time = time.time()
    mfs.deploy_transform_job(
        job_name=job_name,
        model_uri=pretrained_model.model_uri,
        s3_input_data_type="Some Data Type",
        s3_input_uri="Some Input Uri",
        content_type="Some Content Type",
        s3_output_path="Some Output Path",
        archive=True,
        synchronous=False,
    )

    mfs.terminate_transform_job(job_name=job_name, archive=True, synchronous=False)
    termination_end_time = time.time()

    assert (termination_end_time - termination_start_time) < transform_job_termination_latency
    transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
    assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_STOPPING