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

class EndpointHandler:
    def __init__(self, path=""):
        """Initialize the model and tokenizer.
        
        Args:
            path (str): Path to the model directory. Defaults to empty string.
        """
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = AutoModelForCausalLM.from_pretrained(
            path or "merged",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(path or "merged")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle inference requests.
        
        Args:
            data (Dict[str, Any]): The input data. Can be in two formats:
                1. Standard Hugging Face format:
                    {
                        "inputs": str,
                        "parameters": Dict[str, Any]
                    }
                2. Custom format:
                    {
                        "instruction": str,
                        "input": str (optional),
                        "max_new_tokens": int (optional),
                        "temperature": float (optional)
                    }
        
        Returns:
            Dict[str, Any]: The model's response containing:
                - response (str): The generated text
        """
        # Handle standard Hugging Face format
        if "inputs" in data:
            instruction = data["inputs"]
            parameters = data.get("parameters", {})
            input_text = ""
            max_new_tokens = parameters.get("max_new_tokens", 512)
            temperature = parameters.get("temperature", 0.7)
        # Handle custom format
        else:
            instruction = data.get("instruction", "")
            input_text = data.get("input", "")
            max_new_tokens = data.get("max_new_tokens", 512)
            temperature = data.get("temperature", 0.7)

        # Create prompt
        prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{instruction}"""

        if input_text:
            prompt += f"""

### Input:
{input_text}"""

        prompt += """

### Response:"""

        # Generate response
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True,
            pad_token_id=self.tokenizer.eos_token_id
        )
        
        # Decode and extract response
        full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = full_response.split("### Response:")[-1].strip()

        return {"response": response}