import json

from turbofuzzllm.llm.bedrock_model import BedrockModel
from turbofuzzllm.llm.local_model import LocalModel
from turbofuzzllm.llm.model import Model
from turbofuzzllm.llm.model_utils import (
    is_bedrock_model,
    is_local_model,
    is_openai_model,
    is_sagemaker_huggingface_model,
    is_sagemaker_jumpstart_model,
)
from turbofuzzllm.llm.openai_model import OpenAIModel
from turbofuzzllm.llm.sagemaker_huggingface_model import SagemakerHuggingFaceModel
from turbofuzzllm.llm.sagemaker_jumpstart_model import SagemakerJumpStartModel
from turbofuzzllm.utils.cli_utils import cli
from turbofuzzllm.utils.logging import LOGGER


class TargetModel(Model):
    def __init__(self, model_id, api_key, guardrail_id="-", guardrail_version="-"):
        if is_openai_model(model_id):
            self.__class__ = type(
                self.__class__.__name__, (OpenAIModel, Model), dict(self.__class__.__dict__)
            )
            OpenAIModel.__init__(self, model_id, api_key)
        elif is_sagemaker_jumpstart_model(model_id):
            self.__class__ = type(
                self.__class__.__name__,
                (SagemakerJumpStartModel, Model),
                dict(self.__class__.__dict__),
            )
            SagemakerJumpStartModel.__init__(self, model_id, api_key)
        elif is_sagemaker_huggingface_model(model_id):
            self.__class__ = type(
                self.__class__.__name__,
                (SagemakerHuggingFaceModel, Model),
                dict(self.__class__.__dict__),
            )
            hub = {
                "HF_MODEL_ID": model_id,
                "HF_TASK": "question-answering",
                "SM_NUM_GPUS": json.dumps(1),
                "device_map": "auto",
                "SAGEMAKER_MODEL_SERVER_TIMEOUT": json.dumps(300),
                "SAGEMAKER_TS_RESPONSE_TIMEOUT": json.dumps(300),
            }
            SagemakerHuggingFaceModel.__init__(
                self, model_id, api_key, hub, cli().sagemaker_target_instance_type
            )
        elif is_local_model(model_id):
            self.__class__ = type(
                self.__class__.__name__, (LocalModel, Model), dict(self.__class__.__dict__)
            )
            LocalModel.__init__(self, model_id, api_key)
        else:
            assert is_bedrock_model(model_id)
            self.__class__ = type(
                self.__class__.__name__, (BedrockModel, Model), dict(self.__class__.__dict__)
            )
            BedrockModel.__init__(self, model_id, api_key)
        if (guardrail_id != "-") and (guardrail_version != "-"):
            self.guardrail_config = {
                "guardrailIdentifier": guardrail_id,
                "guardrailVersion": guardrail_version,
                "trace": "enabled",
            }

    def get_llm_response(self, system_prompts, messages, comment):
        response = self.send_request(system_prompts, messages)
        LOGGER.debug(f"== TurboFuzzLLM: TargetLLM Response for {comment} ==\n{response}\n==")
        return response
