"""B04 local control slices; run only explicitly budgeted nodes."""
import asyncio
import threading
import sqlite3

import pytest


@pytest.mark.asyncio
async def test_terminal_release(tmp_path, monkeypatch):
    runtime, registry, run = await live_run(tmp_path, monkeypatch)
    closing, release = asyncio.Event(), asyncio.Event()

    class Client:
        async def aclose(self):
            closing.set()
            await release.wait()

    client = Client()
    run._owned_clients[id(client)] = client
    cancel = asyncio.create_task(runtime.cancel(
        run.run_id, reason="done", owner_epoch=1, boot_id="a"))
    try:
        await asyncio.wait_for(closing.wait(), 2)
        assert registry.inspect(run.run_id)["resources_confirmed_exited"] is False
    finally:
        release.set()
        assert await cancel == "cancelled"
    reopened = SQLiteHostRegistry(registry.path, host_id="host", boot_id="b")
    evidence = reopened.inspect(run.run_id)
    assert evidence["status"] == "cancelled"
    assert evidence["resources_confirmed_exited"] is True
    with sqlite3.connect(registry.path) as db:
        assert db.execute("SELECT count(*) FROM active_claims").fetchone()[0] == 0
    original = reopened.lookup_grant(run._admission.grant_id,
                                    evidence["digest"])
    assert original["execution_id"] == run.run_id
    with pytest.raises(EvoRuntimeError, match="EXECUTION_IDENTITY_CONFLICT"):
        reopened.bind(execution_id=run.run_id, grant_id="new", digest="new",
                      thread_id="other", turn_id="other")

from EvoScientist.llm.contracts import EvoRuntimeError, WebHostContext
from EvoScientist.llm.host_execution_registry import SQLiteHostRegistry
from tests.test_web_model_runtime import _runtime, _input, _preparation, _admission, _Sink


@pytest.mark.asyncio
async def test_precommit(tmp_path, monkeypatch):
    runtime, authority = _runtime(tmp_path, monkeypatch)
    registry = SQLiteHostRegistry(tmp_path / "host.sqlite", host_id="host", boot_id="a")
    runtime.host_registry = registry
    value = _input()
    quote = await runtime.prepare_model_run(
        _preparation(authority, value, title_policy="disabled"), value,
        WebHostContext("/tmp", "/tmp", object(), object(), runtime_event_sink=_Sink()))
    assert quote.execution_id
    assert quote.unsigned_payload()["execution_id"] == quote.execution_id
    entered = []
    original = runtime.agent_factory

    def construct(*args):
        assert registry.inspect(quote.execution_id)["execution_id"] == quote.execution_id
        assert registry.inspect(quote.execution_id)["source"] == "host_binding_only"
        entered.append(quote.execution_id)
        return original(*args)

    runtime.agent_factory = construct
    run = await runtime.start_web_run(_admission(authority, quote))
    assert run.run_id == quote.execution_id
    assert entered == [quote.execution_id]
    await run.cancel("test", owner_epoch=1, boot_id="a")


def test_pending_continuation(tmp_path):
    registry = SQLiteHostRegistry(tmp_path / "host.sqlite", host_id="host", boot_id="a")
    registry.bind(execution_id="parent", grant_id="g1", digest="d1",
                  thread_id="thread", turn_id="turn")
    registry.finish("parent", outcome="awaiting_input", checkpoint_id="cp1")
    child = dict(execution_id="child", grant_id="g2", digest="d2",
                 thread_id="thread", turn_id="turn")
    with pytest.raises(EvoRuntimeError, match="CONTINUATION_REQUIRED"):
        registry.bind(**child)
    with pytest.raises(EvoRuntimeError, match="CONTINUATION_INVALID"):
        registry.bind(**child, predecessor_execution_id="parent", predecessor_checkpoint_id="wrong")
    registry.bind(**child, predecessor_execution_id="parent", predecessor_checkpoint_id="cp1")
    assert registry.inspect("parent")["status"] == "awaiting_input"
    registry.finish("child", outcome="completed")
    with pytest.raises(EvoRuntimeError, match="CONTINUATION_INVALID"):
        registry.bind(**{**child, "execution_id": "third", "grant_id": "g3"},
                      predecessor_execution_id="parent", predecessor_checkpoint_id="cp1")


def test_restart_cancel_intent(tmp_path):
    registry = SQLiteHostRegistry(tmp_path / "host.sqlite", host_id="host", boot_id="a")
    registry.bind(execution_id="run", grant_id="grant", digest="digest",
                  thread_id="thread", turn_id="turn")
    registry.accept_cancel("run", owner_epoch=1, boot_id="a", reason="stop")
    restarted = SQLiteHostRegistry(registry.path, host_id="host", boot_id="b")
    evidence = restarted.inspect("run")
    assert evidence["cancel_requested"] is True
    assert evidence["recovery_action"] == "inspect_only"
    assert evidence["resources_confirmed_exited"] is False
    with pytest.raises(EvoRuntimeError, match="EXECUTION_BOOT_MISMATCH"):
        restarted.finish("run", outcome="cancelled")


def test_registry_busy(tmp_path):
    registry = SQLiteHostRegistry(tmp_path / "host.sqlite", host_id="host", boot_id="a")
    with sqlite3.connect(registry.path) as lock:
        lock.execute("BEGIN IMMEDIATE")
        with pytest.raises(EvoRuntimeError, match="HOST_REGISTRY_BUSY"):
            registry.bind(execution_id="run", grant_id="grant", digest="digest",
                          thread_id="thread", turn_id="turn")


@pytest.mark.asyncio
async def test_finish_fault(tmp_path, monkeypatch):
    from types import SimpleNamespace
    for outcome in ("completed", "awaiting_input", "cancelled"):
        case = tmp_path / outcome
        case.mkdir()
        runtime, registry, run = await live_run(case, monkeypatch)
        with sqlite3.connect(registry.path) as db:
            db.execute("CREATE TRIGGER fail_finish BEFORE INSERT ON terminal_evidence "
                       "BEGIN SELECT RAISE(ABORT, 'injected finish failure'); END")
        run._agent_task.cancel()
        await asyncio.gather(run._agent_task, return_exceptions=True)
        async def empty(*args, **kwargs):
            if False:
                yield {}
        async def state(*args):
            return SimpleNamespace(config={"configurable": {
                "thread_id": "thread", "checkpoint_id": "cp"}},
                tasks=(), next=(), interrupts=(SimpleNamespace(id="p", value={}),)
                if outcome == "awaiting_input" else ())
        monkeypatch.setattr("EvoScientist.stream.events.stream_agent_events", empty)
        run._agent.aget_state = state
        if outcome == "cancelled":
            run._request_cancel("fault")
        else:
            run._agent_task = asyncio.create_task(run._run_agent())
        try:
            assert await run.wait_stopped(timeout=0.15) == "unknown"
            terminals = [e for e in run._host.runtime_event_sink.events
                         if e.payload.get("kind") == "run_terminal"]
            assert len(terminals) == 1
            assert terminals[0].payload["outcome"] == outcome
            assert registry.inspect(run.run_id)["status"] == "unknown"
        finally:
            with sqlite3.connect(registry.path) as db:
                db.execute("DROP TRIGGER fail_finish")
        assert await run.wait_stopped(timeout=3) == outcome
        terminals_after = [e for e in run._host.runtime_event_sink.events
                           if e.payload.get("kind") == "run_terminal"]
        assert terminals_after == terminals
        assert registry.inspect(run.run_id)["status"] == outcome
        with sqlite3.connect(registry.path) as db:
            assert db.execute("SELECT count(*) FROM active_claims").fetchone()[0] == 0


@pytest.mark.asyncio
async def test_construction_fault(tmp_path, monkeypatch):
    from EvoScientist.llm.runtime import _construction_owner
    registry_check = SQLiteHostRegistry(tmp_path / "consumed.sqlite", host_id="host", boot_id="a")
    registry_check.bind(execution_id="parent", grant_id="parent", digest="d", thread_id="t", turn_id="t")
    registry_check.finish("parent", outcome="awaiting_input", checkpoint_id="cp")
    registry_check.bind(execution_id="child", grant_id="child", digest="d", thread_id="t", turn_id="t",
                        predecessor_execution_id="parent", predecessor_checkpoint_id="cp")
    registry_check.finish("child", outcome="failed")
    with pytest.raises(EvoRuntimeError, match="CONTINUATION_CONSUMED_FAILURE_REQUIRES_REAUTHORIZATION"):
        registry_check.bind(execution_id="retry", grant_id="retry", digest="d", thread_id="t", turn_id="t",
                            predecessor_execution_id="parent", predecessor_checkpoint_id="cp")
    runtime, authority = _runtime(tmp_path, monkeypatch)
    registry = SQLiteHostRegistry(tmp_path / "host.sqlite", host_id="host", boot_id="a")
    runtime.host_registry = registry
    release = asyncio.Event()
    owners = []
    class Client:
        async def aclose(self):
            await release.wait()
    def broken(*args):
        run = _construction_owner.get()
        owners.append(run)
        client = Client()
        run._owned_clients[id(client)] = client
        raise ValueError("construction fault")
    runtime.agent_factory = broken
    value = _input()
    sink = _Sink()
    quote = await runtime.prepare_model_run(
        _preparation(authority, value, title_policy="disabled"), value,
        WebHostContext("/tmp", "/tmp", object(), object(), runtime_event_sink=sink))
    admission = _admission(authority, quote)
    with pytest.raises((TimeoutError, ValueError)):
        await runtime.start_web_run(admission)
    run = owners[0]
    assert registry.inspect(run.run_id)["status"] == "unknown"
    with sqlite3.connect(registry.path) as db:
        assert db.execute("SELECT count(*) FROM active_claims").fetchone()[0] == 1
        db.execute("CREATE TRIGGER fail_finish BEFORE INSERT ON terminal_evidence "
                   "BEGIN SELECT RAISE(ABORT, 'injected finish failure'); END")
    release.set()
    assert await run.wait_stopped(timeout=0.15) == "unknown"
    with sqlite3.connect(registry.path) as db:
        assert db.execute("SELECT count(*) FROM active_claims").fetchone()[0] == 1
        db.execute("DROP TRIGGER fail_finish")
    assert await run.wait_stopped(timeout=3) == "failed"
    assert registry.inspect(run.run_id)["status"] == "failed"
    with sqlite3.connect(registry.path) as db:
        assert db.execute("SELECT count(*) FROM active_claims").fetchone()[0] == 0
        assert db.execute("SELECT count(*) FROM executions").fetchone()[0] == 1
    terminals = [e for e in sink.events if e.payload.get("kind") == "run_terminal"]
    assert len(terminals) == 1
    assert terminals[0].payload["error_code"] == "RUN_CONSTRUCTION_FAILED"
    with pytest.raises(EvoRuntimeError, match="EXECUTION_FAILED_RETRY_REQUIRES_NEW_ADMISSION"):
        await runtime.start_web_run(admission)


async def live_run(tmp_path, monkeypatch):
    runtime, authority = _runtime(tmp_path, monkeypatch)
    registry = SQLiteHostRegistry(tmp_path / "host.sqlite", host_id="host", boot_id="a")
    runtime.host_registry = registry
    entered = asyncio.Event()

    async def events(*args, **kwargs):
        entered.set()
        await asyncio.Event().wait()
        yield {}

    monkeypatch.setattr("EvoScientist.stream.events.stream_agent_events", events)
    value = _input()
    quote = await runtime.prepare_model_run(
        _preparation(authority, value, title_policy="disabled"), value,
        WebHostContext("/tmp", "/tmp", object(), object(), runtime_event_sink=_Sink()),
    )
    run = await runtime.start_web_run(_admission(authority, quote))
    await asyncio.wait_for(entered.wait(), 2)
    return runtime, registry, run


@pytest.mark.asyncio
async def test_stale_handle(tmp_path, monkeypatch):
    runtime, registry, run = await live_run(tmp_path, monkeypatch)
    peer = SQLiteHostRegistry(registry.path, host_id="host", boot_id="a")
    peer.transfer_control(run.run_id, expected_epoch=1, new_epoch=2)
    try:
        with pytest.raises(EvoRuntimeError, match="OWNER_EPOCH_REQUIRED"):
            await run.cancel("old handle")
        with pytest.raises(EvoRuntimeError, match="OWNER_EPOCH_STALE"):
            await run.cancel("old owner", owner_epoch=1, boot_id="a")
        with pytest.raises(EvoRuntimeError, match="EXECUTION_BOOT_MISMATCH"):
            await runtime.cancel(run.run_id, reason="wrong boot", owner_epoch=2, boot_id="b")
        with pytest.raises(EvoRuntimeError, match="OWNER_EPOCH_STALE"):
            runtime.inspect(run.run_id, owner_epoch=1, boot_id="a")
        assert not run._agent_task.done()
        assert await runtime.cancel(run.run_id, reason="current", owner_epoch=2, boot_id="a") == "cancelled"
    finally:
        if not run._agent_task.done():
            run._agent_task.cancel()
        await asyncio.gather(run._agent_task, return_exceptions=True)


@pytest.mark.asyncio
async def test_transfer_cancel_race(tmp_path, monkeypatch):
    runtime, registry, run = await live_run(tmp_path, monkeypatch)
    peer = SQLiteHostRegistry(registry.path, host_id="host", boot_id="a")
    barrier = threading.Barrier(2)
    original = registry.accept_cancel

    def racing_accept(*args, **kwargs):
        barrier.wait(timeout=2)
        return original(*args, **kwargs)

    monkeypatch.setattr(registry, "accept_cancel", racing_accept)
    def transfer():
        barrier.wait(timeout=2)
        return peer.transfer_control(run.run_id, expected_epoch=1, new_epoch=2)

    worker = asyncio.create_task(asyncio.to_thread(transfer))
    await asyncio.sleep(0)
    try:
        try:
            result = await runtime.cancel(run.run_id, reason="racing", owner_epoch=1, boot_id="a")
        except EvoRuntimeError as exc:
            assert "OWNER_EPOCH_STALE" in str(exc)
            result = "stale"
        assert await asyncio.wait_for(worker, 3) == 2
        monkeypatch.setattr(registry, "accept_cancel", original)
        if result == "stale":
            assert not run._agent_task.done()
            assert await runtime.cancel(run.run_id, reason="current", owner_epoch=2, boot_id="a") == "cancelled"
        else:
            assert result == "cancelled"
        assert run._agent_task.done()
    finally:
        run._agent_task.cancel()
        await asyncio.gather(run._agent_task, worker, return_exceptions=True)


@pytest.mark.asyncio
async def test_accepted_cancel_during_transfer(tmp_path, monkeypatch):
    runtime, registry, run = await live_run(tmp_path, monkeypatch)
    peer = SQLiteHostRegistry(registry.path, host_id="host", boot_id="a")
    closing = asyncio.Event()
    release = asyncio.Event()

    class Client:
        async def aclose(self):
            closing.set()
            await release.wait()

    client = Client()
    run._owned_clients[id(client)] = client
    cancel = asyncio.create_task(runtime.cancel(
        run.run_id, reason="accepted", owner_epoch=1, boot_id="a"))
    try:
        await asyncio.wait_for(closing.wait(), 2)
        assert await asyncio.wait_for(asyncio.to_thread(
            peer.transfer_control, run.run_id, expected_epoch=1, new_epoch=2), 1) == 2
        with sqlite3.connect(registry.path) as db:
            rows = db.execute("SELECT owner_epoch, reason FROM cancel_intents").fetchall()
        assert rows == [(1, "accepted")]
        with pytest.raises(EvoRuntimeError, match="OWNER_EPOCH_STALE"):
            await run.cancel("late", owner_epoch=1, boot_id="a")
        assert not cancel.done()
    finally:
        release.set()
        assert await cancel == "cancelled"