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

class EndpointHandler:
    def __init__(self, path=""):
        base_id = "microsoft/Phi-3-mini-4k-instruct"
        adapter_id = "CarlosMM24/phi3-mini-med-adapter"

        self.tokenizer = AutoTokenizer.from_pretrained(base_id)

        base = AutoModelForCausalLM.from_pretrained(
            base_id,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
        )
        self.model = PeftModel.from_pretrained(base, adapter_id)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        prompt = data.get("inputs", data)

        # 1) Apply chat template so Phi-3 knows this is a user turn and it should respond
        messages = [{"role": "user", "content": prompt}]
        formatted = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,  # IMPORTANT: adds the assistant turn marker
        )

        # 2) Tokenize and move to model device
        inputs = self.tokenizer(formatted, return_tensors="pt")
        device = next(self.model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            out = self.model.generate(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],  # IMPORTANT
                max_new_tokens=64,
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
                repetition_penalty=1.1,
            )

        # 3) Decode ONLY newly generated tokens (completion)
        prompt_len = inputs["input_ids"].shape[1]
        completion_ids = out[0, prompt_len:]
        generated = self.tokenizer.decode(completion_ids, skip_special_tokens=True).strip()

        return [{"generated_text": generated}]