"""Shared text generation utilities."""
from __future__ import annotations

import re
import logging
from typing import List, Dict, Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from src.config import DEFAULT_MAX_NEW_TOKENS, DEFAULT_TEMPERATURE

logger = logging.getLogger(__name__)


def strip_think_blocks(text: str) -> str:
    """Remove <think>...</think> blocks if the model emits them."""
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()


def extract_assistant_reply(text: str) -> str:
    """Extract the assistant reply from a chat-templated output string."""
    # Common delimiters used by chat templates
    markers = [
        "assistant\n",
        "assistant",
        "<|im_start|>assistant",
        "<|start_header_id|>assistant<|end_header_id|>",
    ]
    for marker in markers:
        if marker in text:
            reply = text.split(marker, 1)[-1]
            return strip_think_blocks(reply).strip()
    return strip_think_blocks(text).strip()


def chat_generate(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    messages: List[Dict[str, Any]],
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = DEFAULT_TEMPERATURE,
    do_sample: bool = False,
) -> str:
    """Generate an assistant reply from a list of chat messages."""
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        padding=True,
        truncation=True,
    )
    if hasattr(model, "device") and model.device.type != "meta":
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature if do_sample else None,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated = outputs[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(generated, skip_special_tokens=True)
    return extract_assistant_reply(text)