"""Tests for lifecycle hooks — fire order, error handling, per-tool-call semantics."""

from types import SimpleNamespace
from typing import Any, AsyncIterator, cast
from unittest.mock import AsyncMock, MagicMock

import pytest
from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
    Function,
)

from swarm.agent import Agent
from swarm.context import RunContext
from swarm.hooks import (
    LifecycleHook,
    fire_hook_on_llm_end,
    fire_hook_on_llm_start,
    fire_hook_on_tool_end,
    fire_hook_on_tool_start,
)
from swarm.providers.protocols import ModelRouterProtocol
from swarm.runners.agent import AgentRunner, AgentSignal

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


class _RecordingHook:
    """A concrete LifecycleHook that records every call."""

    def __init__(self) -> None:
        self.calls: list[tuple[str, tuple[Any, ...]]] = []

    async def on_llm_start(
        self, agent: Any, messages: list[Any], context_variables: dict[str, Any]
    ) -> None:
        self.calls.append(("on_llm_start", (agent, messages, context_variables)))

    async def on_llm_end(
        self,
        agent: Any,
        response: Any,
        error: BaseException | None = None,
    ) -> None:
        self.calls.append(("on_llm_end", (agent, response, error)))

    async def on_tool_start(
        self,
        agent: Any,
        tool_call_id: str,
        tool_name: str,
        arguments: dict[str, Any],
    ) -> None:
        self.calls.append(
            ("on_tool_start", (agent, tool_call_id, tool_name, arguments))
        )

    async def on_tool_end(
        self,
        agent: Any,
        tool_call_id: str,
        tool_name: str,
        result: Any,
        error: BaseException | None = None,
    ) -> None:
        self.calls.append(
            ("on_tool_end", (agent, tool_call_id, tool_name, result, error))
        )


class _ExplodingHook(_RecordingHook):
    """Hook that raises on every method — should never crash the runner."""

    async def on_llm_start(
        self, agent: Any, messages: Any, context_variables: Any
    ) -> None:
        await super().on_llm_start(agent, messages, context_variables)
        raise RuntimeError("boom in on_llm_start")

    async def on_llm_end(self, agent: Any, response: Any, error: Any = None) -> None:
        await super().on_llm_end(agent, response, error)
        raise RuntimeError("boom in on_llm_end")


def _make_agent(*, hooks: list[Any] | None = None) -> Agent:
    return Agent(
        name="HookTest",
        model="gpt-4o",
        instructions="test",
        tools=[],
        hooks=hooks or [],
    )


# Reuse fake infrastructure from test_agent_runner
class _FakeContext:
    conversation_id = "conv-1"
    current_context_variables: dict[str, Any] = {}
    workflow_context: dict[str, Any] = {}
    designer_fullname = "Agent"

    def get_history_for_completion(self) -> list[Any]:
        return []


class _FakeStreamManager:
    async def publish_status(self, status: str) -> None:
        pass


class _FakeMessageStore:
    async def persist_messages(self, **kwargs: Any) -> None:
        pass


class _FakeRouter:
    def __init__(self, responses: list[Any]) -> None:
        self._responses = list(responses)

    async def create_completion_async(self, request: Any) -> Any:
        return self._responses.pop(0)

    async def stream_completion(self, request: Any) -> AsyncIterator[Any]:
        _ = request
        raise AssertionError(
            "stream_completion should not be called in AgentRunner tests"
        )
        yield None  # pragma: no cover


def _make_completion(
    *, content: str = "done", tool_calls: Any = None, model: str = "gpt-4o"
) -> Any:
    return SimpleNamespace(
        choices=[
            SimpleNamespace(
                message=SimpleNamespace(content=content, tool_calls=tool_calls)
            )
        ],
        usage=None,
        model=model,
    )


def _make_run_context(router: _FakeRouter) -> RunContext[_FakeContext]:
    return RunContext(
        deps=_FakeContext(),
        stream_manager=_FakeStreamManager(),
        message_store=_FakeMessageStore(),
        router=router,
        retry_policy=None,
    )


# ---------------------------------------------------------------------------
# Protocol check
# ---------------------------------------------------------------------------


class TestLifecycleHookProtocol:
    def test_recording_hook_passes_isinstance_check(self) -> None:
        assert isinstance(_RecordingHook(), LifecycleHook)

    def test_non_hook_fails_isinstance_check(self) -> None:
        assert not isinstance("not a hook", LifecycleHook)


# ---------------------------------------------------------------------------
# Fire helpers
# ---------------------------------------------------------------------------


class TestFireHelpers:
    @pytest.mark.asyncio
    async def test_fire_on_llm_start(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        await fire_hook_on_llm_start(agent, ["msg"], {"key": "val"})
        assert len(hook.calls) == 1
        assert hook.calls[0][0] == "on_llm_start"

    @pytest.mark.asyncio
    async def test_fire_on_llm_end(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        await fire_hook_on_llm_end(agent, "response")
        assert len(hook.calls) == 1
        assert hook.calls[0][0] == "on_llm_end"

    @pytest.mark.asyncio
    async def test_fire_on_llm_end_with_error(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        err = RuntimeError("fail")
        await fire_hook_on_llm_end(agent, None, error=err)
        assert hook.calls[0][1][1] is None  # response is None on failure
        assert hook.calls[0][1][2] is err  # error is passed

    @pytest.mark.asyncio
    async def test_fire_on_tool_start(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        await fire_hook_on_tool_start(agent, "call-1", "my_tool", {"arg": 1})
        assert hook.calls[0] == (
            "on_tool_start",
            (agent, "call-1", "my_tool", {"arg": 1}),
        )

    @pytest.mark.asyncio
    async def test_fire_on_tool_end(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        await fire_hook_on_tool_end(agent, "call-1", "my_tool", "result_val")
        assert hook.calls[0][0] == "on_tool_end"
        assert hook.calls[0][1][3] == "result_val"
        assert hook.calls[0][1][4] is None  # no error

    @pytest.mark.asyncio
    async def test_hook_exception_does_not_propagate(self) -> None:
        hook = _ExplodingHook()
        agent = _make_agent(hooks=[hook])
        # Should NOT raise
        await fire_hook_on_llm_start(agent, [], {})
        await fire_hook_on_llm_end(agent, None)
        assert len(hook.calls) == 2

    @pytest.mark.asyncio
    async def test_invalid_hook_skipped_with_warning(self) -> None:
        agent = _make_agent(hooks=["not_a_hook"])
        # Should NOT raise, just skip
        await fire_hook_on_llm_start(agent, [], {})

    @pytest.mark.asyncio
    async def test_no_hooks_is_noop(self) -> None:
        agent = _make_agent(hooks=[])
        # Should NOT raise
        await fire_hook_on_llm_start(agent, [], {})
        await fire_hook_on_llm_end(agent, None)
        await fire_hook_on_tool_start(agent, "c", "t", {})
        await fire_hook_on_tool_end(agent, "c", "t", None)


# ---------------------------------------------------------------------------
# Batch runner integration — hooks fire around LLM calls
# ---------------------------------------------------------------------------


class TestAgentRunnerHooks:
    @pytest.mark.asyncio
    async def test_hooks_fire_on_llm_start_and_end(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )

        event_names = [c[0] for c in hook.calls]
        assert "on_llm_start" in event_names
        assert "on_llm_end" in event_names
        # on_llm_start fires before on_llm_end
        assert event_names.index("on_llm_start") < event_names.index("on_llm_end")

    @pytest.mark.asyncio
    async def test_on_llm_end_receives_completion_on_success(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])
        completion = _make_completion(content="hello")
        router = _FakeRouter([completion])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )

        llm_end_call = [c for c in hook.calls if c[0] == "on_llm_end"][0]
        # response should be the completion object
        assert llm_end_call[1][1] is completion
        # error should be None
        assert llm_end_call[1][2] is None

    @pytest.mark.asyncio
    async def test_on_llm_end_fires_with_error_on_failure(self) -> None:
        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])

        class _FailRouter:
            async def create_completion_async(self, req: Any) -> Any:
                _ = req
                raise RuntimeError("LLM failed")

            async def stream_completion(self, req: Any) -> AsyncIterator[Any]:
                _ = req
                raise AssertionError("stream_completion should not be called")
                yield None  # pragma: no cover

        run_context = RunContext(
            deps=_FakeContext(),
            stream_manager=_FakeStreamManager(),
            message_store=_FakeMessageStore(),
            router=cast(ModelRouterProtocol, _FailRouter()),
            retry_policy=None,
        )

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        with pytest.raises(RuntimeError, match="LLM failed"):
            await runner.run_batch(
                agent=agent, run_context=cast(Any, run_context), max_turns=1
            )

        llm_end_call = [c for c in hook.calls if c[0] == "on_llm_end"][0]
        assert llm_end_call[1][1] is None  # response is None
        assert isinstance(llm_end_call[1][2], RuntimeError)  # error is the exception

    @pytest.mark.asyncio
    async def test_agent_with_no_hooks_works_identically(self) -> None:
        agent = _make_agent(hooks=[])
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        result = await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )

        assert result.signal == AgentSignal.COMPLETED

    @pytest.mark.asyncio
    async def test_hook_exception_does_not_crash_runner(self) -> None:
        hook = _ExplodingHook()
        agent = _make_agent(hooks=[hook])
        router = _FakeRouter([_make_completion(content="done")])
        run_context = _make_run_context(router)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        # Should NOT raise despite hook explosions
        result = await runner.run_batch(
            agent=agent, run_context=cast(Any, run_context), max_turns=1
        )
        assert result.signal == AgentSignal.COMPLETED


# ---------------------------------------------------------------------------
# Tool hooks — per individual tool call (via execute_tools)
# ---------------------------------------------------------------------------


class TestToolHooksInExecuteTools:
    """Test that tool hooks fire per individual tool call including pre-exec failures."""

    @pytest.mark.asyncio
    async def test_unknown_tool_fires_hooks_with_error(self) -> None:
        from swarm.tool import execute_tools

        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])

        tool_call = ChatCompletionMessageToolCall(
            id="call-1",
            type="function",
            function=Function(name="nonexistent", arguments="{}"),
        )

        stream_mgr = AsyncMock()
        run_ctx = MagicMock()
        run_ctx.stream_manager = stream_mgr

        await execute_tools(
            [tool_call], [], run_ctx, "test-agent", hooks=agent.hooks, agent=agent
        )

        # Should have on_tool_start + on_tool_end with error
        events = [c[0] for c in hook.calls]
        assert events == ["on_tool_start", "on_tool_end"]
        # on_tool_end should have an error
        tool_end = hook.calls[1]
        assert isinstance(tool_end[1][4], ValueError)

    @pytest.mark.asyncio
    async def test_invalid_json_args_fires_hooks_with_error(self) -> None:
        from swarm.tool import execute_tools

        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])

        def dummy_tool() -> None:
            pass

        dummy_tool.__name__ = "dummy_tool"

        tool_call = ChatCompletionMessageToolCall(
            id="call-1",
            type="function",
            function=Function(name="dummy_tool", arguments="{invalid json}"),
        )

        stream_mgr = AsyncMock()
        run_ctx = MagicMock()
        run_ctx.stream_manager = stream_mgr

        await execute_tools(
            [tool_call],
            [dummy_tool],
            run_ctx,
            "test-agent",
            hooks=agent.hooks,
            agent=agent,
        )

        events = [c[0] for c in hook.calls]
        assert events == ["on_tool_start", "on_tool_end"]
        # on_tool_start should get empty args
        tool_start = hook.calls[0]
        assert tool_start[1][3] == {}  # arguments is {}
        # on_tool_end should have JSONDecodeError
        tool_end = hook.calls[1]
        assert tool_end[1][4] is not None  # error exists

    @pytest.mark.asyncio
    async def test_malformed_tool_call_no_hooks(self) -> None:
        """Malformed tool calls (missing function/name or missing id) skip hooks."""
        from swarm.tool import execute_tools

        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])

        # Missing function name
        malformed = ChatCompletionMessageToolCall(
            id="call-1",
            type="function",
            function=Function(name="", arguments="{}"),
        )

        stream_mgr = AsyncMock()
        run_ctx = MagicMock()
        run_ctx.stream_manager = stream_mgr

        await execute_tools(
            [malformed], [], run_ctx, "test-agent", hooks=agent.hooks, agent=agent
        )

        # No hooks should fire for malformed calls
        assert len(hook.calls) == 0

    @pytest.mark.asyncio
    async def test_multiple_tool_calls_fire_hooks_per_call(self) -> None:
        """Each valid tool call gets its own on_tool_start/on_tool_end pair."""
        from swarm.tool import execute_tools

        hook = _RecordingHook()
        agent = _make_agent(hooks=[hook])

        # Two unknown tools — both should fire hooks
        tc1 = ChatCompletionMessageToolCall(
            id="call-1",
            type="function",
            function=Function(name="tool_a", arguments="{}"),
        )
        tc2 = ChatCompletionMessageToolCall(
            id="call-2",
            type="function",
            function=Function(name="tool_b", arguments="{}"),
        )

        stream_mgr = AsyncMock()
        run_ctx = MagicMock()
        run_ctx.stream_manager = stream_mgr

        await execute_tools(
            [tc1, tc2], [], run_ctx, "test-agent", hooks=agent.hooks, agent=agent
        )

        events = [c[0] for c in hook.calls]
        # 2 tool calls × (start + end) = 4 events
        assert events == [
            "on_tool_start",
            "on_tool_end",
            "on_tool_start",
            "on_tool_end",
        ]

        # Verify correlation IDs
        assert hook.calls[0][1][1] == "call-1"  # first start: call-1
        assert hook.calls[1][1][1] == "call-1"  # first end: call-1
        assert hook.calls[2][1][1] == "call-2"  # second start: call-2
        assert hook.calls[3][1][1] == "call-2"  # second end: call-2
