"""CPU checks of pooled counters, uncertainty and actual four-wave driver sequencing."""
from __future__ import annotations

import copy
import json
import time
from pathlib import Path
from types import SimpleNamespace
from typing import cast

import pytest

from offload.gb10_budget import window_cells as wc
from offload.launch.windows import pp4_cells as c
from offload.launch.windows import pp4_metrics as m
from offload.launch.windows import pp4_window as w
from offload.launch.windows.tests.test_pp4_window import complete
from offload.prefill_maxlen.evidence import Json, number, record


def test_pool_uses_round_weighted_totals_and_milliseconds() -> None:
    rows: list[dict[str, Json]] = [
        {'spec_request_rounds': 10., 'accepted_tokens': 0., 'server_decode_seconds': .1},
        {'spec_request_rounds': 90., 'accepted_tokens': 180., 'server_decode_seconds': .9}]
    result = m.pool(rows)
    assert result['acceptance_per_round'] == 1.8
    assert result['server_decode_ms_per_spec_request_round'] == 10.
    assert result['expected_tps'] == 280.


@pytest.mark.parametrize('key,value', [('spec_request_rounds', 0.), ('spec_request_rounds', 1.2),
    ('accepted_tokens', -1.), ('accepted_tokens', 501.), ('server_decode_seconds', float('nan')),
    ('server_decode_seconds', float('inf')), ('server_decode_seconds', 0.)])
def test_invalid_counter_evidence_refuses(key: str, value: float) -> None:
    row: dict[str, Json] = {'spec_request_rounds': 100., 'accepted_tokens': 100., 'server_decode_seconds': 2.}
    row[key] = value
    with pytest.raises(ValueError):
        m.pool([row])


def test_client_rate_cannot_determine_round3_decision() -> None:
    rows = complete()
    for arm in w.ARMS:
        for cell in record(record(rows[arm.name])['cells']).values():
            record(cell)['median_decode_tps'] = 1. if arm.role == 'B' else 10000.
    decision = w.decision(rows)
    assert decision['status'] == 'CANDIDATE FOR REVIEW'
    cell = record(record(decision['cells'])['short-b1'])
    assert record(cast(list[Json], cell['starts'])[0])['expected_tps'] == 100.
    assert record(cast(list[Json], cell['starts'])[1])['expected_tps'] == 110.
    assert record(record(cast(list[Json], cell['starts'])[1])['pooled_95pct'])['acceptance_per_round'] == [1.2, 1.2]


@pytest.mark.parametrize('change', ['three-waves', 'duplicate-prompt', 'source-order', 'old-schema'])
def test_independent_request_and_registration_gates(change: str) -> None:
    rows = complete()
    cell = record(record(record(rows[w.ARMS[1].name])['cells'])['short-b1'])
    waves = cast(list[Json], cell['waves'])
    if change == 'three-waves': waves.pop()
    if change == 'duplicate-prompt': record(waves[2])['streams'] = copy.deepcopy(record(waves[0])['streams'])
    if change == 'source-order': record(waves[0])['family'] = 'filings'
    if change == 'old-schema': cell['metric_registration'] = 'old'
    with pytest.raises(ValueError):
        w.decision(rows)


def test_interval_preserves_request_wave_and_source_variation() -> None:
    rows = complete()
    for arm in (w.ARMS[1],):
        for cell in record(record(rows[arm.name])['cells']).values():
            for wave in cast(list[Json], record(cell)['waves']):
                record(wave)['accepted_tokens'] = 500. if record(wave)['family'] == 'opinions' else 2000.
    result = w.decision(rows)
    assert result['status'] == 'INCONCLUSIVE'
    assert number(record(record(result['cells'])['short-b1'])['point_gain']) > .05
    assert number(cast(list[Json], record(record(result['cells'])['short-b1'])['paired_delta_95pct_tps'])[0]) < 0
    for cell in record(result['cells']).values():
        bounds = cast(list[Json], record(record(cast(list[Json], record(cell)['starts'])[1])['pooled_95pct'])['acceptance_per_round'])
        assert number(bounds[0]) < 1.25 < number(bounds[1])


@pytest.mark.parametrize('name', list(c.CELLS))
def test_four_actual_driver_waves_and_incremental_receipts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, name: str) -> None:
    state = {'rounds': 100., 'accepted': 100., 'seconds': 2., 'completed': 0.}
    calls: list[tuple[str, tuple[str, ...]]] = []
    class Api:
        def metrics(self) -> str:
            return json.dumps({'vllm:spec_decode_num_drafts_total': state['rounds'],
                'vllm:spec_decode_num_accepted_tokens_total': state['accepted'],
                'vllm:request_decode_time_seconds_sum': state['seconds'], 'vllm:request_decode_time_seconds_count': state['completed']})
    def run(api: wc.Api, cell: wc.Cell, texts: dict[str, tuple[str, int]], plan: wc.Plan) -> dict[str, Json]:
        calls.append((cell.name, cell.prompts))
        assert isinstance(api, c.SeededApi) and api.seed == 42+len(calls)-1
        assert cell.max_tokens == 512
        state['completed'] += cell.batch
        state['rounds'] += 100.; state['accepted'] += 120.; state['seconds'] += 2.
        return {'cell': cell.name, 'batch': cell.batch, 'streams': [
            {'usage': {'completion_tokens': 512}, 'finish': 'length', 'first_by_kind_s': {'reasoning': 1.},
             'last_by_kind_s': {'reasoning': 4.}, 'ttft_s': 1., 'prompt_sha256': texts[key][0], 'cache_salt_sha256': plan.salt_prefix+key,
             'sampling': {'temperature': 1, 'top_p': 1, 'seed': 42, 'priority': 10, 'max_tokens': 512}} for key in cell.prompts]}
    monkeypatch.setattr(c, 'load_bench', lambda: SimpleNamespace(parse_metrics=json.loads, spec_delta=lambda *a: {}))
    monkeypatch.setattr(c, 'family_corpus', lambda root, family: (family, {}))
    monkeypatch.setattr(c, 'texts_from', lambda cuts, family, hashes: {key: (family+'-'+key, 32000) for key in c.PROMPTS})
    monkeypatch.setattr(wc, 'run_cell', run)
    monkeypatch.setattr(time, 'sleep', lambda seconds: None)
    output = tmp_path/'partial.json'
    result = c.repeated_cell(cast(wc.Api, Api()), name, {f: {} for f in c.FAMILIES}, tmp_path,
                             wc.Plan(9999999999., 0, [], 'fixed', None), output)
    assert len(calls) == m.group_count(name) and len({key for key, _ in calls}) == m.group_count(name)
    assert [keys for _, keys in calls] == [c.wave_prompts(name, i) for i in range(m.group_count(name))]
    assert result['independent_requests'] == m.group_count(name)*c.CELLS[name].batch
    assert result['acceptance_per_round'] == pytest.approx(1.2)
    assert record(result['pooled_95pct'])['acceptance_per_round'] == [1.2, 1.2]
    assert record(cast(Json, json.loads(output.read_text())))['status'] == 'PARTIAL'
    assert len(m.waves(result, c.CELLS[name].batch)) == m.group_count(name)


@pytest.mark.parametrize('kind', ['ttft', 'server-cost'])
def test_family_regression_cannot_hide_in_pooled_median(kind: str) -> None:
    rows = complete()
    for arm in (w.ARMS[1],):
        for cell in record(record(rows[arm.name])['cells']).values():
            data = record(cell)
            if kind == 'ttft':
                for wave in cast(list[Json], data['waves']):
                    if record(wave)['family'] == 'opinions':
                        record(wave)['median_ttft_s'] = 1.051
            else:
                for wave in cast(list[Json], data['waves']):
                    if record(wave)['family'] == 'opinions':
                        record(wave)['server_decode_seconds'] = 21.01
    assert w.decision(rows)['status'] == 'INCONCLUSIVE'


def test_retired_calibration_gate_cannot_admit_measurement() -> None:
    from offload.launch.windows import pp4_prerequisites as p
    assert 'measured_complete_block_budget' not in p.GATES
    assert 'measured_complete_block_budget' in p.RETIRED_GATES