# handler.py
import os
from typing import Dict, Any, List, Union
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel
import torch

class EndpointHandler:
    def __init__(self, path: str = ""):
        # Load tokenizer and base model
        model_id = path or os.getenv("HF_MODEL_ID")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
        base_model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="auto",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        )
        # Load LoRA adapter
        self.model = PeftModel.from_pretrained(base_model, model_id)
        self.model.eval()
        self.model.config.use_cache = True

        # Create text-generation pipeline
        self.pipe = pipeline(
            "text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            device_map="auto",
            return_full_text=False,
        )

    def __call__(
        self,
        request: Union[Dict[str, Any], List[Dict[str, Any]]]
    ) -> List[Dict[str, str]]:
        # Normalize to list of request dicts
        reqs = request if isinstance(request, list) else [request]
        responses: List[Dict[str, str]] = []

        for req in reqs:
            # Support 'inputs' or 'prompt'
            raw = req.get("inputs") or req.get("prompt")
            if raw is None:
                raise ValueError("No 'inputs' or 'prompt' field in request")
            # Normalize to list of strings
            texts: List[str] = raw if isinstance(raw, list) else [raw]

            params = req.get("parameters", {})
            # Ensure pad_token_id set
            if "pad_token_id" not in params:
                params["pad_token_id"] = self.tokenizer.eos_token_id

            for text in texts:
                # Generate
                out = self.pipe(text, **params)
                # pipeline returns list of dicts [{'generated_text': ...}]
                if isinstance(out, list) and out and isinstance(out[0], dict):
                    responses.append({"generated_text": out[0]["generated_text"].strip()})
                else:
                    # Fallback: stringify
                    responses.append({"generated_text": str(out).strip()})
        return responses