import threading
import time

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError

from turbofuzzllm.llm.model import Model
from turbofuzzllm.utils.cli_utils import cli
from turbofuzzllm.utils.logging import LOGGER


class BedrockModel(Model):
    def __init__(self, model_id, api_key="-"):
        super().__init__(model_id, api_key)
        # we will need a lock to synchronize creation of boto3 session and clients.
        # note that boto3 clients are thread-safe, but creation of the session and client objects are not
        # see example: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html#general-example
        self.lock = threading.Lock()
        self.client_rt = None
        self.client = None
        self.session = None
        # Guardrail config to use.
        self.guardrail_config = None

        self.setup()

    def setup(self):
        # synchronize creation of boto3 session and client objects
        with self.lock:
            # Check model path is a valid model id
            LOGGER.debug(f"\tSetting up client for model with id: {self.model_id}")
            self.cleanup()

            self.session = boto3.Session(profile_name="default", region_name=cli().region)
            self.client = self.session.client("bedrock")
            self.client_rt = self.session.client("bedrock-runtime", config=Config(read_timeout=300))

    def generate(self, message, comment):
        system_prompts = self.get_system_prompts()
        messages = self.get_messages(message)

        try:
            return self.get_llm_response(system_prompts, messages, comment)
        except ClientError as e:
            LOGGER.debug(f"Bedrock API call failed due to {e}. Retrying ...")
            if e.response["Error"]["Code"] == "ThrottlingException":
                time.sleep(60)
            self.setup()
            return self.get_llm_response(system_prompts, messages, comment)
        except Exception as e:
            assert False, f"Bedrock API call failed due to {e}."

    def send_request(self, system_prompts, messages):
        # synchronize access self.client_rt attribute object by each thread
        with self.lock:
            client_rt = self.client_rt

        # LOGGER.debug(f"== TurboFuzzLLM: SYSTEM PROMPT ==\n{system_prompts}\n==")
        # LOGGER.debug(f"== TurboFuzzLLM: MESSAGE ==\n{messages}\n==")

        # Send the message
        if self.guardrail_config is None:
            response = client_rt.converse(
                modelId=self.model_id,
                messages=messages,
                system=system_prompts,
                inferenceConfig=self.inference_config,
                additionalModelRequestFields=self.additional_model_fields,
            )
        else:
            response = client_rt.converse(
                modelId=self.model_id,
                messages=messages,
                system=system_prompts,
                inferenceConfig=self.inference_config,
                additionalModelRequestFields=self.additional_model_fields,
                guardrailConfig=self.guardrail_config,
            )

        LOGGER.debug(f"== TurboFuzzLLM: RESPONSE ==")
        # LOGGER.debug(f"Text: {self.get_text_response(response)}")
        LOGGER.debug(f"Usage: {response['usage']}")
        LOGGER.debug(f"StopReason: {response['stopReason']}")
        LOGGER.debug(f"Metrics: {response['metrics']}")
        LOGGER.debug("==")

        if response["stopReason"] == "guardrail_intervened":
            LOGGER.debug("\t\tBedrock Guardrail: BLOCKED")
            trace = response["trace"]
            LOGGER.debug(f"== Guardrail trace ==\n{trace['guardrail']}\n==")

        result = {"text": self.get_text_response(response), "status": response["stopReason"]}
        return result

    def format_as_system_prompt(self, prompt):
        return [{"text": prompt}]

    def format_as_user_prompt(self, prompt):
        return [{"role": "user", "content": [{"text": prompt}]}]

    @staticmethod
    def get_text_response(response):
        # Extract the response text.
        response_text = response["output"]["message"]["content"][0]["text"]
        assert isinstance(response_text, str), f"response: {response}"
        return response_text
