"""Tests for SubagentManager lifecycle — spawn, run, announce, cancel."""

import asyncio
import time
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from nanobot.agent import SubagentManager
from nanobot.agent.hook import AgentHookContext
from nanobot.agent.runner import AgentRunResult
from nanobot.agent.subagent import (
    SubagentStatus,
    _SubagentHook,
)
from nanobot.agent.tools.context import current_request_context
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import GenerationSettings, LLMProvider, LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _manager(tmp_path: Path, **kw) -> SubagentManager:
    defaults = dict(
        workspace=tmp_path,
        bus=MessageBus(),
        max_tool_result_chars=16_000,
    )
    defaults.update(kw)
    return SubagentManager(**defaults)


def _runtime(*, model: str = "test-model", temperature: float = 0.1) -> LLMRuntime:
    provider = MagicMock(spec=LLMProvider)
    provider.generation = GenerationSettings(temperature=temperature, max_tokens=4096)
    return LLMRuntime.capture(
        provider,
        model,
        context_window_tokens=128_000,
    )


def _make_hook_context(**overrides) -> AgentHookContext:
    defaults = dict(
        iteration=1,
        tool_calls=[],
        tool_events=[],
        messages=[],
        usage=None,
        error=None,
        stop_reason="completed",
        final_content="ok",
    )
    defaults.update(overrides)
    return AgentHookContext(**defaults)


async def _drain_subagent_tasks(sm: SubagentManager) -> None:
    tasks = list(sm._running_tasks.values())
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)
    await asyncio.sleep(0)


@pytest.mark.asyncio
async def test_close_cancels_tasks_before_closing_exec_sessions(tmp_path):
    sm = _manager(tmp_path)
    task = asyncio.create_task(asyncio.Event().wait())
    sm._running_tasks["t1"] = task

    async def close_exec_sessions() -> int:
        assert task.done()
        return 0

    sm._exec_session_manager.close_all = AsyncMock(side_effect=close_exec_sessions)

    await sm.close()

    assert task.cancelled()
    sm._exec_session_manager.close_all.assert_awaited_once()


# ---------------------------------------------------------------------------
# SubagentStatus defaults
# ---------------------------------------------------------------------------


class TestSubagentStatus:
    def test_defaults(self):
        s = SubagentStatus(
            task_id="abc", label="test", task_description="do stuff",
            started_at=time.monotonic(),
        )
        assert s.phase == "initializing"
        assert s.iteration == 0
        assert s.tool_events == []
        assert s.usage is None
        assert s.stop_reason is None
        assert s.error is None


# ---------------------------------------------------------------------------
# Runtime ownership
# ---------------------------------------------------------------------------


class TestRuntimeOwnership:
    def test_manager_has_no_provider_model_mirrors(self, tmp_path):
        sm = _manager(tmp_path)
        assert not hasattr(sm, "provider")
        assert not hasattr(sm, "model")
        assert not hasattr(sm, "context_window_tokens")
        assert not hasattr(sm.runner, "provider")


class TestLegacyCompatibility:
    def test_accepts_exported_legacy_constructor_positionally(self, tmp_path):
        provider = MagicMock(spec=LLMProvider)
        provider.generation = GenerationSettings(temperature=0.2, max_tokens=2048)

        with pytest.warns(DeprecationWarning, match="provider"):
            sm = SubagentManager(
                provider,
                tmp_path,
                MessageBus(),
                16_000,
                "legacy-model",
            )

        assert sm.workspace == tmp_path
        assert sm.max_tool_result_chars == 16_000
        assert not hasattr(sm, "provider")
        assert not hasattr(sm, "model")

    @pytest.mark.asyncio
    async def test_legacy_spawn_captures_runtime_at_admission(self, tmp_path):
        provider = MagicMock(spec=LLMProvider)
        provider.generation = GenerationSettings(temperature=0.2, max_tokens=2048)
        with pytest.warns(DeprecationWarning, match="provider"):
            sm = SubagentManager(
                provider=provider,
                workspace=tmp_path,
                bus=MessageBus(),
                max_tool_result_chars=16_000,
                model="legacy-model",
            )
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="done", messages=[], stop_reason="completed",
        ))
        provider.generation = GenerationSettings(temperature=0.8, max_tokens=512)

        with pytest.warns(DeprecationWarning, match="runtime"):
            await sm.spawn("legacy task")
        await _drain_subagent_tasks(sm)

        runtime = sm.runner.run.await_args.args[0].runtime
        assert runtime.provider is provider
        assert runtime.model == "legacy-model"
        assert runtime.generation == GenerationSettings(0.8, 512, None)

    @pytest.mark.asyncio
    async def test_set_provider_supports_future_legacy_spawns(self, tmp_path):
        sm = _manager(tmp_path)
        provider = MagicMock(spec=LLMProvider)
        provider.generation = GenerationSettings(temperature=0.3, max_tokens=1024)
        with pytest.warns(DeprecationWarning, match="set_provider"):
            sm.set_provider(provider, "replacement-model")
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="done", messages=[], stop_reason="completed",
        ))

        with pytest.warns(DeprecationWarning, match="runtime"):
            await sm.spawn("legacy task")
        await _drain_subagent_tasks(sm)

        runtime = sm.runner.run.await_args.args[0].runtime
        assert runtime.provider is provider
        assert runtime.model == "replacement-model"

    @pytest.mark.asyncio
    async def test_new_constructor_still_requires_explicit_spawn_runtime(self, tmp_path):
        sm = _manager(tmp_path)

        with pytest.raises(TypeError, match="runtime"):
            await sm.spawn("task")


# ---------------------------------------------------------------------------
# spawn
# ---------------------------------------------------------------------------


class TestSpawn:
    @pytest.mark.asyncio
    async def test_returns_string_with_task_id(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="done", messages=[], stop_reason="completed",
        ))
        result = await sm.spawn("do something", runtime=_runtime())
        assert "started" in result
        assert "id:" in result

    @pytest.mark.asyncio
    async def test_creates_task_in_running_tasks(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        await sm.spawn("task", runtime=_runtime(), session_key="s1")
        assert len(sm._running_tasks) == 1

        block.set()
        await _drain_subagent_tasks(sm)
        assert len(sm._running_tasks) == 0

    @pytest.mark.asyncio
    async def test_creates_status(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="done", messages=[], stop_reason="completed",
        ))
        await sm.spawn("my task", runtime=_runtime())
        await _drain_subagent_tasks(sm)
        # Status cleaned up after task completes
        assert len(sm._task_statuses) == 0

    @pytest.mark.asyncio
    async def test_registers_in_session_tasks(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        await sm.spawn("task", runtime=_runtime(), session_key="s1")
        assert "s1" in sm._session_tasks
        assert len(sm._session_tasks["s1"]) == 1

        block.set()
        await _drain_subagent_tasks(sm)
        assert "s1" not in sm._session_tasks

    @pytest.mark.asyncio
    async def test_no_session_key_no_registration(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        await sm.spawn("task", runtime=_runtime())
        assert len(sm._session_tasks) == 0

        block.set()
        await _drain_subagent_tasks(sm)

    @pytest.mark.asyncio
    async def test_label_defaults_to_truncated_task(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        long_label_source = "A" * 50
        await sm.spawn(long_label_source, runtime=_runtime(), session_key="s1")
        status = next(iter(sm._task_statuses.values()))
        assert status.label == long_label_source[:30] + "..."

        block.set()
        await _drain_subagent_tasks(sm)

    @pytest.mark.asyncio
    async def test_custom_label(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        await sm.spawn(
            "task", runtime=_runtime(), label="Custom Label", session_key="s1"
        )
        status = next(iter(sm._task_statuses.values()))
        assert status.label == "Custom Label"

        block.set()
        await _drain_subagent_tasks(sm)

    @pytest.mark.asyncio
    async def test_cleanup_callback_removes_all_entries(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="done", messages=[], stop_reason="completed",
        ))
        await sm.spawn("task", runtime=_runtime(), session_key="s1")
        await _drain_subagent_tasks(sm)
        assert len(sm._running_tasks) == 0
        assert len(sm._task_statuses) == 0
        assert len(sm._session_tasks) == 0

    @pytest.mark.asyncio
    async def test_runtime_is_captured_before_background_task_starts(self, tmp_path):
        sm = _manager(tmp_path)
        runtime = _runtime(temperature=0.2)
        entered = asyncio.Event()
        release = asyncio.Event()
        seen: dict[str, object] = {}

        async def observe(spec):
            seen["spec_runtime"] = spec.runtime
            request_ctx = current_request_context()
            seen["context_runtime"] = request_ctx.runtime if request_ctx else None
            entered.set()
            await release.wait()
            return AgentRunResult(
                final_content="done",
                messages=[],
                stop_reason="completed",
            )

        sm.runner.run = observe
        await sm.spawn("task", runtime=runtime, session_key="s1")
        runtime.provider.generation = GenerationSettings(
            temperature=0.9,
            max_tokens=128,
        )
        await asyncio.wait_for(entered.wait(), timeout=1)

        assert seen["spec_runtime"] is runtime
        assert seen["context_runtime"] is runtime
        assert runtime.generation.temperature == 0.2

        release.set()
        await _drain_subagent_tasks(sm)


# ---------------------------------------------------------------------------
# _run_subagent
# ---------------------------------------------------------------------------


class TestRunSubagent:
    @pytest.mark.asyncio
    async def test_successful_run(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="Task done!", messages=[], stop_reason="completed",
        ))
        with patch.object(sm, "_announce_result", new_callable=AsyncMock) as mock_announce:
            await sm._run_subagent(
                "t1", "do task", "label",
                {"channel": "cli", "chat_id": "direct"},
                SubagentStatus(task_id="t1", label="label", task_description="do task", started_at=time.monotonic()),
                _runtime(),
            )
            mock_announce.assert_called_once()
            assert mock_announce.call_args.args[-2] == "ok"

    @pytest.mark.asyncio
    async def test_exception_run(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(side_effect=RuntimeError("LLM down"))
        status = SubagentStatus(task_id="t1", label="label", task_description="do task", started_at=time.monotonic())
        with patch.object(sm, "_announce_result", new_callable=AsyncMock) as mock_announce:
            await sm._run_subagent(
                "t1", "do task", "label",
                {"channel": "cli", "chat_id": "direct"}, status, _runtime(),
            )
            assert status.phase == "error"
            assert "LLM down" in status.error
            assert mock_announce.call_args.args[-2] == "error"

    @pytest.mark.asyncio
    async def test_status_updated_on_success(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="ok", messages=[], stop_reason="completed",
        ))
        status = SubagentStatus(task_id="t1", label="label", task_description="do task", started_at=time.monotonic())
        with patch.object(sm, "_announce_result", new_callable=AsyncMock):
            await sm._run_subagent(
                "t1", "do task", "label",
                {"channel": "cli", "chat_id": "direct"}, status, _runtime(),
            )
            assert status.phase == "done"
            assert status.stop_reason == "completed"


# ---------------------------------------------------------------------------
# _announce_result
# ---------------------------------------------------------------------------


class TestAnnounceResult:
    @pytest.mark.asyncio
    async def test_publishes_inbound_message(self, tmp_path):
        sm = _manager(tmp_path)
        published = []
        sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg))

        await sm._announce_result(
            "t1", "label", "task", "result text",
            {"channel": "cli", "chat_id": "direct"}, "ok",
        )

        assert len(published) == 1
        msg = published[0]
        assert msg.channel == "system"
        assert msg.sender_id == "subagent"
        assert msg.metadata["injected_event"] == "subagent_result"
        assert msg.metadata["subagent_task_id"] == "t1"

    @pytest.mark.asyncio
    async def test_session_key_override(self, tmp_path):
        sm = _manager(tmp_path)
        published = []
        sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg))

        await sm._announce_result(
            "t1", "label", "task", "result",
            {"channel": "telegram", "chat_id": "123", "session_key": "s1"}, "ok",
        )

        assert published[0].session_key_override == "s1"

    @pytest.mark.asyncio
    async def test_session_key_override_fallback(self, tmp_path):
        sm = _manager(tmp_path)
        published = []
        sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg))

        await sm._announce_result(
            "t1", "label", "task", "result",
            {"channel": "telegram", "chat_id": "123"}, "ok",
        )

        assert published[0].session_key_override == "telegram:123"

    @pytest.mark.asyncio
    async def test_ok_status_text(self, tmp_path):
        sm = _manager(tmp_path)
        published = []
        sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg))

        await sm._announce_result(
            "t1", "label", "task", "result",
            {"channel": "cli", "chat_id": "direct"}, "ok",
        )

        assert "completed successfully" in published[0].content

    @pytest.mark.asyncio
    async def test_error_status_text(self, tmp_path):
        sm = _manager(tmp_path)
        published = []
        sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg))

        await sm._announce_result(
            "t1", "label", "task", "error details",
            {"channel": "cli", "chat_id": "direct"}, "error",
        )

        assert "failed" in published[0].content

    @pytest.mark.asyncio
    async def test_origin_message_id_in_metadata(self, tmp_path):
        sm = _manager(tmp_path)
        published = []
        sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg))

        await sm._announce_result(
            "t1", "label", "task", "result",
            {"channel": "cli", "chat_id": "direct"}, "ok",
            origin_message_id="msg-123",
        )

        assert published[0].metadata["origin_message_id"] == "msg-123"


# ---------------------------------------------------------------------------
# cancel_by_session
# ---------------------------------------------------------------------------


class TestCancelBySession:
    @pytest.mark.asyncio
    async def test_cancels_running_tasks(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        runtime = _runtime()
        await sm.spawn("task1", runtime=runtime, session_key="s1")
        await sm.spawn("task2", runtime=runtime, session_key="s1")
        assert len(sm._session_tasks.get("s1", set())) == 2

        count = await sm.cancel_by_session("s1")
        assert count == 2
        block.set()
        await _drain_subagent_tasks(sm)

    @pytest.mark.asyncio
    async def test_no_tasks_returns_zero(self, tmp_path):
        sm = _manager(tmp_path)
        count = await sm.cancel_by_session("nonexistent")
        assert count == 0

    @pytest.mark.asyncio
    async def test_cancels_active_and_queued_tasks(self, tmp_path):
        sm = _manager(tmp_path, max_concurrent_subagents=1)
        active_entered = asyncio.Event()
        queued_entered = asyncio.Event()

        async def _blocked_run(spec):
            task = spec.initial_messages[-1]["content"]
            if task == "active":
                active_entered.set()
            else:
                queued_entered.set()
            await asyncio.Event().wait()

        sm.runner.run = _blocked_run
        runtime = _runtime()
        await sm.spawn("active", runtime=runtime, session_key="s1")
        await asyncio.wait_for(active_entered.wait(), timeout=1.0)
        await sm.spawn("queued", runtime=runtime, session_key="s1")
        await asyncio.sleep(0)

        assert not queued_entered.is_set()
        assert await sm.cancel_by_session("s1") == 2
        await asyncio.sleep(0)

        assert not queued_entered.is_set()
        assert sm._running_tasks == {}
        assert sm._session_tasks == {}

    @pytest.mark.asyncio
    async def test_already_done_not_counted(self, tmp_path):
        sm = _manager(tmp_path)
        sm.runner.run = AsyncMock(return_value=AgentRunResult(
            final_content="done", messages=[], stop_reason="completed",
        ))
        await sm.spawn("task1", runtime=_runtime(), session_key="s1")
        await _drain_subagent_tasks(sm)

        count = await sm.cancel_by_session("s1")
        assert count == 0


# ---------------------------------------------------------------------------
# get_running_count / get_running_count_by_session
# ---------------------------------------------------------------------------


class TestRunningCounts:
    @pytest.mark.asyncio
    async def test_running_count_zero(self, tmp_path):
        sm = _manager(tmp_path)
        assert sm.get_running_count() == 0

    @pytest.mark.asyncio
    async def test_running_count_tracks_tasks(self, tmp_path):
        sm = _manager(tmp_path)
        block = asyncio.Event()
        async def _slow_run(spec):
            await block.wait()
            return AgentRunResult(final_content="done", messages=[], stop_reason="completed")
        sm.runner.run = _slow_run

        runtime = _runtime()
        await sm.spawn("t1", runtime=runtime, session_key="s1")
        await sm.spawn("t2", runtime=runtime, session_key="s1")
        assert sm.get_running_count() == 2
        assert sm.get_running_count_by_session("s1") == 2

        block.set()
        await _drain_subagent_tasks(sm)
        assert sm.get_running_count() == 0

    @pytest.mark.asyncio
    async def test_running_count_by_session_nonexistent(self, tmp_path):
        sm = _manager(tmp_path)
        assert sm.get_running_count_by_session("nonexistent") == 0


# ---------------------------------------------------------------------------
# _SubagentHook
# ---------------------------------------------------------------------------


class TestSubagentHook:
    @pytest.mark.asyncio
    async def test_before_execute_tools_logs(self, tmp_path):
        hook = _SubagentHook("t1")
        tool_call = MagicMock()
        tool_call.name = "read_file"
        tool_call.arguments = {"path": "/tmp/test"}
        ctx = _make_hook_context(tool_calls=[tool_call])
        result = await hook.before_execute_tools(ctx)
        assert result is None
        assert ctx.tool_calls == [tool_call]

    @pytest.mark.asyncio
    async def test_after_iteration_updates_status(self):
        status = SubagentStatus(
            task_id="t1", label="test", task_description="do", started_at=time.monotonic(),
        )
        hook = _SubagentHook("t1", status)
        ctx = _make_hook_context(
            iteration=3,
            tool_events=[{"name": "read_file", "status": "ok", "detail": ""}],
            usage=LLMUsage.reported(input_tokens=100, output_tokens=0),
        )
        await hook.after_iteration(ctx)
        assert status.iteration == 3
        assert len(status.tool_events) == 1
        assert status.usage == LLMUsage.reported(input_tokens=100, output_tokens=0)

    @pytest.mark.asyncio
    async def test_after_iteration_no_status_noop(self):
        hook = _SubagentHook("t1", status=None)
        ctx = _make_hook_context(iteration=5)
        result = await hook.after_iteration(ctx)
        assert result is None
        assert ctx.iteration == 5

    @pytest.mark.asyncio
    async def test_after_iteration_sets_error(self):
        status = SubagentStatus(
            task_id="t1", label="test", task_description="do", started_at=time.monotonic(),
        )
        hook = _SubagentHook("t1", status)
        ctx = _make_hook_context(error="something broke")
        await hook.after_iteration(ctx)
        assert status.error == "something broke"
