import os
from typing import Any, Dict, List
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline

# Ensure your template includes escaped newlines correctly
CHAT_TEMPLATE = (
    "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"
    "{% for message in messages %}"
    "{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}"
    "{% endfor %}"
    "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}"
)

class EndpointHandler:
    def __init__(self, model_dir: str, **kwargs: Any):
        # Load tokenizer and model from provided model_dir
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
        # Assign the chat_template under 'default'
        self.tokenizer.chat_template = {"default": CHAT_TEMPLATE}
        self.model = AutoModelForCausalLM.from_pretrained(
            model_dir, trust_remote_code=True, device_map="auto"
        )
        self.pipeline = TextGenerationPipeline(
            model=self.model, tokenizer=self.tokenizer, return_full_text=False
        )

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # Validate input structure
        messages: List[Dict[str, str]] = data.get("messages")
        if messages is None:
            raise ValueError("Request body must include 'messages' array.")

        # Format prompt with proper controlling flag
        inputs = self.tokenizer.apply_chat_template(
            messages=messages, add_generation_prompt=True, return_tensors=None
        )
        # Generate text
        gen = self.pipeline(inputs, max_new_tokens=data.get("parameters", {}).get("max_new_tokens", 128))
        content = gen[0]["generated_text"]

        return {"choices": [{"message": {"role": "assistant", "content": content}}]}