import io
import json
import threading
import time

import boto3
import sagemaker
from botocore.exceptions import ClientError
from sagemaker.jumpstart.model import JumpStartModel

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


class SagemakerJumpStartModel(Model):
    def __init__(self, model_id, api_key="-"):
        super().__init__(model_id, api_key)
        self.lock = threading.Lock()
        self.model = None
        self.predictor = None

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

        self.setup()

    def cleanup(self):
        if self.predictor is not None:
            self.predictor.delete_predictor()
            self.predictor = None

    def setup(self):
        # synchronize creation of boto3 session and client objects
        with self.lock:
            if not is_sagemaker_jumpstart_model(self.model_id):
                raise ValueError(f"SagemakerJumpStartModel model id {self.model_id} not identified")

            # Check model path is a valid model id
            LOGGER.debug(f"\tSetting up sagemaker jumpstart for model with id: {self.model_id}")
            self.cleanup()

            boto_session = boto3.Session(profile_name="default", region_name=cli().region)
            sagemaker_session = sagemaker.Session(boto_session=boto_session)
            self.model = JumpStartModel(
                model_id=self.model_id,
                role=cli().sagemaker_role,
                region=cli().region,
                sagemaker_session=sagemaker_session,
                name=self.model_name,
            )
            self.predictor = self.model.deploy(accept_eula=True, endpoint_name=self.endpoint_name)
            LOGGER.debug(f"\tSagemaker jumpstart model deployed")

    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 = self.predictor

        # example_payloads = self.model.retrieve_all_examples()
        # for payload in example_payloads:
        #     response = predictor.predict(payload.body)
        #     response = response[0] if isinstance(response, list) else response
        #     print("Input:\n", payload.body, end="\n\n")
        #     print("Output:\n", response["generated_text"].strip(), end="\n\n\n")

        # 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):
        return {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": self.inference_config["maxTokens"],
                "return_full_text": False,
            },
        }
