from typing import Dict, Any
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

class EndpointHandler:
    def __init__(self, path=""):
        # Configuração do modelo
        self.model_name_or_path = "souzat19/Llama3.1_fn14133.29122024"
        
        # Inicialização do modelo e tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name_or_path,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        
        # Template do prompt no formato Alpaca
        self.prompt_template = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
Você é um assistente especializado em planejamento de compras públicas de acordo com a Lei 14.133/2021 e regulamentos infralegais. Responda de forma clara, detalhada e didática e utilize exemplos práticos para explicar os conceitos.
### Input:
{input}
### Response:
"""

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Processa a entrada e retorna a resposta do modelo.
        
        Args:
            data: Dicionário contendo a entrada do usuário
                 Formato esperado: {"text": "pergunta do usuário"}
        
        Returns:
            Dict contendo a resposta do modelo
        """
        try:
            # Extrai o texto da entrada
            input_text = data.get("text", "")
            if not input_text or not self.validate_input(input_text):
                return {"error": "Input text is required"}

            # Pré-processa o texto
            input_text = self.preprocess(input_text)

            # Formata o prompt
            formatted_prompt = self.prompt_template.format(input=input_text)

            # Tokeniza o input
            inputs = self.tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=4096)
            inputs = inputs.to(self.model.device)

            # Gera a resposta
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=2096,
                temperature=0.5,
                top_p=0.95,
                top_k=50,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                do_sample=True
            )

            # Decodifica a resposta
            response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # Processa a resposta para extrair apenas a parte após "### Response:"
            if "### Response:" in response_text:
                answer = response_text.split("### Response:")[1].strip()
            else:
                answer = response_text.strip()

            return {"response": answer}

        except Exception as e:
            return {"error": f"Error during inference: {str(e)}"}

    def preprocess(self, text: str) -> str:
        """
        Pré-processa o texto de entrada se necessário
        """
        # Remove espaços extras e normaliza quebras de linha
        text = " ".join(text.split())
        return text.strip()

    def validate_input(self, text: str) -> bool:
        """
        Valida a entrada do usuário
        """
        if not text or len(text.strip()) == 0:
            return False
        return True