"""Task 3 acceptance: conversation identity wired into chat_completions.

These tests exercise the integration the wiring added end to end: a client
``X-Router-Conversation`` header must scope the recorded ``session_key`` (and
therefore the session incumbent, the cloud ``energy_observations`` ledger and
the local-dispatch energy ledger), while ``X-Router-Agent``/``X-Router-Parent``
must land in their columns. A headerless request must still behave exactly as
before — one fingerprint key, NULL agent/parent, byte-identical to the frozen
pre-wiring snapshot (``tests/fixtures/route_decision_no_header.json``).

The commit that added the snapshot fixture deliberately precedes this file, so
a clean checkout can read it without a gitignored ``.omo/evidence/`` path.
"""

import sqlite3
from types import SimpleNamespace

import pytest

import dispatcher

from test_no_header_snapshot import (
    assert_matches_snapshot,
    make_router,
    no_header_payload,
)

CHEAP = "cheap-model"
DEAR = "dear-model"


@pytest.fixture(autouse=True, scope="module")
def _clear_prefix_probe_state():
    """Clean up ``prefix_probe._store`` after our module so the snapshot
    test (runs alphabetically after us) sees a clean prefix-probe cache.
    The probe stores per-session fingerprints in process-memory and a
    POST with the snapshot's payload here would otherwise make a later
    test see a prior turn and report non-NULL prefix divergence."""
    yield
    import prefix_probe
    prefix_probe._store.clear()

_SYS = "You work in /home/user/proj and follow the design docs."


def _messages(*roles):
    return [
        {"role": r, "content": c}
        for r, c in zip(roles, ["system prompt", "user query one", "assistant reply", "user query two"])
    ][: len(roles)]


def _rows(db_path, table, where=None, order="id"):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    q = f"SELECT * FROM {table}"
    if where:
        q += f" WHERE {where}"
    q += f" ORDER BY {order}"
    rows = [dict(r) for r in conn.execute(q).fetchall()]
    conn.close()
    return rows


def _latest_decision(db_path):
    rows = _rows(db_path, "route_decisions")
    assert rows, "no route_decisions row"
    return rows[-1]


def _post(client, *, headers=None, model="auto", msgs=None):
    payload = {
        "model": model,
        "messages": msgs if msgs is not None else _messages("system", "user"),
    }
    return client.post(
        "/v1/chat/completions", json=payload, headers=headers or {}
    )


# --- Acceptance 1: same prompt, different conversation -> distinct keys -----


def test_same_system_prompt_distinct_conversations_get_distinct_session_keys(
    tmp_path, monkeypatch
):
    client, db_path = make_router(tmp_path, monkeypatch)
    msgs = _messages("system", "user")

    r1 = _post(client, headers={"X-Router-Conversation": "a"}, msgs=msgs)
    r2 = _post(client, headers={"X-Router-Conversation": "b"}, msgs=msgs)
    assert r1.status_code == 200 and r2.status_code == 200

    keys = {r["session_key"] for r in _rows(db_path, "route_decisions")}
    assert keys == {"c:a", "c:b"}, keys


# --- Acceptance 2: each conversation's incumbent is its OWN last model -----


def test_incumbent_is_conversation_scoped_not_global(
    tmp_path, monkeypatch
):
    client, db_path = make_router(tmp_path, monkeypatch)
    conn = sqlite3.connect(db_path)
    for sk, model in (
        ("c:a", CHEAP),   # interleave so the GLOBAL last chat row is cheap
        ("c:b", DEAR),
        ("c:a", CHEAP),
    ):
        conn.execute(
            """INSERT INTO route_decisions (
                observed_at, kind, session_key, selected_provider,
                selected_model, task_category, task_tier
            ) VALUES ('2026-09-13T00:00:00+00:00', 'chat', ?, 'neuralwatt',
                      ?, 'coding_general', 2)""",
            (sk, model),
        )
    conn.commit()
    conn.close()

    monkeypatch.setattr(
        dispatcher.cfg.objective, "incumbent_cache_pricing", True
    )

    captured = {}

    class Capturer:
        def __call__(self, rows, **kwargs):
            captured["incumbent"] = kwargs.get("incumbent")
            from routing import rank_candidates as _real
            return _real(rows, **kwargs)

    monkeypatch.setattr(dispatcher, "rank_candidates", Capturer())

    r = _post(client, headers={"X-Router-Conversation": "b"})
    assert r.status_code == 200, r.text

    # Conversation b's own last chat row is dear-model, even though the
    # overall-last chat row belongs to conversation a (cheap-model).
    assert captured["incumbent"] == ("neuralwatt", DEAR), captured["incumbent"]


# --- Acceptance 3: headerless requests land under one fingerprint, NULL
#     agent/parent, byte-identical to the committed snapshot -----------------


def test_headerless_requests_reproduce_committed_snapshot(
    tmp_path, monkeypatch
):
    client, db_path = make_router(tmp_path, monkeypatch)
    # The exact payload the committed snapshot was frozen from: the row must
    # be byte-identical, so the messages (and therefore the fingerprint,
    # pinch and prefix-probe values) must match the fixture's.
    payload = no_header_payload()

    # First POST: fresh session, no prior prefix-probe state → the row has
    # NULL prefix columns and is byte-identical to the pre-wiring fixture.
    r1 = _post(client, msgs=payload["messages"])
    assert r1.status_code == 200, r1.text
    assert_matches_snapshot(db_path)

    # Second POST in the same session: the prefix probe now sees a prior
    # turn and writes non-NULL prefix columns, so the byte-identical
    # assertion cannot include this row. But both rows must share one
    # fingerprint key and carry NULL agent/parent.
    r2 = _post(client, msgs=payload["messages"])
    assert r2.status_code == 200, r2.text
    rows = _rows(db_path, "route_decisions")
    assert len(rows) == 2
    assert rows[0]["session_key"] == rows[1]["session_key"]
    assert rows[0]["session_key"] is not None
    for r in rows:
        assert r["agent"] is None
        assert r["parent_key"] is None


# --- Acceptance 4: agent and parent land in their columns -------------------


def test_agent_and_parent_land_in_columns(tmp_path, monkeypatch):
    client, db_path = make_router(tmp_path, monkeypatch)
    r = _post(
        client,
        headers={
            "X-Router-Conversation": "conv1",
            "X-Router-Agent": "my-agent",
            "X-Router-Parent": "parent1",
        },
    )
    assert r.status_code == 200, r.text
    row = _latest_decision(db_path)
    assert row["session_key"] == "c:conv1"
    assert row["agent"] == "my-agent"
    assert row["parent_key"] == "c:parent1"


# --- Acceptance 5: invalid conversation header falls back to fingerprint ----


def test_invalid_conversation_header_falls_back_to_fingerprint(
    tmp_path, monkeypatch
):
    client, db_path = make_router(tmp_path, monkeypatch)
    msgs = _messages("system", "user")

    bad = _post(
        client,
        headers={"X-Router-Conversation": "has space and /slash"},
        msgs=msgs,
    )
    assert bad.status_code == 200, bad.text

    good = _post(client, msgs=msgs)
    assert good.status_code == 200, good.text

    rows = _rows(db_path, "route_decisions")
    bad_key, fp_key = rows[0]["session_key"], rows[1]["session_key"]
    # The invalid header is treated as absent: same fingerprint as no header.
    assert bad_key == fp_key
    assert not bad_key.startswith("c:")


# --- Acceptance 6: energy ledgers get the conversation key ------------------

def test_cloud_energy_session_key_starts_with_c(tmp_path, monkeypatch):
    client, db_path = make_router(tmp_path, monkeypatch)
    r = _post(
        client,
        headers={"X-Router-Conversation": "ener"},
        model="auto",
        msgs=_messages("system", "user"),
    )
    assert r.status_code == 200, r.text
    rows = _rows(db_path, "energy_observations")
    assert rows, "expected an energy_observations row"
    for row in rows:
        assert (row["session_key"] or "").startswith("c:"), row


def _local_dispatch_harness(tmp_path, monkeypatch, session_key):
    """Drive one real _run_local_dispatch and return the produced row."""
    client, db_path = make_router(tmp_path, monkeypatch)
    model_id = "local-q"

    monkeypatch.setattr(dispatcher.cfg.local_energy, "enabled", True)
    # local_energy_dispatch_models is a read-only cached property on the
    # Pydantic config; shadow it on the class so the metering gate sees our
    # model as meterable.
    monkeypatch.setattr(
        type(dispatcher.cfg), "local_energy_dispatch_models",
        property(lambda self: frozenset({model_id})),
    )
    monkeypatch.setattr(dispatcher.cfg.circuit_breaker, "enabled", False)

    class FakeMeasurement:
        avg_power_watts = 100.0
        duration_seconds = 2.0

        def __enter__(self):
            return self

        def __exit__(self, *args):
            return False

    monkeypatch.setattr(
        dispatcher.local_energy, "measure",
        lambda *a, **k: FakeMeasurement(),
    )

    def fake_post(url, headers=None, json=None, stream=False, timeout=None):
        resp = SimpleNamespace()
        resp.status_code = 200
        resp.json = lambda: {
            "id": "ollama-local-1",
            "choices": [
                {"message": {"role": "assistant", "content": "local answer"}}
            ],
            "usage": {"prompt_tokens": 20, "completion_tokens": 8},
        }
        resp.text = "{}"
        return resp

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)

    entry = SimpleNamespace(
        model_id=model_id,
        base_url="http://localhost:11434/v1",
        api_key_env=None,
        max_output_tokens=2048,
        timeout_seconds=60.0,
    )
    result = dispatcher._run_local_dispatch(
        entry,
        _messages("system", "user"),
        category=None,
        body={},
        session_key=session_key,
    )
    assert result["payload"]["choices"][0]["message"]["content"] == "local answer"
    assert result["request_id"].startswith("local-dispatch-")
    rows = _rows(db_path, "local_energy_observations")
    assert rows, "expected a local_energy_observations row"
    return rows, client


def test_local_dispatch_answer_records_conversation_key(
    tmp_path, monkeypatch
):
    rows, _ = _local_dispatch_harness(tmp_path, monkeypatch, session_key="c:ener")
    assert rows[-1]["session_key"] == "c:ener"


# --- Acceptance 7 [H5]: /dispatch (no header) keeps local energy NULL -------


def test_no_header_local_dispatch_keeps_energy_null_and_resolves(
    tmp_path, monkeypatch
):
    # /dispatch passes no session_key to _run_local_dispatch, so the keyword
    # defaults to None and the local energy row must stay NULL.
    rows, _ = _local_dispatch_harness(tmp_path, monkeypatch, session_key=None)
    assert rows[-1]["session_key"] is None


# --- Acceptance 8 [H4]: cold start — sensible routing on a short history ----


def test_cold_start_conversation_routes_sensibly_on_short_history(
    tmp_path, monkeypatch
):
    client, db_path = make_router(tmp_path, monkeypatch)
    # A short sub-agent history (3 messages) with a fresh conversation key
    # that has NO prior route_decisions rows: must still resolve to a real
    # model instead of failing or guessing an incumbent that cannot exist.
    msgs = _messages("system", "user", "assistant")
    r = _post(
        client,
        headers={"X-Router-Conversation": "cold"},
        msgs=msgs,
    )
    assert r.status_code == 200, r.text
    row = _latest_decision(db_path)
    assert row["kind"] == "chat"
    assert row["session_key"] == "c:cold"
    assert row["selected_model"] is not None