# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

import json
from abc import ABC, abstractmethod
from typing import Optional

from agenteval.conversation import Conversation
from agenteval.evaluators.model_config.bedrock_model_config import BedrockModelConfig
from agenteval.hook import Hook
from agenteval.targets import BaseTarget
from agenteval.test import Test, TestResult
from agenteval.trace import Trace
from agenteval.utils import create_boto3_client, import_class

_BOTO3_SERVICE_NAME = "bedrock-runtime"


class BaseEvaluator(ABC):
    """The `BaseEvaluator` abstract base class defines the common interface for evaluator
    classes.

    Attributes:
        test (Test): The test case.
        target (BaseTarget): The target agent being evaluated.
        conversation (Conversation): Captures the interaction between a user and an agent.
        trace (Trace): Captures steps during evaluation.
        test_result (TestResult): The result of the test which is set in `BaseEvaluator.run`.
        input_token_count (int): Number of input tokens processed by the evaluator.
        output_token_count (int): Number of output tokens generated by the evaluator.
        model_config (BedrockModelConfig): A configuration of the bedrock model being used. If `provisioned_throughput_arn` is provided,
            then the model_id will be set to the ARN of the provisioned throughput.
        boto3_client (BaseClient): A `boto3` client representing Amazon Bedrock Runtime.
    """

    def __init__(
        self,
        test: Test,
        target: BaseTarget,
        work_dir: str,
        model_config: BedrockModelConfig,
        provisioned_throughput_arn: Optional[str] = None,
        aws_profile: Optional[str] = None,
        aws_region: Optional[str] = None,
        endpoint_url: Optional[str] = None,
        max_retry: int = 10,
    ):
        """Initialize the evaluator.

        Args:
            test (Test): The test case.
            target (BaseTarget): The target agent being evaluated.
            work_dir (str): The directory where the test result and trace will be
                generated.
            model_config (BedrockModelConfig): The config of the Bedrock model used to run evaluation.
            provisioned_throughput_arn (Optional[str]): The ARN of the provisioned throughput.
            aws_profile (Optional[str]): The AWS profile name.
            aws_region (Optional[str]): The AWS region.
            endpoint_url (Optional[str]): The endpoint URL for the AWS service.
            max_retry (int): The maximum number of retry attempts.
        """
        # overwrite the model_id with the provisioned_throughput_arn if provided, keep the request_config the same.
        if provisioned_throughput_arn:
            model_config = BedrockModelConfig(
                model_id=provisioned_throughput_arn,
                request_body=model_config.request_body,
            )
        self.test = test
        self.target = target
        self.conversation = Conversation()
        self.trace = Trace(work_dir=work_dir, test_name=test.name)
        self.test_result = None
        self.input_token_count = 0
        self.output_token_count = 0
        self.model_config = model_config
        self.bedrock_runtime_client = create_boto3_client(
            boto3_service_name=_BOTO3_SERVICE_NAME,
            aws_profile=aws_profile,
            aws_region=aws_region,
            endpoint_url=endpoint_url,
            max_retry=max_retry,
        )

    @abstractmethod
    def evaluate(self) -> TestResult:
        """Conduct the test.

        Returns:
            TestResult
        """
        pass

    def _get_hook_cls(self, hook: Optional[str]) -> Optional[type[Hook]]:
        if hook:
            hook_cls = import_class(hook, parent_class=Hook)
            return hook_cls

    def invoke_model(self, request_body: dict) -> dict:
        """
        Invoke the Bedrock model using the `boto3_client`. This method will convert
        a request dictionary to a JSON string before passing it to the `InvokeModel` API.

        Refer to the `boto3` documentation for more details.

        Args:
            request_body (dict): The request payload as a dictionary.

        Returns:
            dict: The response from the model invocation.

        """
        response = self.bedrock_runtime_client.invoke_model(
            modelId=self.model_config.model_id, body=json.dumps(request_body)
        )

        self._incr_token_counts(response)

        return response

    def _incr_token_counts(self, response: dict):
        headers = response["ResponseMetadata"]["HTTPHeaders"]

        self.input_token_count += int(
            headers.get("x-amzn-bedrock-input-token-count", 0)
        )
        self.output_token_count += int(
            headers.get("x-amzn-bedrock-output-token-count", 0)
        )

    def run(self) -> TestResult:
        """
        Run the evaluator within a trace context manager and run hooks
        if provided.

        Returns:
            TestResult
        """

        hook_cls = self._get_hook_cls(self.test.hook)

        with self.trace:
            if hook_cls:
                hook_cls.pre_evaluate(self.test, self.trace)
            self.test_result = self.evaluate()
            if hook_cls:
                hook_cls.post_evaluate(self.test, self.test_result, self.trace)

        return self.test_result
