# one-shot inference helper. loads the model (with optional lora adapter) and
# produces a cleaned string. greedy decode, max_new_tokens capped near the
# input length so the model cannot balloon into a chat reply.

import torch

from cleanup.prompts import build_messages


def load_model(base_model: str, adapter_dir=None, dtype=None):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    src = adapter_dir if adapter_dir else base_model
    tokenizer = AutoTokenizer.from_pretrained(src, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if dtype is None:
        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        torch_dtype=dtype,
        device_map="auto" if torch.cuda.is_available() else None,
    )
    if adapter_dir:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, adapter_dir)
    model.eval()
    return model, tokenizer


def clean_text(model, tokenizer, raw: str, max_new_factor: float = 1.6) -> str:
    messages = build_messages(raw)
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    raw_tokens = len(tokenizer.encode(raw))
    max_new = min(256, max(8, int(raw_tokens * max_new_factor)))
    with torch.no_grad():
        out_ids = model.generate(
            **inputs,
            do_sample=False,
            max_new_tokens=max_new,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    new_tokens = out_ids[0][inputs.input_ids.shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()