"""SageMaker component for RLEstimator."""
# 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
import os
from sagemaker.rl import RLEstimator, RLToolkit, RLFramework
from rlestimator.src.sagemaker_rlestimator_spec import (
    SageMakerRLEstimatorSpec,
    SageMakerRLEstimatorInputs,
    SageMakerRLEstimatorOutputs,
)
from common.sagemaker_component import (
    SageMakerComponent,
    ComponentMetadata,
    SageMakerJobStatus,
    DebugRulesStatus,
)
from common.boto3_manager import Boto3Manager
from common.common_inputs import SageMakerComponentCommonInputs
from common.spec_input_parsers import SpecInputParsers


@ComponentMetadata(
    name="SageMaker - RLEstimator Training Job",
    description="Handle end-to-end training and deployment of custom RLEstimator code.",
    spec=SageMakerRLEstimatorSpec,
)
class SageMakerRLEstimatorComponent(SageMakerComponent):
    """SageMaker component for RLEstimator."""

    def Do(self, spec: SageMakerRLEstimatorSpec):
        self._rlestimator_job_name = (
            spec.inputs.job_name
            if spec.inputs.job_name
            else SageMakerComponent._generate_unique_timestamped_id(
                prefix="RLEstimatorJob"
            )
        )
        super().Do(spec.inputs, spec.outputs, spec.output_paths)

    def _configure_aws_clients(self, inputs: SageMakerComponentCommonInputs):
        """Configures the internal AWS clients for the component.

        Args:
            inputs: A populated list of user inputs.
        """
        self._sm_client = Boto3Manager.get_sagemaker_client(
            self._get_component_version(),
            inputs.region,
            endpoint_url=inputs.endpoint_url,
            assume_role_arn=inputs.assume_role,
        )
        self._cw_client = Boto3Manager.get_cloudwatch_client(
            inputs.region, assume_role_arn=inputs.assume_role
        )
        self._sagemaker_session = Boto3Manager.get_sagemaker_session(
            self._get_component_version(),
            inputs.region,
            assume_role_arn=inputs.assume_role,
        )

    def _get_job_status(self) -> SageMakerJobStatus:
        response = self._sm_client.describe_training_job(
            TrainingJobName=self._rlestimator_job_name
        )
        status = response["TrainingJobStatus"]

        if status == "Completed" or status == "Stopped":
            return self._get_debug_rule_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_debug_rule_status(self) -> SageMakerJobStatus:
        """Get the job status of the training debugging rules.

        Returns:
            SageMakerJobStatus: A status object.
        """
        response = self._sm_client.describe_training_job(
            TrainingJobName=self._rlestimator_job_name
        )

        # No debugging configured
        if "DebugRuleEvaluationStatuses" not in response:
            return SageMakerJobStatus(is_completed=True, has_error=False, raw_status="")

        raw_status = DebugRulesStatus.from_describe(response)
        if raw_status != DebugRulesStatus.INPROGRESS:
            logging.info("Rules have ended with status:\n")
            self._print_debug_rule_status(response, True)
            return SageMakerJobStatus(
                is_completed=True,
                has_error=(raw_status == DebugRulesStatus.ERRORED),
                raw_status=raw_status,
            )

        self._print_debug_rule_status(response)
        return SageMakerJobStatus(is_completed=False, raw_status=raw_status)

    def _print_debug_rule_status(self, response, last_print=False):
        """Prints the status of each debug rule.

            Example of DebugRuleEvaluationStatuses:
            response['DebugRuleEvaluationStatuses'] =
                [{
                    "RuleConfigurationName": "VanishingGradient",
                    "RuleEvaluationStatus": "IssuesFound",
                    "StatusDetails": "There was an issue."
                }]
            If last_print is False:
            INFO:root: - LossNotDecreasing: InProgress
            INFO:root: - Overtraining: NoIssuesFound
            ERROR:root:- CustomGradientRule: Error
            If last_print is True:
            INFO:root: - LossNotDecreasing: IssuesFound
            INFO:root:   - RuleEvaluationConditionMet: Evaluation of the rule LossNotDecreasing at step 10 resulted in the condition being met

        Args:
            response: A describe training job response.
            last_print: If true, prints each of the debug rule issues if found.
        """
        for debug_rule in response["DebugRuleEvaluationStatuses"]: