# handler.py

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

class EndpointHandler:
    def __init__(self, path: str = ""):
        """
        Load model and tokenizer at startup.
        `path` is the model repo on the Hub or local directory.
        """
        # Determine device
        self.device = 0 if torch.cuda.is_available() else -1

        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(path).to(self.device)

        # Set up a text-generation pipeline
        self.generator = pipeline(
            task="text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            device=self.device
        )

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Handle each inference request. Must return a non-empty list.
        `data` will always contain at least the key "inputs".
        """
        # Extract input messages or fall back to raw prompt
        raw = data.get("inputs", data)
        prompt = ""
        # Support both string inputs and OpenAI-style messages
        if isinstance(raw, list):
            for msg in raw:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                prompt += f"{role}: {content}\n"
        else:
            prompt = str(raw)

        # Run generation (returns a list of dicts)
        outputs = self.generator(prompt, max_new_tokens=128)

        # Ensure we output a list of {"generated_text": ...}
        return [{"response": out["response"]} for out in outputs]