"""
HuggingFace Inference Endpoint custom handler for sinllama-mcq-merged-2.0.

The deployed model (itsjorigo/sinllama-mcq-merged) is a fully merged
AutoModelForCausalLM — NOT a PEFT adapter stack. Load it directly.

The tokenizer must come from polyglots/SinLlama_v01 (trust_remote_code=True)
because that repo defines the custom TokenizersBackend class.

Request:  {"inputs": {"passage": "<Sinhala text>", "entity_block": "<optional>"}}
Response: {"mcq": "ප්‍රශ්නය: ..."}
"""

import re
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

SINLLAMA_ID = "polyglots/SinLlama_v01"

PROMPT_TEMPLATE = (
    "පහත ඉතිහාස ඡේදය කියවා, ඒ ගැන බහු-විකල්ප ප්‍රශ්නයක් සාදන්න.\n\n"
    "ඡේදය: {passage}\n\n"
    "{entity_block}"
    "MCQ:"
)


class EndpointHandler:
    def __init__(self, path=""):
        # Tokenizer from SinLlama — defines the TokenizersBackend class
        print("Loading tokenizer from SinLlama repo...", flush=True)
        self.tokenizer = AutoTokenizer.from_pretrained(SINLLAMA_ID, trust_remote_code=True)

        # Model loaded directly — itsjorigo/sinllama-mcq-merged is already fully merged
        print(f"Loading merged model from {path!r}...", flush=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            torch_dtype=torch.float16,
            device_map="auto",
            low_cpu_mem_usage=True,
            attn_implementation="sdpa",
        )
        self.model.eval()
        print("EndpointHandler ready.", flush=True)

    def __call__(self, data: dict) -> dict:
        inputs       = data.get("inputs", {})
        passage      = inputs.get("passage", "")
        entity_block = inputs.get("entity_block", "")

        if not passage:
            return {"error": 'No passage provided. Send {"inputs": {"passage": "..."}}'}

        prompt = PROMPT_TEMPLATE.format(
            passage=passage.strip(),
            entity_block=entity_block,
        )

        enc = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

        with torch.no_grad():
            out = self.model.generate(
                **enc,
                max_new_tokens=280,
                temperature=0.7,
                do_sample=True,
                repetition_penalty=1.1,
                eos_token_id=self.tokenizer.eos_token_id,
                pad_token_id=self.tokenizer.pad_token_id,
            )

        new_ids = out[0][enc.input_ids.shape[1]:]
        text    = self.tokenizer.decode(new_ids, skip_special_tokens=True).strip()

        # Ensure each option starts on its own line
        for tag in ["A)", "B)", "C)", "D)", "නිවැරදි පිළිතුර:"]:
            text = re.sub(rf"(?<!\n)({re.escape(tag)})", r"\n\1", text)

        return {"mcq": text.strip()}