"""Tests for scheduler context propagation in team response flows."""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager, suppress
from typing import TYPE_CHECKING
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from mindroom.bot import AgentBot
from mindroom.config.agent import AgentConfig
from mindroom.config.main import Config
from mindroom.config.models import ModelConfig, RouterConfig
from mindroom.constants import STREAM_STATUS_ERROR, STREAM_STATUS_KEY
from mindroom.dispatch_source import MESSAGE_SOURCE_KIND
from mindroom.final_delivery import StreamTransportOutcome
from mindroom.hooks import MessageEnvelope
from mindroom.matrix.client import DeliveredMatrixEvent
from mindroom.matrix.users import AgentMatrixUser
from mindroom.message_target import MessageTarget
from mindroom.orchestration.runtime import SYNC_RESTART_CANCEL_MSG
from mindroom.response_runner import ResponseRequest, ResponseRunner
from mindroom.streaming import _INTERRUPTED_RESPONSE_NOTE, build_restart_interrupted_body
from mindroom.tool_system.runtime_context import get_tool_runtime_context
from tests.conftest import (
    TEST_ACCESS_TOKEN,
    TEST_PASSWORD,
    bind_runtime_paths,
    install_runtime_cache_support,
    message_origin,
    patch_response_runner_module,
    runtime_paths_for,
    test_runtime_paths,
)
from tests.identity_helpers import entity_ids, persist_entity_accounts

if TYPE_CHECKING:
    from collections.abc import AsyncIterator
    from pathlib import Path

    from mindroom.matrix.identity import MatrixID


@asynccontextmanager
async def _noop_typing_indicator(*_args: object, **_kwargs: object) -> AsyncIterator[None]:
    yield


def _response_envelope() -> MessageEnvelope:
    return MessageEnvelope(
        source_event_id="$user_event",
        target=MessageTarget.resolve("!team:localhost", "$thread_root", "$user_event"),
        body="Please coordinate and schedule a reminder",
        attachment_ids=(),
        mentioned_agents=(),
        agent_name="general",
        origin=message_origin(
            sender_id="@user:localhost",
            requester_id="@user:localhost",
            source_kind=MESSAGE_SOURCE_KIND,
        ),
    )


def _make_bot(tmp_path: Path) -> AgentBot:
    runtime_paths = test_runtime_paths(tmp_path)
    config = bind_runtime_paths(
        Config(
            agents={
                "general": AgentConfig(display_name="General Agent", rooms=["!team:localhost"]),
                "research": AgentConfig(display_name="Research Agent", rooms=["!team:localhost"]),
            },
            models={"default": ModelConfig(provider="ollama", id="test-model")},
            router=RouterConfig(model="default"),
        ),
        runtime_paths,
    )
    persist_entity_accounts(config, runtime_paths)
    agent_user = AgentMatrixUser(
        agent_name="general",
        user_id=entity_ids(config, runtime_paths)["general"].full_id,
        display_name="General Agent",
        password=TEST_PASSWORD,
        access_token=TEST_ACCESS_TOKEN,
    )
    bot = AgentBot(
        agent_user=agent_user,
        storage_path=tmp_path,
        config=config,
        runtime_paths=runtime_paths_for(config),
        rooms=["!team:localhost"],
    )
    bot.client = AsyncMock()
    bot.client.user_id = agent_user.user_id
    bot.client.rooms = {"!team:localhost": MagicMock(room_id="!team:localhost")}
    bot.orchestrator = MagicMock(config=config)
    return install_runtime_cache_support(bot)


def _team_agents(bot: AgentBot) -> list[MatrixID]:
    ids = entity_ids(bot.config, runtime_paths_for(bot.config))
    return [ids["general"], ids["research"]]


@pytest.mark.asyncio
async def test_team_non_streaming_has_scheduler_context(tmp_path: Path) -> None:
    """Team non-streaming flow should expose scheduler context to tool calls."""
    bot = _make_bot(tmp_path)
    team_agents = _team_agents(bot)
    response_run_id: str | None = None

    async def fake_run_cancellable_response(**kwargs: object) -> None:
        nonlocal response_run_id
        assert isinstance(kwargs["run_id"], str)
        assert kwargs["run_id"]
        response_run_id = kwargs["run_id"]
        response_function = kwargs["response_function"]
        await response_function(None)

    async def fake_team_response(*_args: object, **_kwargs: object) -> str:
        assert get_tool_runtime_context() is not None
        assert _kwargs["ctx"].session_id == "!team:localhost:$thread_root"
        assert _kwargs["user_id"] == "@user:localhost"
        assert _kwargs["ctx"].run_id == response_run_id
        return "team non-streaming response"

    with (
        patch.object(
            ResponseRunner,
            "_run_cancellable_response",
            new=AsyncMock(side_effect=fake_run_cancellable_response),
        ),
        patch("mindroom.response_runner.typing_indicator", new=_noop_typing_indicator),
        patch_response_runner_module(
            should_use_streaming=AsyncMock(return_value=False),
            team_response=fake_team_response,
        ),
    ):
        await bot._response_runner.generate_team_response_helper(
            ResponseRequest(
                thread_history=[],
                prompt="Please coordinate and schedule a reminder",
                user_id="@user:localhost",
                response_envelope=_response_envelope(),
                correlation_id="corr-team-non-streaming",
            ),
            team_agents=team_agents,
            team_mode="coordinate",
        )


@pytest.mark.asyncio
async def test_team_non_streaming_cancellation_edits_placeholder(tmp_path: Path) -> None:
    """Generic team interruptions should replace the thinking placeholder with an interruption note."""
    bot = _make_bot(tmp_path)
    team_agents = _team_agents(bot)

    async def fake_run_cancellable_response(**kwargs: object) -> None:
        response_function = kwargs["response_function"]
        with suppress(asyncio.CancelledError):
            await response_function("$thinking")

    async def fake_team_response(*_args: object, **_kwargs: object) -> str:
        raise asyncio.CancelledError

    with (
        patch.object(
            ResponseRunner,
            "_run_cancellable_response",
            new=AsyncMock(side_effect=fake_run_cancellable_response),
        ),
        patch("mindroom.response_runner.typing_indicator", new=_noop_typing_indicator),
        patch(
            "mindroom.delivery_gateway.edit_message_result",
            new=AsyncMock(
                return_value=DeliveredMatrixEvent(
                    event_id="$thinking",
                    content_sent={"body": _INTERRUPTED_RESPONSE_NOTE},
                ),
            ),
        ) as mock_edit,
        patch_response_runner_module(
            should_use_streaming=AsyncMock(return_value=False),
            team_response=fake_team_response,
        ),
    ):
        await bot._response_runner.generate_team_response_helper(
            ResponseRequest(
                thread_history=[],
                prompt="Please coordinate and schedule a reminder",
                user_id="@user:localhost",
                response_envelope=_response_envelope(),
                correlation_id="corr-team-cancelled",
            ),
            team_agents=team_agents,
            team_mode="coordinate",
        )

    assert mock_edit.await_args.args[2] == "$thinking"
    assert mock_edit.await_args.args[4] == _INTERRUPTED_RESPONSE_NOTE
    assert mock_edit.await_args.args[3][STREAM_STATUS_KEY] == STREAM_STATUS_ERROR


@pytest.mark.asyncio
async def test_team_non_streaming_sync_restart_edits_placeholder_with_restart_note(tmp_path: Path) -> None:
    """Sync restarts should mark team placeholders as interrupted, not user-cancelled."""
    bot = _make_bot(tmp_path)
    team_agents = _team_agents(bot)

    async def fake_run_cancellable_response(**kwargs: object) -> None:
        response_function = kwargs["response_function"]
        with suppress(asyncio.CancelledError):
            await response_function("$thinking")

    async def fake_team_response(*_args: object, **_kwargs: object) -> str:
        raise asyncio.CancelledError(SYNC_RESTART_CANCEL_MSG)

    with (
        patch.object(
            ResponseRunner,
            "_run_cancellable_response",
            new=AsyncMock(side_effect=fake_run_cancellable_response),
        ),
        patch("mindroom.response_runner.typing_indicator", new=_noop_typing_indicator),
        patch(
            "mindroom.delivery_gateway.edit_message_result",
            new=AsyncMock(
                return_value=DeliveredMatrixEvent(
                    event_id="$thinking",
                    content_sent={"body": build_restart_interrupted_body("")},
                ),
            ),
        ) as mock_edit,
        patch_response_runner_module(
            should_use_streaming=AsyncMock(return_value=False),
            team_response=fake_team_response,
        ),
    ):
        await bot._response_runner.generate_team_response_helper(
            ResponseRequest(
                thread_history=[],
                prompt="Please coordinate and schedule a reminder",
                user_id="@user:localhost",
                response_envelope=_response_envelope(),
                correlation_id="corr-team-restart",
            ),
            team_agents=team_agents,
            team_mode="coordinate",
        )

    assert mock_edit.await_args.args[2] == "$thinking"
    assert mock_edit.await_args.args[4] == build_restart_interrupted_body("")
    assert mock_edit.await_args.args[3][STREAM_STATUS_KEY] == STREAM_STATUS_ERROR


@pytest.mark.asyncio
async def test_team_streaming_has_scheduler_context(tmp_path: Path) -> None:
    """Team streaming flow should expose scheduler context to tool calls."""
    bot = _make_bot(tmp_path)
    team_agents = _team_agents(bot)
    response_run_id: str | None = None

    async def fake_run_cancellable_response(**kwargs: object) -> None:
        nonlocal response_run_id
        assert isinstance(kwargs["run_id"], str)
        assert kwargs["run_id"]
        response_run_id = kwargs["run_id"]
        response_function = kwargs["response_function"]
        await response_function(None)

    async def fake_send_streaming_response(*args: object, **_kwargs: object) -> StreamTransportOutcome:
        response_stream = args[4]
        chunks = [str(chunk) async for chunk in response_stream]
        return StreamTransportOutcome(
            last_physical_stream_event_id="$stream_event",
            terminal_status="completed",
            rendered_body="".join(chunks),
            visible_body_state="visible_body",
        )

    async def fake_team_response_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]:
        assert get_tool_runtime_context() is not None
        assert _kwargs["ctx"].session_id == "!team:localhost:$thread_root"
        assert _kwargs["user_id"] == "@user:localhost"
        assert _kwargs["ctx"].run_id == response_run_id
        yield "stream chunk"

    with (
        patch.object(
            ResponseRunner,
            "_run_cancellable_response",
            new=AsyncMock(side_effect=fake_run_cancellable_response),
        ),
        patch("mindroom.response_runner.typing_indicator", new=_noop_typing_indicator),
        patch(
            "mindroom.delivery_gateway.send_streaming_response",
            new=AsyncMock(side_effect=fake_send_streaming_response),
        ),
        patch_response_runner_module(
            should_use_streaming=AsyncMock(return_value=True),
            team_response_stream=fake_team_response_stream,
        ),
    ):
        await bot._response_runner.generate_team_response_helper(
            ResponseRequest(
                thread_history=[],
                prompt="Please collaborate and schedule a reminder",
                user_id="@user:localhost",
                response_envelope=_response_envelope(),
                correlation_id="corr-team-streaming",
            ),
            team_agents=team_agents,
            team_mode="collaborate",
        )


@pytest.mark.asyncio
async def test_team_late_cancellation_during_post_effects_propagates(tmp_path: Path) -> None:
    """Late cancellation should still cancel the team path while post-effects are running."""
    bot = _make_bot(tmp_path)
    team_agents = _team_agents(bot)
    started = asyncio.Event()
    release = asyncio.Event()

    async def fake_run_cancellable_response(**kwargs: object) -> str:
        response_function = kwargs["response_function"]
        await response_function(None)
        return "$team_response"

    async def fake_post_effects(*_args: object, **_kwargs: object) -> None:
        started.set()
        await release.wait()

    with (
        patch.object(
            ResponseRunner,
            "_run_cancellable_response",
            new=AsyncMock(side_effect=fake_run_cancellable_response),
        ),
        patch("mindroom.response_runner.typing_indicator", new=_noop_typing_indicator),
        patch_response_runner_module(
            should_use_streaming=AsyncMock(return_value=False),
            team_response=AsyncMock(return_value="team non-streaming response"),
            apply_post_response_effects=AsyncMock(side_effect=fake_post_effects),
        ),
    ):
        task = asyncio.create_task(
            bot._response_runner.generate_team_response_helper(
                ResponseRequest(
                    thread_history=[],
                    prompt="Please coordinate and schedule a reminder",
                    user_id="@user:localhost",
                    response_envelope=_response_envelope(),
                    correlation_id="corr-team-late-cancel",
                ),
                team_agents=team_agents,
                team_mode="coordinate",
            ),
        )
        await started.wait()
        task.cancel()
        release.set()
        with pytest.raises(asyncio.CancelledError):
            await task
