import io
import json
import threading
import time

import boto3
import sagemaker
from botocore.config import Config
from botocore.exceptions import ClientError
from sagemaker.huggingface import (
    HuggingFaceModel,
    HuggingFacePredictor,
    get_huggingface_llm_image_uri,
)

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


class SagemakerHuggingFaceModel(Model):
    def __init__(self, model_id, api_key, hub, instance_type):
        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.hub = hub
        self.instance_type = instance_type

        base_name = (
            self.model_id.replace("/", "-").replace(".", "-").replace("_", "-").replace(":", "-")
        )
        self.model_name = base_name
        self.endpoint_name = f"{base_name}"

        self.client_rt = None
        self.client = None
        self.session = None
        self.endpoint_created = False

        self.setup()

    def endpoint_config_exists(self):
        try:
            self.client.describe_endpoint_config(EndpointConfigName=self.endpoint_name)
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "ValidationException":
                return False
            elif e.response["Error"]["Code"] == "ThrottlingException":
                LOGGER.debug(f"Sagemaker API call failed due to {e}. Retrying ...")
                time.sleep(60)
                return self.endpoint_config_exists()
            else:
                raise

    def endpoint_exists(self, attempt_number=1):
        try:
            self.client.describe_endpoint(EndpointName=self.endpoint_name)
            # describe endpoint creation status
            status = self.client.describe_endpoint(EndpointName=self.endpoint_name)[
                "EndpointStatus"
            ]
            if status == "InService":
                return True
            elif status == "Failed":
                self.delete_endpoint()
                time.sleep(10)
                return False
            else:
                if attempt_number > 3:
                    assert (
                        False
                    ), f"Status stuck at {status} for sagemaker endpoint {self.endpoint_name}"
                else:
                    LOGGER.debug(
                        f"\tAttempt {attempt_number}: waiting for sagemaker endpoint: {self.endpoint_name} (status: {status})"
                    )
                    time.sleep(300)
                    return self.endpoint_exists(attempt_number + 1)
        except ClientError as e:
            if e.response["Error"]["Code"] == "ValidationException":
                return False
            elif e.response["Error"]["Code"] == "ThrottlingException":
                LOGGER.debug(f"Sagemaker API call failed due to {e}. Retrying ...")
                time.sleep(60)
                return self.endpoint_exists(attempt_number + 1)
            else:
                raise

    def create_endpoint(self):
        image_uri = get_huggingface_llm_image_uri(backend="huggingface", region=cli().region)

        # create Hugging Face Model Class
        LOGGER.debug(f"\tSetting up sagemaker model: {self.model_name}")
        model = HuggingFaceModel(
            name=self.model_name,
            image_uri=image_uri,
            env=self.hub,
            role=cli().sagemaker_role,
            sagemaker_session=sagemaker.Session(boto_session=self.session),
        )

        if self.endpoint_config_exists():
            self.client.delete_endpoint_config(EndpointConfigName=self.endpoint_name)

        # deploy model to SageMaker Inference
        LOGGER.debug(f"\tSetting up sagemaker endpoint: {self.endpoint_name}")
        model.deploy(
            initial_instance_count=1,
            instance_type=self.instance_type,
            endpoint_name=self.endpoint_name,
            container_startup_health_check_timeout=1000,
            update_endpoint=True,
        )

        LOGGER.debug(f"\tSagemaker endpoint created")
        self.endpoint_created = True

    def setup_endpoint(self):
        if not self.endpoint_exists():
            self.create_endpoint()
        assert self.endpoint_exists()
        self.client_rt = self.session.client("sagemaker-runtime", config=Config(read_timeout=1000))

    def delete_endpoint(self):
        LOGGER.debug(f"\tDeleting sagemaker endpoint: {self.endpoint_name}")
        self.client.delete_endpoint(EndpointName=self.endpoint_name)
        self.client.delete_endpoint_config(EndpointConfigName=self.endpoint_name)

    def cleanup(self):
        if self.endpoint_created and self.endpoint_exists():
            # self.client.delete_model(ModelName=self.model_name)
            # self.delete_endpoint()
            pass

    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 sagemaker for model with id: {self.model_id}")
            self.session = boto3.Session(profile_name="default", region_name=cli().region)
            self.client = self.session.client("sagemaker")
            self.cleanup()
            self.setup_endpoint()

    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:
            if e.response["Error"]["Code"] == "ThrottlingException":
                LOGGER.debug(f"Sagemaker API call failed due to {e}. Retrying ...")
                time.sleep(60)
            self.setup()
            return self.get_llm_response(system_prompts, message, comment)
        except Exception as e:
            assert False, f"Sagemaker API call failed due to {e}."

    def send_request(self, _, messages):
        # synchronize access self.predictor attribute object by each thread
        with self.lock:
            predictor = HuggingFacePredictor(
                endpoint_name=self.endpoint_name,
                sagemaker_session=sagemaker.Session(boto_session=self.session),
            )

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

        response_full = predictor.predict_stream(messages)
        # LOGGER.debug(f"== TurboFuzzLLM: RESPONSE full ==\n{response_full}\n==")
        event_stream = response_full
        response_list = []
        for event in event_stream:
            response_list.append(event)
        response_bytes = b"".join(response_list)
        # LOGGER.debug(f"== TurboFuzzLLM: RESPONSE bytes ==\n{response_bytes}\n==")
        response = json.load(io.BytesIO(response_bytes))
        # LOGGER.debug(f"== TurboFuzzLLM: RESPONSE ==\n{response}\n==")

        response = response[0] if isinstance(response, list) else response
        generated_text = response["generated_text"].strip()
        result = {"text": generated_text}
        return result

    def format_as_user_prompt(self, prompt):
        if self.model_id.startswith("huggingface-llm-gemma-7b-instruct-finetuned"):
            return {
                "inputs": f"<bos><start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model",
                "parameters": {
                    "max_new_tokens": self.inference_config["maxTokens"],
                    "return_full_text": False,
                },
            }
        return {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": self.inference_config["maxTokens"],
                "return_full_text": False,
            },
        }
