"""
HuggingFace Transformers generation engine — drop-in replacement for
GenerationEngine that loads any model from the HuggingFace Hub.

Uses ``AutoModelForCausalLM`` and ``AutoTokenizer`` so any causal LM
works without architecture-specific code.  Exposes the same
``generate_answer()`` / ``generate_stream()`` interface so eval scripts
can swap engines with a single line.

Designed for the multi-model eval benchmark: load a model, run the
eval, delete the engine, free memory, load the next model.
"""

from __future__ import annotations

import re

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from src.utils.device import get_default_device


# Pattern to strip <think>...</think> blocks (possibly spanning lines)
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
# Unclosed <think> block at the end of a generation
_THINK_UNCLOSED_RE = re.compile(r"<think>.*", re.DOTALL)


class TransformersEngine:
    """Architecture-agnostic generation engine backed by HuggingFace Transformers.

    Parameters
    ----------
    model_id : str
        HuggingFace model identifier (e.g. ``"Qwen/Qwen3-4B"``).
    dtype : torch.dtype
        Weight precision.  Defaults to ``torch.bfloat16``.
    max_memory_gb : float | None
        Optional memory cap (not enforced, just for documentation).
    """

    def __init__(
        self,
        model_id: str,
        dtype: torch.dtype = torch.bfloat16,
        enable_thinking: bool = False,
        trust_remote_code: bool = True,
    ):
        self.model_id = model_id
        self.device = get_default_device()
        self.enable_thinking = enable_thinking

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_id, trust_remote_code=trust_remote_code,
        )

        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        ).to(self.device)
        self.model.eval()

    def generate_answer(
        self,
        messages: list[dict],
        max_tokens: int = 512,
        temperature: float = 0.7,
        top_k: int = 20,
        top_p: float = 0.95,
    ) -> str:
        """Generate a complete response (non-streaming).

        Compatible with ``GenerationEngine.generate_answer()``.
        """
        # Apply chat template
        template_kwargs = {
            "tokenize": False,
            "add_generation_prompt": True,
        }
        # Disable thinking mode when supported (Qwen3/3.5) to avoid
        # wasting the token budget on <think> blocks.
        if not self.enable_thinking:
            try:
                text = self.tokenizer.apply_chat_template(
                    messages, enable_thinking=False, **template_kwargs,
                )
            except TypeError:
                # Model doesn't support enable_thinking param
                text = self.tokenizer.apply_chat_template(
                    messages, **template_kwargs,
                )
        else:
            text = self.tokenizer.apply_chat_template(
                messages, **template_kwargs,
            )
        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
        input_len = inputs["input_ids"].shape[1]

        # Generate — fall back to greedy if sampling hits NaN logits
        # (common on MPS + bfloat16 with long contexts).
        do_sample = temperature > 0.01
        gen_kwargs: dict = dict(
            max_new_tokens=max_tokens,
            pad_token_id=self.tokenizer.eos_token_id,
        )
        if do_sample:
            gen_kwargs.update(
                do_sample=True,
                temperature=max(temperature, 1e-4),
                top_k=top_k,
                top_p=top_p,
            )
        else:
            gen_kwargs["do_sample"] = False
        with torch.no_grad():
            try:
                outputs = self.model.generate(**inputs, **gen_kwargs)
            except RuntimeError as e:
                if "probability tensor" in str(e) and do_sample:
                    # NaN/Inf logits — retry with greedy decoding
                    outputs = self.model.generate(
                        **inputs,
                        max_new_tokens=max_tokens,
                        pad_token_id=self.tokenizer.eos_token_id,
                        do_sample=False,
                    )
                else:
                    raise

        # Decode only the generated tokens
        generated = outputs[0][input_len:]
        answer = self.tokenizer.decode(generated, skip_special_tokens=True)

        # Strip <think>...</think> blocks
        answer = _THINK_RE.sub("", answer)
        answer = _THINK_UNCLOSED_RE.sub("", answer)

        return answer.strip()

    def generate_stream(
        self,
        messages: list[dict],
        max_tokens: int = 512,
        temperature: float = 0.7,
        top_k: int = 20,
        top_p: float = 0.95,
    ):
        """Yield the full answer as a single chunk.

        Provided for interface compatibility with ``GenerationEngine``.
        True token-level streaming is not needed for eval.
        """
        yield self.generate_answer(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
        )