"""Continuous batching and bounded admission (DGR-012).

These tests drive the node-local continuous-batching scheduler with the *same*
pure-numpy KV-cached dense-Llama reference the Hot KV State manager uses
(DGR-007), imported from ``test_hot_kv_state``. That keeps the whole gate
deterministic, download-free, GPU-free, and API-credit-free while exercising the
real KV isolation path (``KvBoundaryAdapter`` + ``HotKvStateManager``) rather than
a mock.

Coverage maps to the story's acceptance criteria:

* bounded admission against weight/KV/scratch/queue budgets,
* compatible decode steps batched with per-session positions/outputs preserved,
* prefill never starving in-flight decode (explicit decode-first policy),
* backpressure when the bounded queue is full,
* capability telemetry reporting every required signal,
* a deterministic 1/2/4/8 concurrency sweep showing saturation and no
  cross-session corruption.
"""

from __future__ import annotations

import numpy as np
import pytest

from meshnet_node.hot_kv_state import (
    HotKvStateConfig,
    HotKvStateManager,
    KvBoundaryAdapter,
    kv_recipe_for,
)
from meshnet_node.batch_scheduler import (
    AdmissionReason,
    ContinuousBatchScheduler,
    GenerationRequest,
    KvBatchEngine,
    NodeBudget,
    Phase,
    run_concurrency_sweep,
)

# Reuse the certified numpy dense-Llama reference and shard from the DGR-007 gate.
from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard


# --------------------------------------------------------------------------- #
# Helpers.
# --------------------------------------------------------------------------- #


class _FakeClock:
    def __init__(self) -> None:
        self.now = 0.0

    def __call__(self) -> float:
        return self.now

    def advance(self, delta: float) -> None:
        self.now += delta


def _make_engine(
    model: _KvDenseLlama | None = None,
    *,
    config: HotKvStateConfig | None = None,
) -> KvBatchEngine:
    """A full-shard KV batch engine over the deterministic numpy dense-Llama."""
    model = model or _KvDenseLlama()
    shard = _KvReferenceShard(model, 0, model.n_layers - 1)
    manager = HotKvStateManager(kv_recipe_for(shard), config=config)
    adapter = KvBoundaryAdapter(shard, manager)
    return KvBatchEngine(adapter)


def _reference_tokens(model: _KvDenseLlama, prompt, n_new: int) -> list[int]:
    return model.stateless_greedy(list(prompt), n_new)


def _generation(session_id: str, prompt, n_new: int, epoch: int = 0) -> GenerationRequest:
    return GenerationRequest(
        session_id=session_id,
        route_epoch=epoch,
        prompt_token_ids=tuple(prompt),
        max_new_tokens=n_new,
    )


# --------------------------------------------------------------------------- #
# Bounded admission (weight / KV / scratch / queue budgets).
# --------------------------------------------------------------------------- #


def test_admission_respects_active_scratch_and_queue_budgets():
    "Admission fills active slots, queues the overflow, then rejects a full queue.\n\nTags: node, scheduler, admission"
    engine = _make_engine()
    budget = NodeBudget(
        max_active_sessions=2,
        scratch_bytes_per_session=1,
        scratch_budget_bytes=2,  # scratch also caps at 2 concurrent
        max_queue_depth=1,
        max_batch_size=2,
    )
    scheduler = ContinuousBatchScheduler(engine, budget)

    a = scheduler.submit(_generation("a", [1, 2, 3], 4))
    b = scheduler.submit(_generation("b", [4, 5, 6], 4))
    assert a.reason is AdmissionReason.ADMITTED
    assert b.reason is AdmissionReason.ADMITTED

    # Two active slots full -> the next goes to the bounded queue.
    c = scheduler.submit(_generation("c", [7, 8, 9], 4))
    assert c.reason is AdmissionReason.QUEUED

    # Queue depth 1 is now full -> backpressure rejection.
    d = scheduler.submit(_generation("d", [1, 1, 1], 4))
    assert d.reason is AdmissionReason.REJECTED_QUEUE_FULL
    assert d.rejected

    telem = scheduler.telemetry()
    assert telem.active_sessions == 2
    assert telem.queue_depth == 1
    assert telem.rejected_admissions_total == 1
    assert telem.rejected_by_reason[AdmissionReason.REJECTED_QUEUE_FULL.value] == 1


def test_admission_rejects_a_session_that_cannot_fit_the_kv_budget():
    "A generation whose whole KV cannot fit the node budget is rejected up front.\n\nTags: node, scheduler, admission"
    engine = _make_engine()
    per_token = engine._manager.recipe.bytes_per_token()
    # Budget holds only 3 positions; a prompt(4)+7 new = 10 final positions cannot fit.
    budget = NodeBudget(kv_budget_bytes=per_token * 3)
    scheduler = ContinuousBatchScheduler(engine, budget)
    decision = scheduler.submit(_generation("big", [1, 2, 3, 4], 7))
    assert decision.reason is AdmissionReason.REJECTED_KV_BUDGET
    assert scheduler.telemetry().rejected_admissions_total == 1


def test_admission_rejects_when_per_session_scratch_exceeds_budget():
    "A per-session scratch larger than the whole scratch envelope is rejected.\n\nTags: node, scheduler, admission"
    engine = _make_engine()
    budget = NodeBudget(scratch_bytes_per_session=1024, scratch_budget_bytes=512)
    scheduler = ContinuousBatchScheduler(engine, budget)
    decision = scheduler.submit(_generation("s", [1, 2], 2))
    assert decision.reason is AdmissionReason.REJECTED_SCRATCH_BUDGET


def test_duplicate_submission_is_rejected():
    "Submitting a session id that is already scheduled is rejected as a duplicate.\n\nTags: node, scheduler, admission"
    engine = _make_engine()
    scheduler = ContinuousBatchScheduler(engine, NodeBudget(max_active_sessions=4))
    assert scheduler.submit(_generation("dup", [1, 2], 3)).reason is AdmissionReason.ADMITTED
    assert scheduler.submit(_generation("dup", [3, 4], 3)).reason is AdmissionReason.REJECTED_DUPLICATE


def test_weight_budget_is_reported_in_telemetry():
    "The resident weight footprint is surfaced as a capability signal.\n\nTags: node, scheduler, telemetry"
    engine = _make_engine()
    budget = NodeBudget(weight_bytes=123_456)
    scheduler = ContinuousBatchScheduler(engine, budget)
    assert scheduler.telemetry().weight_bytes == 123_456


# --------------------------------------------------------------------------- #
# Continuous batching preserves per-session positions and outputs.
# --------------------------------------------------------------------------- #


def test_batched_decode_preserves_per_session_positions_and_outputs():
    "Four sessions batched together each reproduce their own stateless tokens.\n\nTags: node, scheduler, batching"
    model = _KvDenseLlama()
    engine = _make_engine(model)
    budget = NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4)
    scheduler = ContinuousBatchScheduler(engine, budget)

    prompts = {
        "alpha": [1, 2, 3, 4],
        "bravo": [40, 39, 2, 15],
        "charlie": [7, 7, 7, 7],
        "delta": [31, 5, 18, 22],
    }
    n_new = 10
    references = {sid: _reference_tokens(model, p, n_new) for sid, p in prompts.items()}
    # The four references must diverge, else "no cross-talk" would be vacuous.
    assert len({tuple(v) for v in references.values()}) == 4

    for sid, prompt in prompts.items():
        assert scheduler.submit(_generation(sid, prompt, n_new)).running

    outputs = scheduler.run_to_completion()
    for sid in prompts:
        assert outputs[sid] == references[sid], sid

    telem = scheduler.telemetry()
    # A genuine batch formed: at least one decode tick carried all four sessions.
    assert telem.batch_occupancy_max == 4
    assert telem.completed_sessions == 4
    assert telem.active_sessions == 0


def test_positions_are_isolated_across_different_prompt_lengths():
    "Sessions with different prompt lengths keep independent positions when batched.\n\nTags: node, scheduler, batching"
    model = _KvDenseLlama()
    engine = _make_engine(model)
    scheduler = ContinuousBatchScheduler(
        engine, NodeBudget(max_active_sessions=3, max_batch_size=3, max_queue_depth=3)
    )
    jobs = {
        "short": ([5], 6),
        "medium": ([2, 9, 14], 6),
        "long": ([1, 2, 3, 4, 5, 6, 7], 6),
    }
    refs = {sid: _reference_tokens(model, p, n) for sid, (p, n) in jobs.items()}
    for sid, (prompt, n) in jobs.items():
        scheduler.submit(_generation(sid, prompt, n))
    outputs = scheduler.run_to_completion()
    for sid in jobs:
        assert outputs[sid] == refs[sid], sid


# --------------------------------------------------------------------------- #
# Prefill does not starve decode.
# --------------------------------------------------------------------------- #


def test_prefill_does_not_starve_in_flight_decode():
    "A burst of new prefills never stalls an already-decoding session.\n\nTags: node, scheduler, fairness"
    model = _KvDenseLlama()
    engine = _make_engine(model)
    # One prefill per tick (budget == a single prompt) so prefill is throttled and
    # we can observe that decode still advances every tick.
    budget = NodeBudget(
        max_active_sessions=8,
        max_batch_size=8,
        max_queue_depth=8,
        scratch_bytes_per_session=1,
        scratch_budget_bytes=8,
        max_prefill_tokens_per_tick=4,
    )
    scheduler = ContinuousBatchScheduler(engine, budget)

    # Session A starts and prefills on tick 1.
    scheduler.submit(_generation("A", [3, 14, 1, 5], 12))
    scheduler.run_tick()
    a_state = scheduler.session_result("A")
    assert a_state.phase is Phase.DECODING
    a_len = len(a_state.generated)
    assert a_len == 1

    # Burst of new work arrives while A is decoding.
    for sid in ("B", "C", "D", "E"):
        scheduler.submit(_generation(sid, [2, 27, 18, 4], 12))

    # Over the next few ticks A must decode on *every* tick (never starved),
    # while at most one new session prefills per tick (prefill is bounded).
    prefill_counts = []
    for _ in range(4):
        report = scheduler.run_tick()
        new_a_len = len(scheduler.session_result("A").generated)
        assert new_a_len == a_len + 1, "decode of A stalled while prefills were pending"
        a_len = new_a_len
        assert "A" in report.decoded
        prefill_counts.append(len(report.prefilled))

    assert max(prefill_counts) <= 1, "prefill was not bounded per tick"


def test_decode_first_policy_is_explicit_in_a_single_tick():
    "In one tick decode of active sessions precedes prefill of new ones.\n\nTags: node, scheduler, fairness"
    model = _KvDenseLlama()
    engine = _make_engine(model)
    scheduler = ContinuousBatchScheduler(
        engine,
        NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4,
                   scratch_bytes_per_session=1, scratch_budget_bytes=4),
    )
    scheduler.submit(_generation("live", [1, 2, 3], 8))
    scheduler.run_tick()  # 'live' prefills, now decoding
    scheduler.submit(_generation("fresh", [9, 8, 7], 8))
    report = scheduler.run_tick()
    assert "live" in report.decoded
    assert "fresh" in report.prefilled


# --------------------------------------------------------------------------- #
# Backpressure and bounded memory.
# --------------------------------------------------------------------------- #


def test_backpressure_signals_when_queue_full_then_recovers():
    "A full queue rejects new work; a completed session frees a slot for the queue.\n\nTags: node, scheduler, backpressure"
    engine = _make_engine()
    budget = NodeBudget(
        max_active_sessions=1,
        max_batch_size=1,
        max_queue_depth=1,
        scratch_bytes_per_session=1,
        scratch_budget_bytes=1,
    )
    scheduler = ContinuousBatchScheduler(engine, budget)
    assert scheduler.submit(_generation("first", [1, 2], 2)).running
    assert scheduler.submit(_generation("second", [3, 4], 2)).reason is AdmissionReason.QUEUED
    # Both a slot and the queue are full now.
    assert scheduler.submit(_generation("third", [5, 6], 2)).reason is AdmissionReason.REJECTED_QUEUE_FULL

    # Drain 'first'; the queued 'second' must be pulled into the freed slot.
    scheduler.run_to_completion()
    outputs = scheduler.outputs()
    assert set(outputs) == {"first", "second"}


def test_completed_sessions_release_kv_so_growth_is_bounded():
    "Finished sessions release their KV, so total KV returns to zero.\n\nTags: node, scheduler, backpressure"
    engine = _make_engine()
    scheduler = ContinuousBatchScheduler(
        engine, NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=8)
    )
    for sid in ("a", "b", "c", "d"):
        scheduler.submit(_generation(sid, [1, 2, 3], 4))
    scheduler.run_to_completion()
    telem = scheduler.telemetry()
    assert telem.kv_total_bytes == 0, "KV not released after completion"
    assert telem.active_sessions == 0
    assert telem.completed_sessions == 4


# --------------------------------------------------------------------------- #
# Telemetry.
# --------------------------------------------------------------------------- #


def test_telemetry_reports_every_required_signal():
    "The capability snapshot reports sessions, queue, batch, KV, rates, rejections.\n\nTags: node, scheduler, telemetry"
    model = _KvDenseLlama()
    engine = _make_engine(model)
    clock = _FakeClock()
    budget = NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=1)
    scheduler = ContinuousBatchScheduler(engine, budget, clock=clock)

    scheduler.submit(_generation("x", [1, 2, 3], 4))
    scheduler.submit(_generation("y", [4, 5, 6], 4))
    scheduler.submit(_generation("z", [7, 8, 9], 4))  # queued
    rejected = scheduler.submit(_generation("w", [1, 1, 1], 4))  # queue full
    assert rejected.rejected

    clock.advance(1.0)
    scheduler.run_tick()  # both prefill
    clock.advance(1.0)
    scheduler.run_tick()  # both decode as a batch of 2

    clock.advance(2.0)
    telem = scheduler.telemetry()
    snap = telem.to_dict()
    for key in (
        "active_sessions", "queue_depth", "batch_occupancy_last",
        "batch_occupancy_avg", "batch_occupancy_max", "weight_bytes",
        "kv_total_bytes", "kv_budget_bytes", "kv_pressure",
        "scratch_used_bytes", "scratch_budget_bytes", "scratch_pressure",
        "prefill_tokens_total", "decode_tokens_total",
        "prefill_tokens_per_sec", "decode_tokens_per_sec",
        "rejected_admissions_total", "rejected_by_reason",
        "completed_sessions", "ticks",
    ):
        assert key in snap, key

    assert telem.batch_occupancy_max == 2
    assert telem.prefill_tokens_total == 6  # two prompts of length 3
    assert telem.decode_tokens_total == 2  # one batched decode step, two sessions
    assert telem.rejected_admissions_total == 1
    # Rates are deterministic under the injected clock: 4 seconds elapsed.
    assert telem.decode_tokens_per_sec == pytest.approx(2 / 4.0)
    assert telem.prefill_tokens_per_sec == pytest.approx(6 / 4.0)
    assert 0.0 < telem.kv_pressure <= 1.0


# --------------------------------------------------------------------------- #
# Concurrency 1/2/4/8 sweep: saturation and no corruption.
# --------------------------------------------------------------------------- #


def test_concurrency_sweep_identifies_saturation_without_corruption():
    "A 1/2/4/8 sweep raises batch occupancy, cuts ticks, and never corrupts output.\n\nTags: node, scheduler, benchmark"
    model = _KvDenseLlama()
    prompts = {
        "s0": [1, 2, 3, 4], "s1": [5, 6, 7, 8], "s2": [9, 10, 11, 12],
        "s3": [13, 14, 15, 16], "s4": [17, 18, 19, 20], "s5": [21, 22, 23, 24],
        "s6": [25, 26, 27, 28], "s7": [29, 30, 31, 32],
    }
    n_new = 8
    requests = [_generation(sid, p, n_new) for sid, p in prompts.items()]

    sweep = run_concurrency_sweep(
        lambda: _make_engine(model),
        requests,
        concurrency_levels=(1, 2, 4, 8),
    )

    assert sweep.corruption_free
    assert [r.concurrency for r in sweep.results] == [1, 2, 4, 8]

    # No session hit a cache miss (budgets are sized to never evict here).
    assert all(r.cache_misses == 0 for r in sweep.results)
    assert all(r.rejected_admissions == 0 for r in sweep.results)

    # Each per-session stream matches the serialized (concurrency-1) reference.
    for sid, prompt in prompts.items():
        assert list(sweep.reference_outputs[sid]) == _reference_tokens(model, prompt, n_new)

    occupancies = [r.avg_batch_occupancy for r in sweep.results]
    ticks = [r.ticks for r in sweep.results]
    tokens_per_tick = [r.tokens_per_tick for r in sweep.results]

    # Batching packs more sessions per decode step as concurrency rises, so
    # average occupancy strictly increases and total ticks strictly decrease.
    assert occupancies == sorted(occupancies) and len(set(occupancies)) == 4
    assert ticks == sorted(ticks, reverse=True) and len(set(ticks)) == 4
    # Aggregate work per tick rises with concurrency (the throughput win).
    assert tokens_per_tick == sorted(tokens_per_tick)

    # For eight equal-length jobs the node keeps saturating up to the top level.
    assert sweep.saturation_concurrency == 8

    # The report is JSON-safe for durable evidence.
    import json

    json.dumps(sweep.to_dict())


def test_concurrency_sweep_saturates_below_max_when_load_is_small():
    "With fewer concurrent jobs than slots, saturation is found below the top level.\n\nTags: node, scheduler, benchmark"
    model = _KvDenseLlama()
    # Only three jobs: at concurrency 4 and 8 the batch can never exceed 3, so
    # occupancy stops rising past the load and saturation is detected early.
    requests = [
        _generation("j0", [1, 2, 3], 6),
        _generation("j1", [4, 5, 6], 6),
        _generation("j2", [7, 8, 9], 6),
    ]
    sweep = run_concurrency_sweep(
        lambda: _make_engine(model), requests, concurrency_levels=(1, 2, 4, 8)
    )
    assert sweep.corruption_free
    assert sweep.saturation_concurrency <= 4
    # Levels at or above the load size share the same occupancy/tick profile.
    top = [r for r in sweep.results if r.concurrency >= 4]
    assert len({r.ticks for r in top}) == 1


# --------------------------------------------------------------------------- #
# Engine contract guards.
# --------------------------------------------------------------------------- #


def test_kv_batch_engine_requires_a_full_shard():
    "The batch engine rejects a partial (non head+tail) shard.\n\nTags: node, scheduler"
    model = _KvDenseLlama()
    head = _KvReferenceShard(model, 0, 2)  # head only, not tail
    manager = HotKvStateManager(kv_recipe_for(head))
    adapter = KvBoundaryAdapter(head, manager)
    with pytest.raises(Exception):
        KvBatchEngine(adapter)


def test_run_to_completion_is_bounded_against_misconfiguration():
    "run_to_completion raises rather than looping forever when work cannot drain.\n\nTags: node, scheduler"
    engine = _make_engine()
    scheduler = ContinuousBatchScheduler(
        engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4)
    )
    scheduler.submit(_generation("only", [1, 2], 3))
    # A tiny explicit tick ceiling is exceeded deterministically.
    with pytest.raises(Exception):
        scheduler.run_to_completion(max_ticks=1)