"""WIN2 greedy refinement (Fable review): distribution check at the first greedy divergence.

A near-tie only says the two competing tokens were close. This adds the question "did the whole next-token
distribution change?": in a control/candidate/control window the candidate arm's greedy check records the top-20 and
the prefix at each divergence (nbw_cells.greedy(record_distribution=True)); the following control arm re-probes the
same prefixes (fresh cache salt) and top20_kl measures KL(control || candidate) over the shared top-20 tokens plus one
OTHER bucket, the same coarse estimator UP2's screen uses (up2_checks.numerical). The bound is a registration choice
of each runner; there is no default.
"""
from __future__ import annotations

import math
from collections.abc import Mapping, Sequence
from typing import cast

from offload.gb10_budget import window_cells as wc
from offload.launch.windows import nbw_cells as nbw
from offload.launch.windows import up2_checks as ck
from offload.prefill_maxlen.evidence import Json, number, record

MIN_COMMON_MASS = .99


def _probabilities(top: Mapping[str, Json]) -> dict[int, float]:
    ck.require(len(top) == 20, 'exactly twenty top-logprob tokens required')
    out: dict[int, float] = {}
    for key, value in top.items():
        ck.require(key.startswith('token_id:') and key[9:].isdigit(), 'malformed top-20 token id')
        lp = number(value)
        ck.require(-9999 < lp <= 0, 'clamped or invalid logprob')
        out[int(key[9:])] = math.exp(lp)
    ck.require(sum(out.values()) <= 1 + 1e-6, 'top-20 probabilities exceed one')
    return out


def top20_kl(reference: Mapping[str, Json], experiment: Mapping[str, Json]) -> dict[str, float]:
    """KL(reference || experiment) over shared top-20 tokens + OTHER; common_mass = smaller shared probability mass."""
    p, q = _probabilities(reference), _probabilities(experiment)
    keys = sorted(p.keys() & q.keys())
    px, qx = [p[k] for k in keys], [q[k] for k in keys]
    common = min(sum(px), sum(qx))
    px.append(max(0.0, 1 - sum(px)))
    qx.append(max(0.0, 1 - sum(qx)))
    kl = sum(v * math.log(v / w) if w > 0 else math.inf for v, w in zip(px, qx) if v > 0)
    ck.require(math.isfinite(kl), 'nonfinite coarse KL at the divergence')
    return {'kl_nats': kl, 'common_mass': common, 'shared_tokens': float(len(keys))}


def passes(result: Mapping[str, float], bound: float) -> bool:
    if not (bound > 0 and math.isfinite(bound)):
        raise ValueError('a positive registered KL bound is required')
    return result['common_mass'] >= MIN_COMMON_MASS and result['kl_nats'] <= bound


def probe(api: wc.Http, prefix: Sequence[int], salt: str) -> dict[str, Json]:
    """Top-20 at an exact prefix from the server under test, with a fresh cache salt."""
    ck.require(bool(salt.strip()), 'probe cache salt required')
    r = api.post('/v1/completions', {'model': nbw.MODEL, 'prompt': cast(Json, list(prefix)), 'max_tokens': 1,
                                     'temperature': 0, 'logprobs': 20, 'return_tokens_as_token_ids': True,
                                     'priority': 10, 'cache_salt': salt})
    tops = record(record(ck.rows(r['choices'])[0])['logprobs'])['top_logprobs']
    ck.require(isinstance(tops, list) and len(tops) == 1, 'one top-logprob position expected')
    return record(cast(list[Json], tops)[0])


def compare_recorded(api: wc.Http, candidate_rows: Sequence[Mapping[str, Json]], salt_prefix: str,
                     bound: float) -> list[dict[str, Json]]:
    """Re-probe every recorded candidate divergence on the CURRENT (control) server and measure the KL."""
    out: list[dict[str, Json]] = []
    for row in candidate_rows:
        if row.get('identical') is True:
            continue
        ck.require('top20_logprobs' in row and 'probe_prefix_token_ids' in row,
                   'candidate divergence lacks a recorded distribution (record_distribution=True)')
        prefix = nbw.token_ids(row['probe_prefix_token_ids'])
        control = probe(api, prefix, f'{salt_prefix}:kl:{row.get("i")}')
        result = top20_kl(control, record(row['top20_logprobs']))
        out.append({'i': row.get('i'), **{k: v for k, v in result.items()}, 'bound': bound,
                    'pass': passes(result, bound)})
    return out