"""Stop and Steer must share a deterministic stream-ownership edge."""
import queue
import threading
from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock

import pytest

from api import config, streaming
from tests.test_real_steer import _captured_response, _make_handler


class ObservedLock:
    """Report a contender before acquisition, without timing-based guesses."""

    def __init__(self):
        self.lock = threading.Lock()
        self.owner = None
        self.contender = threading.Event()
        # Optional unlocked hooks: before acquisition and after release.
        self.before_acquire = None
        self.after_release = None

    def __enter__(self):
        if self.before_acquire is not None:
            self.before_acquire()
        if self.owner is not None and self.owner != threading.get_ident():
            self.contender.set()
        assert self.lock.acquire(timeout=5), "stream lock acquisition timed out"
        self.owner = threading.get_ident()
        return self

    def __exit__(self, *args):
        self.owner = None
        self.lock.release()
        if self.after_release is not None:
            self.after_release()


@pytest.fixture
def scene(monkeypatch):
    lock = ObservedLock()
    agent = Mock(session_id="original")
    agent.steer.return_value = True
    maps = {
        "STREAMS": {"run": queue.Queue()},
        "CANCEL_FLAGS": {"run": threading.Event()},
        "AGENT_INSTANCES": {"run": agent},
        "STREAM_PARTIAL_TEXT": {"run": "partial answer"},
        "STREAM_REASONING_TEXT": {"run": "partial reasoning"},
        "STREAM_LIVE_TOOL_CALLS": {"run": [{"name": "example"}]},
        "STREAM_SESSION_OWNERS": {"run": "original"},
        "ACTIVE_RUNS": {"run": {"session_id": "original", "backend": "legacy", "phase": "running"}},
        "SESSION_AGENT_CACHE": {"original": (agent, "sig")},
    }
    for name, value in maps.items():
        monkeypatch.setattr(config, name, value)
        if hasattr(streaming, name):
            monkeypatch.setattr(streaming, name, value)
    monkeypatch.setattr(config, "STREAMS_LOCK", lock)
    monkeypatch.setattr(streaming, "STREAMS_LOCK", lock)
    monkeypatch.setattr(streaming, "get_session", lambda sid: SimpleNamespace(active_stream_id="run", messages=[]))
    monkeypatch.setattr(streaming, "_get_session_agent_lock", lambda sid: nullcontext())
    # Persistence is outside this ordering test; do not touch real sessions.
    monkeypatch.setattr(streaming, "_stream_writeback_is_current", lambda *args: False)
    from api import clarify
    monkeypatch.setattr(clarify, "clear_pending", lambda sid: None)
    return lock, agent


def steer():
    handler = _make_handler()
    streaming._handle_chat_steer(handler, {"session_id": "original", "text": "guidance"})
    return _captured_response(handler)


@pytest.mark.parametrize("registered", [True, False])
def test_cancel_snapshot_claims_stream_before_later_steer(scene, monkeypatch, registered):
    """Pause at the exact old snapshot -> cancellation publication gap."""
    lock, agent = scene
    if not registered:
        config.AGENT_INSTANCES.clear()  # Matching-cache compatibility path.
    reached = threading.Event()
    release = threading.Event()
    update = streaming.update_active_run
    held_at_publication = []

    def publish(*args, **kwargs):
        held_at_publication.append(lock.owner == threading.get_ident())
        reached.set()
        assert release.wait(5), "cancel publication barrier timed out"
        return update(*args, **kwargs)

    monkeypatch.setattr(streaming, "update_active_run", publish)
    with ThreadPoolExecutor(max_workers=2) as pool:
        cancel = pool.submit(streaming.cancel_stream, "run")
        try:
            assert reached.wait(5)
            guidance = pool.submit(steer)
            # On the old code there is no held lock: force Steer to complete
            # within the publication gap, proving the original accepted=True.
            if held_at_publication == [False]:
                assert guidance.result(timeout=5)["accepted"] is False
            else:
                assert lock.contender.wait(5)
                assert not guidance.done()
        finally:
            release.set()
        assert cancel.result(timeout=5) is True
        assert guidance.result(timeout=5) == {"accepted": False, "fallback": "stream_dead", "stream_id": None}
    assert held_at_publication == [True]
    agent.steer.assert_not_called()
    agent.interrupt.assert_called_once()
    assert config.ACTIVE_RUNS["run"]["phase"] == "cancelling"
    assert "run" not in config.STREAMS
    assert "run" not in config.AGENT_INSTANCES
    # Buffers are still owned by worker teardown, not eager cancellation.
    assert config.STREAM_PARTIAL_TEXT["run"] == "partial answer"
    assert config.STREAM_REASONING_TEXT["run"] == "partial reasoning"
    assert config.STREAM_LIVE_TOOL_CALLS["run"] == [{"name": "example"}]


@pytest.mark.parametrize("registered", [True, False])
def test_steer_claimed_first_enqueues_before_cancel(scene, registered):
    lock, agent = scene
    if not registered:
        config.AGENT_INSTANCES.clear()  # Matching-cache compatibility path.
    reached = threading.Event()
    release = threading.Event()
    order = []

    def enqueue(text):
        assert lock.owner == threading.get_ident()
        reached.set()
        assert release.wait(5), "steer enqueue barrier timed out"
        order.append("steer")
        return True

    def interrupt(reason):
        assert lock.owner != threading.get_ident(), "interrupt must be outside registry lock"
        assert config.ACTIVE_RUNS["run"]["phase"] == "cancelling"
        assert "run" not in config.STREAMS
        assert "run" not in config.AGENT_INSTANCES
        order.append("cancel")

    agent.steer.side_effect = enqueue
    agent.interrupt.side_effect = interrupt
    with ThreadPoolExecutor(max_workers=2) as pool:
        guidance = pool.submit(steer)
        try:
            assert reached.wait(5)
            cancel = pool.submit(streaming.cancel_stream, "run")
            assert lock.contender.wait(5)
            assert not cancel.done()
        finally:
            release.set()
        assert guidance.result(timeout=5)["accepted"] is True
        assert cancel.result(timeout=5) is True
    assert order == ["steer", "cancel"]


@pytest.mark.parametrize("gap", ["before_acquire", "after_release"])
def test_cache_only_steer_is_atomic_with_stop(scene, monkeypatch, gap):
    """Cache-only Steer (no AGENT_INSTANCES entry) must never enqueue after Stop claims.

    Steer selects the matching cached worker, then pauses at an unlocked gap
    around its next stream-lock edge while Stop runs to completion. Exactly one
    complete outcome is allowed: Steer already queued under the stream lock
    before Stop claimed, or Stop claimed first and Steer reports ``stream_dead``
    without ever calling ``agent.steer()`` on the cached worker.
    """
    lock, agent = scene
    config.AGENT_INSTANCES.clear()
    decoy = Mock(session_id="other-session")
    config.SESSION_AGENT_CACHE["other-session"] = (decoy, "sig")
    reached = threading.Event()
    release = threading.Event()
    state = {"steer_ident": None, "selected": None, "paused": False, "enqueue": None}
    order = []

    matches = streaming._cached_agent_matches_session

    def select(candidate, sid):
        result = matches(candidate, sid)
        if threading.get_ident() == state["steer_ident"] and result:
            state["selected"] = candidate
        return result

    def gate(edge):
        # Pause the Steer thread once, at the requested unlocked edge, only
        # after it has selected its cache candidate.
        if (threading.get_ident() != state["steer_ident"]
                or state["selected"] is None or state["paused"] or edge != gap):
            return
        state["paused"] = True
        reached.set()
        assert release.wait(5), "cache-only steer barrier timed out"

    def enqueue(text):
        state["enqueue"] = {
            "held": lock.owner == threading.get_ident(),
            "alive": "run" in config.STREAMS,
            "phase": config.ACTIVE_RUNS["run"]["phase"],
        }
        order.append("steer")
        return True

    def interrupt(reason):
        assert lock.owner != threading.get_ident(), "interrupt must be outside registry lock"
        order.append("cancel")

    def cache_only_steer():
        state["steer_ident"] = threading.get_ident()
        return steer()

    agent.steer.side_effect = enqueue
    agent.interrupt.side_effect = interrupt
    monkeypatch.setattr(streaming, "_cached_agent_matches_session", select)
    lock.before_acquire = lambda: gate("before_acquire")
    lock.after_release = lambda: gate("after_release")
    with ThreadPoolExecutor(max_workers=1) as pool:
        guidance = pool.submit(cache_only_steer)
        try:
            assert reached.wait(5), "steer never reached the requested gap"
            assert lock.owner is None, "steer must not pause while holding the stream lock"
            assert streaming.cancel_stream("run") is True
        finally:
            release.set()
        result = guidance.result(timeout=5)
    assert state["selected"] is agent, "the selected cache object must be the intended worker"
    decoy.steer.assert_not_called()
    assert config.ACTIVE_RUNS["run"]["phase"] == "cancelling"
    assert "run" not in config.STREAMS
    assert "run" not in config.AGENT_INSTANCES
    agent.interrupt.assert_called_once()
    if order == ["steer", "cancel"]:
        # Steer queued before Stop claimed cancellation, under the stream lock.
        assert result == {"accepted": True, "fallback": None, "stream_id": "run"}
        agent.steer.assert_called_once_with("guidance")
        assert state["enqueue"] == {"held": True, "alive": True, "phase": "running"}
    else:
        # Stop claimed first: the cached worker must not be steered afterwards.
        assert order == ["cancel"]
        assert result == {"accepted": False, "fallback": "stream_dead", "stream_id": None}
        agent.steer.assert_not_called()


def test_cache_only_selection_revalidates_phase_before_enqueue(scene, monkeypatch):
    """Finalization can win after initial resolution but before cache enqueue."""
    lock, agent = scene
    config.AGENT_INSTANCES.clear()
    matches = streaming._cached_agent_matches_session
    selected = []

    def select_then_finalize(candidate, sid):
        matched = matches(candidate, sid)
        if matched:
            selected.append(candidate)
            assert lock.owner != threading.get_ident()
            with config.STREAMS_LOCK:
                streaming.update_active_run("run", phase="finalizing")
        return matched

    monkeypatch.setattr(streaming, "_cached_agent_matches_session", select_then_finalize)
    assert steer() == {"accepted": False, "fallback": "not_running", "stream_id": "run"}
    assert selected == [agent]
    agent.steer.assert_not_called()