import os

import torch
from langchain_core.output_parsers import StrOutputParser
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace, HuggingFacePipeline


from schema_miner.services.LLM_Inference.LLM_inference import LLM_Inference


class HuggingFace_LLM_Inference(LLM_Inference):
    """
    HuggingFace_LLM_Inference is a subclass of LLM_Inference which has a task of doing inference from the specific Large Language Model(LLM) giving a specified prompt.
    """

    def __init__(self, model_name: str, temperature: float = 0.3):
        """
        Initialize the object with HuggingFace configuration (Model name etc.)
        """
        super().__init__(model_name=model_name, temperature=temperature)

        # HuggingFace Access Token
        assert self.config.HUGGINGFACE_access_token is not None
        self.access_token = self.config.HUGGINGFACE_access_token
        os.environ["HUGGINGFACEHUB_API_TOKEN"] = self.access_token

        if self.config.HUGGINGFACE_USE_LOCAL:
            self.logger.info(f"Using local HuggingFace model: {self.model_name} for LLM Inference")
            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

            if "Ministral-3" in self.model_name:
                from transformers import Mistral3ForConditionalGeneration, MistralCommonBackend
                tokenizer = MistralCommonBackend.from_pretrained(self.model_name, token = self.access_token)
                model = Mistral3ForConditionalGeneration.from_pretrained(self.model_name, token = self.access_token, dtype=torch.bfloat16, device_map="auto")
            else:
                tokenizer = AutoTokenizer.from_pretrained(self.model_name, token = self.access_token)
                model = AutoModelForCausalLM.from_pretrained(self.model_name, token = self.access_token, dtype=torch.bfloat16, device_map="auto")
            
            # Create the text generation pipeline using the loaded model and tokenizer
            pipe = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                trust_remote_code=True,
                max_new_tokens=8192,
                temperature=self.temperature,
                max_length=None,
                do_sample=self.temperature > 0
            )

            # Initialize the HuggingFacePipeline with the created pipeline
            self.model = HuggingFacePipeline(pipeline=pipe)
        else:
            self.logger.info(f"Using HuggingFace Inference API with model: {self.model_name} for LLM Inference")

            # The Large Language Model to use for Inference
            self.model = HuggingFaceEndpoint(
                repo_id=self.model_name,
                task="text-generation",
                max_new_tokens=8192,
                temperature=self.temperature
            )

        # Initialize the Chat Model with the specified HuggingFace model
        self.chat_model = ChatHuggingFace(llm = self.model)

    def __str__(self):
        """
        Returns a human-readable representation of an object
        """
        return f"HuggingFace - LLM Inference with Model: {self.model_name}"

    def completion(self, prompt_template, var_dict):
        """
        Requests the completion endpoint of the HuggingFace with the specified prompt/message. Returns the parsed output from the model.

        :param prompt_template: The prompt template with placeholders for dynamic values.
        :param dict var_dict: The dictionary containing variable name and its corresponding value to format the prompt.
        :returns model_output: The parsed LLM output
        """
        try:
            # Formatting the prompt
            prompt = super().format_prompt_template(prompt_template, var_dict)

            # Invoking the model's completion API with the prompt
            model_output = self.chat_model.invoke(prompt.to_messages())

            # Parsing the LLM's output to extract the final output
            model_output = StrOutputParser().invoke(model_output)

            # Returns the LLM's output
            return model_output
        except Exception as e:
            self.logger.debug(
                f"Exception Occurred while calling the Completion API of model: {self.model_name}"
            )
            self.logger.debug(f"Exception: {e}")
            return None
