"""
Custom handler for HuggingFace Inference Endpoints.

This handler loads the base model with the LoRA adapter for inference.
"""

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# Configuration
BASE_MODEL = "Qwen/Qwen3-Coder-Next"

SYSTEM_PROMPT = """You are a Univer API expert. Generate JavaScript code to perform spreadsheet operations using the Univer API.
The code should be an async arrow function that uses the univerAPI global object.

IMPORTANT PRINCIPLES:
1. ALWAYS use spreadsheet formulas (setFormula) for calculations like sum, average, max, min, count, etc.
   - Use setFormula('=SUM(...)'), setFormula('=AVERAGE(...)'), etc.
   - Do NOT use JavaScript calculations (reduce, Math.max, etc.) for these operations
   - Formulas are preferred because they update dynamically when data changes
2. Use the correct API for the task:
   - To delete rows: use deleteRows(), NOT clear()
   - To insert rows: use insertRows(), NOT manual approaches
3. ONLY use documented Univer API methods. Do NOT invent or guess method names."""


class EndpointHandler:
    def __init__(self, path: str = ""):
        """
        Initialize the handler by loading the model and tokenizer.

        Args:
            path: Path to the model directory (contains adapter files)
        """
        # Load tokenizer from base model to avoid tokenizer_config compatibility issues.
        # The adapter repo's tokenizer_config.json may have extra_special_tokens as a list,
        # which causes AttributeError with some transformers versions.
        print(f"Loading tokenizer from base model: {BASE_MODEL}...")
        self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        print(f"Loading base model: {BASE_MODEL}...")
        # Use 8-bit quantization for memory efficiency
        bnb_config = BitsAndBytesConfig(
            load_in_8bit=True,
            llm_int8_enable_fp32_cpu_offload=False,
        )

        self.model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
        )

        print(f"Loading LoRA adapter from {path}...")
        self.model = PeftModel.from_pretrained(self.model, path)
        self.model.eval()

        print("Model loaded successfully!")

    def __call__(self, data: dict) -> dict:
        """
        Handle inference requests.

        Args:
            data: Request data with 'inputs' key containing the instruction

        Returns:
            Generated code response
        """
        # Extract input
        inputs = data.get("inputs", data.get("input", ""))
        parameters = data.get("parameters", {})

        # Build messages
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": inputs},
        ]

        # Apply chat template
        prompt = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

        # Tokenize
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
        input_ids = input_ids.to(self.model.device)

        # Generation parameters
        max_new_tokens = parameters.get("max_new_tokens", 1024)
        temperature = parameters.get("temperature", 0.1)
        top_p = parameters.get("top_p", 0.9)
        do_sample = parameters.get("do_sample", temperature > 0)

        # Generate
        with torch.no_grad():
            outputs = self.model.generate(
                input_ids,
                max_new_tokens=max_new_tokens,
                temperature=temperature if do_sample else None,
                top_p=top_p if do_sample else None,
                do_sample=do_sample,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            )

        # Decode response (only new tokens)
        response = self.tokenizer.decode(
            outputs[0][input_ids.shape[1] :],
            skip_special_tokens=True,
        )

        return {"generated_text": response.strip()}