from typing import Any
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

SYSTEM_PROMPT = """You are a financial planning assistant specializing in US tax optimization for early retirees pursuing FIRE.

Key 2024 tax facts:
- 0% LTCG bracket: up to $94,050 MFJ
- 12% ordinary income ceiling: $94,300 MFJ
- Roth conversion ladder: 5-year seasoning rule applies
- RMDs begin at age 73
- ACA cliff: ~$81,760 for couples"""

class EndpointHandler:
    def __init__(self, path=""):
        # Load tokenizer from base Mistral to avoid TokenizersBackend error
        self.tokenizer = AutoTokenizer.from_pretrained(
            "mistralai/Mistral-7B-v0.1"
        )
        self.tokenizer.pad_token = self.tokenizer.eos_token

        # No quantization_config here — transformers reads it from config.json
        # bitsandbytes handles it automatically
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            device_map="auto",
            torch_dtype=torch.bfloat16,
        )
        self.model.eval()
        print("✅ Model loaded!")

    def __call__(self, data: Any) -> Any:
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", {})
        prompt = f"<s>[INST] {SYSTEM_PROMPT}\n\n{inputs} [/INST]"
        encoded = self.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=1024
        ).to(self.model.device)
        with torch.no_grad():
            outputs = self.model.generate(
                **encoded,
                max_new_tokens=parameters.get("max_new_tokens", 512),
                temperature=parameters.get("temperature", 0.3),
                do_sample=True,
                repetition_penalty=1.2,
                pad_token_id=self.tokenizer.eos_token_id,
            )
        generated = outputs[0][encoded["input_ids"].shape[1]:]
        return [{"generated_text": self.tokenizer.decode(generated, skip_special_tokens=True)}]