"""SageMaker component for ground truth."""
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from typing import Dict

from ground_truth.src.sagemaker_ground_truth_spec import (
    SageMakerGroundTruthSpec,
    SageMakerGroundTruthInputs,
    SageMakerGroundTruthOutputs,
)
from common.sagemaker_component import (
    SageMakerComponent,
    ComponentMetadata,
    SageMakerJobStatus,
)


@ComponentMetadata(
    name="SageMaker - Ground Truth",
    description="Ground Truth Jobs in SageMaker",
    spec=SageMakerGroundTruthSpec,
)
class SageMakerGroundTruthComponent(SageMakerComponent):
    """SageMaker component for ground truth."""

    def Do(self, spec: SageMakerGroundTruthSpec):
        self._labeling_job_name = (
            spec.inputs.job_name
            if spec.inputs.job_name
            else self._generate_unique_timestamped_id(prefix="LabelingJob")
        )
        super().Do(spec.inputs, spec.outputs, spec.output_paths)

    def _get_job_status(self) -> SageMakerJobStatus:
        response = self._sm_client.describe_labeling_job(
            LabelingJobName=self._labeling_job_name
        )
        status = response["LabelingJobStatus"]

        if status == "Completed":
            return SageMakerJobStatus(
                is_completed=True, has_error=False, raw_status=status
            )
        if status == "Failed":
            message = response["FailureReason"]
            return SageMakerJobStatus(
                is_completed=True,
                has_error=True,
                error_message=message,
                raw_status=status,
            )

        return SageMakerJobStatus(is_completed=False, raw_status=status)

    def _get_labeling_job_outputs(self, auto_labeling):
        """Get labeling job output values.

        Args:
            auto_labeling: Determines whether to return a active learning model ARN.

        Returns:
            tuple: (Output manifest S3 URI, active learning model ARN)
        """
        ### Get and return labeling job outputs
        info = self._sm_client.describe_labeling_job(
            LabelingJobName=self._labeling_job_name
        )
        output_manifest = info["LabelingJobOutput"]["OutputDatasetS3Uri"]
        if auto_labeling:
            active_learning_model_arn = info["LabelingJobOutput"][
                "FinalActiveLearningModelArn"
            ]
        else:
            active_learning_model_arn = " "
        return output_manifest, active_learning_model_arn

    def _after_job_complete(
        self,
        job: object,
        request: Dict,
        inputs: SageMakerGroundTruthInputs,
        outputs: SageMakerGroundTruthOutputs,
    ):
        (
            outputs.output_manifest_location,
            outputs.active_learning_model_arn,
        ) = self._get_labeling_job_outputs(inputs.enable_auto_labeling)

    def _on_job_terminated(self):
        self._sm_client.stop_labeling_job(LabelingJobName=self._labeling_job_name)

    def _create_job_request(
        self, inputs: SageMakerGroundTruthInputs, outputs: SageMakerGroundTruthOutputs,
    ) -> Dict:
        ### Documentation: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_labeling_job
        request = self._get_request_template("gt")

        # Mapping are extracted from ARNs listed in https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_labeling_job
        algorithm_arn_map = {
            "us-west-2": "081040173940",
            "us-east-1": "432418664414",
            "us-east-2": "266458841044",
            "ca-central-1": "918755190332",
            "eu-west-1": "568282634449",
            "eu-west-2": "487402164563",
            "eu-central-1": "203001061592",
            "ap-northeast-1": "477331159723",
            "ap-northeast-2": "845288260483",
            "ap-south-1": "565803892007",
            "ap-southeast-1": "377565633583",
            "ap-southeast-2": "454466003867",
        }

        task_map = {
            "bounding box": "BoundingBox",
            "image classification": "ImageMultiClass",
            "semantic segmentation": "SemanticSegmentation",
            "text classification": "TextMultiClass",
        }

        auto_labeling_map = {
            "bounding box": "object-detection",
            "image classification": "image-classification",
            "text classification": "text-classification",
        }

        task = inputs.task_type.lower()

        request["LabelingJobName"] = self._labeling_job_name

        if inputs.label_attribute_name: