"""
vLLM generation wrapper shared by pool building, eval, and smoke tests.

Captures per-story cumulative logprob, which E4's `divpo-prob` variant needs:
DivPO's probability criterion picks the LOWEST-logprob passing story as the
diverse chosen and the HIGHEST-logprob failing story as the common rejected.
The highest-probability sample in a temperature pool is, by construction, the
near-greedy one -- which is why this is the principled version of "reject the
greedy decode".

Length normalization matters here. Raw cumulative logprob scales with token
count, so ranking by it would just rank by length (short stories win). We keep
both `cumlogprob` and `mean_logprob = cumlogprob / n_tokens` and use the mean
for DivPO ranking.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, asdict
from pathlib import Path

# NEVER truncate a story. The prior run died of a max_new_tokens wall, so this
# is sized for real headroom, not for the target length: a 500-word story is
# ~650-700 Qwen tokens, so 768 would clip anything that runs long and would
# reproduce the exact bug we are fixing. 1280 lets the model finish and overrun
# a little; overrun is then DETECTED by the >600-word gate rather than silently
# chopped. Measure the failure, don't manufacture it.
#
# Nothing downstream truncates either: the judge receives the complete story
# text, and embeddings are computed over the complete story text.
MAX_NEW_TOKENS = 1280
MAX_MODEL_LEN = 2048


@dataclass
class Gen:
    prompt_id: str
    prompt: str
    idx: int
    text: str
    n_tokens: int
    cumlogprob: float
    mean_logprob: float
    finish_reason: str

    def as_dict(self) -> dict:
        return asdict(self)


def build_llm(
    model: str,
    gpu_mem_util: float = 0.85,
    max_model_len: int = MAX_MODEL_LEN,
    seed: int = 0,
    enable_lora: bool = False,
):
    from vllm import LLM
    kw = dict(
        model=model,
        dtype="bfloat16",
        gpu_memory_utilization=gpu_mem_util,
        max_model_len=max_model_len,
        seed=seed,
        enforce_eager=False,
        disable_log_stats=True,
    )
    if enable_lora:
        kw.update(enable_lora=True, max_lora_rank=32)
    return LLM(**kw)


def sampling_params(
    n: int, temperature: float, top_p: float, seed: int | None,
    max_tokens: int = MAX_NEW_TOKENS,
):
    from vllm import SamplingParams
    return SamplingParams(
        n=n,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
        seed=seed,
        logprobs=0,            # cumulative_logprob only; per-token table not needed
        skip_special_tokens=True,
    )


def render_chat(tokenizer, prompt: str) -> str:
    """Render the chat template, with thinking OFF where the model has it.

    Qwen3-8B is a hybrid-thinking model: left at its default it emits a <think>
    block before the story, which would (a) eat the token budget, (b) pollute
    the embedding with reasoning text, and (c) make the 8B arm incomparable to
    the 4B-Instruct arms, which have no thinking mode at all. Qwen3-4B-Instruct
    ignores the kwarg, so the same call is correct for both.
    """
    from data import chat_messages
    msgs = chat_messages(prompt)
    try:
        return tokenizer.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False,
        )
    except TypeError:
        return tokenizer.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True,
        )


def generate(
    llm,
    tokenizer,
    prompts: list[dict],
    n: int,
    temperature: float,
    top_p: float,
    seed: int | None = None,
    max_tokens: int = MAX_NEW_TOKENS,
    lora_path: str | None = None,
) -> list[Gen]:
    """prompts = [{'id':..., 'prompt':...}]  ->  flat list of n*len(prompts) Gens."""
    texts = [render_chat(tokenizer, p["prompt"]) for p in prompts]
    sp = sampling_params(n, temperature, top_p, seed, max_tokens)

    kw = {}
    if lora_path:
        from vllm.lora.request import LoRARequest
        kw["lora_request"] = LoRARequest("adapter", 1, lora_path)

    outs = llm.generate(texts, sp, **kw)

    gens: list[Gen] = []
    for p, out in zip(prompts, outs):
        for j, o in enumerate(out.outputs):
            ntok = len(o.token_ids)
            cum = float(o.cumulative_logprob) if o.cumulative_logprob is not None else 0.0
            gens.append(Gen(
                prompt_id=p["id"], prompt=p["prompt"], idx=j,
                text=o.text.strip(), n_tokens=ntok,
                cumlogprob=cum,
                mean_logprob=(cum / ntok) if ntok else 0.0,
                finish_reason=o.finish_reason or "",
            ))
    return gens


def save_gens(gens: list[Gen], path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        for g in gens:
            f.write(json.dumps(g.as_dict()) + "\n")


def load_gens(path: str | Path) -> list[Gen]:
    return [Gen(**json.loads(l)) for l in open(path) if l.strip()]