import threading

from transformers import AutoModelForCausalLM, AutoTokenizer

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


# Uncomment this class to use for Gemma models
class LocalModel(Model):
    def __init__(self, model_id, api_key="-"):
        super().__init__(model_id, api_key)
        self.lock = threading.Lock()
        # self.pipe = pipeline("text-generation", model=self.model_id, torch_dtype=torch.bfloat16, device_map="auto")

        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")

    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 Exception as e:
            LOGGER.debug(f"== LocalModel call failed due to {e}.\n==")
            return {"text": "<bos>This is a harmless response"}

    def send_request(self, system_prompts, messages):

        inputs = self.tokenizer(messages[0]["content"], return_tensors="pt")

        device = next(self.model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        outputs = self.model.generate(
            **inputs,
            pad_token_id=self.tokenizer.eos_token_id,
            do_sample=True,
            max_new_tokens=self.inference_config["maxTokens"],
            temperature=self.inference_config["temperature"],
        )  # Adjust max_new_tokens as needed

        result = {"text": self.tokenizer.decode(outputs[0])}
        return result

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

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

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


# Uncomment this class to use for Zephyr models
# class LocalModel(Model):
#     def __init__(self, model_id, api_key="-"):
#         super().__init__(model_id, api_key)
#         self.lock = threading.Lock()
#         self.pipe = pipeline("text-generation", model=self.model_id, torch_dtype=torch.bfloat16, device_map="auto")

#     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 Exception as e:
#             assert False, f"LocalModel 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:
#             pipe = self.pipe

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

#         response = pipe(messages, pad_token_id=pipe.tokenizer.eos_token_id, do_sample=True, max_new_tokens=self.inference_config["maxTokens"], temperature=self.inference_config["temperature"])
#         # LOGGER.debug(f"== TurboFuzzLLM: RESPONSE ==\n{response}\n==")

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

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

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

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