import threading

import openai

from turbofuzzllm.llm.model import Model
from turbofuzzllm.llm.model_utils import is_openai_model
from turbofuzzllm.utils.logging import LOGGER


class OpenAIModel(Model):
    def __init__(self, model_id, api_key):
        super().__init__(model_id, api_key)
        self.client = None
        self.lock = threading.Lock()
        self.setup()

    def setup(self):
        # synchronize creation of session and client objects
        with self.lock:
            if not self.api_key.startswith("sk-"):
                raise ValueError("OpenAI API key should start with sk-")
            if not is_openai_model(self.model_id):
                raise ValueError(f"OpenAI model id {self.model_id} not identified")

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

            self.client = openai.OpenAI(api_key=self.api_key)

    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 openai.BadRequestError as e:
            if (
                "Invalid prompt: your prompt was flagged as potentially violating our usage policy."
                in f"{e}"
            ):
                LOGGER.debug(
                    f"== TurboFuzzLLM: TargetLLM prompt rejected for {comment} ==\n{e}\n=="
                )
                result = {
                    "text": "Invalid prompt: your prompt was flagged as potentially violating our usage policy."
                }
                return result
            else:
                LOGGER.debug(
                    f"== TurboFuzzLLM: TargetLLM prompt rejected for {comment} ==\n{e}\n=="
                )
                result = {
                    "text": "Invalid prompt: Not sure why but let's move on to the next one!."
                }
                return result
        except Exception as e:
            LOGGER.debug(f"== TurboFuzzLLM: TargetLLM prompt failed for {comment} ==\n{e}\n==")
            raise

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

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

        openai_messages = system_prompts + messages

        # Send the message
        response = client.chat.completions.create(
            model=self.model_id,
            messages=openai_messages,
            temperature=self.inference_config["temperature"],
            n=1,
        )

        LOGGER.debug(f"== TurboFuzzLLM: RESPONSE ==")
        # LOGGER.debug(f"Text: {self.get_text_response(response)}")

        result = {"text": self.get_text_response(response)}
        return result

    def format_as_system_prompt(self, prompt):
        return [{"role": "system", "content": [{"type": "text", "text": prompt}]}]

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

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