import json
from unittest.mock import MagicMock

import pytest

import lemurian.instrumentation as inst
from lemurian.agent import Agent
from lemurian.context import Context
from lemurian.handoff import Handoff
from lemurian.message import Message, MessageRole, ToolCallResultMessage
from lemurian.runner import Runner
from lemurian.session import Session
from lemurian.state import State
from lemurian.tools import LLMRecoverableError, tool

from lemurian.provider import CompletionResult

from tests.conftest import (
    MockCapability,
    MockFunction,
    MockToolCall,
    make_multi_tool_call_response,
    make_text_response,
    make_tool_call_response,
)


# ---------------------------------------------------------------------------
# Tool fixtures (module-level, reused across test classes)
# ---------------------------------------------------------------------------

@tool
def echo(text: str):
    """Echo text back."""
    return text


@tool
def get_data():
    """Return structured data."""
    return {"items": [1, 2, 3]}


@tool
def context_reader(context: Context, query: str):
    """Reads agent name from context."""
    return f"agent={context.agent.name}, query={query}"


@tool
async def async_echo(text: str):
    """Async echo."""
    return f"async: {text}"


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

class TestRunnerSimpleResponse:
    @pytest.mark.asyncio
    async def test_appends_assistant_message(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [make_text_response("Hello!")]
        agent = make_agent()
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Hello!"
        assert result.last_message.role == MessageRole.ASSISTANT
        assert result.hand_off is None
        assert len(session.transcript) == 1

    @pytest.mark.asyncio
    async def test_system_prompt_not_in_transcript(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [make_text_response("Hi")]
        agent = make_agent(system_prompt="Be helpful.")
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        for msg in session.transcript:
            assert msg.role != MessageRole.SYSTEM

        sent_messages = mock_provider.call_log[0]["messages"]
        assert sent_messages[0]["role"] == "system"
        assert sent_messages[0]["content"] == "Be helpful."


class TestRunnerToolCalls:
    @pytest.mark.asyncio
    async def test_tool_call_then_response(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response("echo", {"text": "hi"}),
            make_text_response("Done"),
        ]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Done"
        # transcript: tool_call_request, tool_result, assistant
        assert len(session.transcript) == 3
        assert session.transcript[1].role == MessageRole.TOOL
        assert session.transcript[1].content == "hi"

    @pytest.mark.asyncio
    async def test_context_injection(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response(
                "context_reader", {"query": "test"}
            ),
            make_text_response("Done"),
        ]
        agent = make_agent(name="myagent", tools=[context_reader])
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        tool_output = session.transcript[1].content
        assert "agent=myagent" in tool_output
        assert "query=test" in tool_output

    @pytest.mark.asyncio
    async def test_tool_not_found(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response("nonexistent", {}),
            make_text_response("Recovered"),
        ]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Recovered"
        error_msg = session.transcript[1]
        assert isinstance(error_msg, ToolCallResultMessage)
        assert "not found" in error_msg.content

    @pytest.mark.asyncio
    async def test_non_string_output_serialized(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response("get_data", {}),
            make_text_response("Done"),
        ]
        agent = make_agent(tools=[get_data])
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        tool_result = session.transcript[1]
        assert json.loads(tool_result.content) == {
            "items": [1, 2, 3]
        }

    @pytest.mark.asyncio
    async def test_async_tool(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response(
                "async_echo", {"text": "hello"}
            ),
            make_text_response("Done"),
        ]
        agent = make_agent(tools=[async_echo])
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        assert session.transcript[1].content == "async: hello"


class TestRunnerHandoff:
    @pytest.mark.asyncio
    async def test_handoff_detection(
        self, make_agent, mock_provider
    ):
        """Runner classifies a handoff tool call by name map."""
        billing_handoff = Handoff(
            tool_name="transfer_to_billing",
            tool_description="Hand off to billing.",
            target_agent="billing",
        )
        mock_provider.responses = [
            make_tool_call_response(
                "transfer_to_billing",
                {"message": "Customer needs billing help"},
            ),
        ]
        agent = make_agent()
        session = Session(session_id="s1")

        result = await Runner().run(
            agent, session, State(),
            handoffs=[billing_handoff],
        )

        assert result.hand_off is not None
        assert result.hand_off.target_agent == "billing"
        assert (
            result.hand_off.message
            == "Customer needs billing help"
        )
        assert (
            "Transferring to billing"
            in session.transcript[-1].content
        )


class TestRunnerWindowing:
    @pytest.mark.asyncio
    async def test_context_start_windows_transcript(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [make_text_response("Reply")]
        agent = make_agent()
        session = Session(session_id="s1")
        session.transcript = [
            Message(role=MessageRole.USER, content="old msg 0"),
            Message(
                role=MessageRole.ASSISTANT, content="old msg 1"
            ),
            Message(role=MessageRole.USER, content="new msg"),
        ]

        await Runner().run(agent, session, State(), context_start=2)

        sent_messages = mock_provider.call_log[0]["messages"]
        # system prompt + transcript[2:]
        assert len(sent_messages) == 2
        assert sent_messages[0]["role"] == "system"
        assert sent_messages[1]["content"] == "new msg"


class TestRunnerMaxTurns:
    @pytest.mark.asyncio
    async def test_max_turns_exceeded(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response("echo", {"text": "loop"})
            for _ in range(5)
        ]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        result = await Runner(max_turns=3).run(
            agent, session, State()
        )

        assert "Maximum turns" in result.last_message.content
        assert result.hand_off is None


# ---------------------------------------------------------------------------
# Error paths
# ---------------------------------------------------------------------------

class TestRunnerErrorPaths:
    @pytest.mark.asyncio
    async def test_tool_exception_recorded_in_transcript(
        self, make_agent, mock_provider
    ):
        """When a tool raises, the error is recorded and the loop continues."""

        @tool
        def explode(msg: str):
            """Always fails."""
            raise RuntimeError(f"boom: {msg}")

        mock_provider.responses = [
            make_tool_call_response("explode", {"msg": "test"}),
            make_text_response("Recovered"),
        ]
        agent = make_agent(tools=[explode])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Recovered"
        error_msg = session.transcript[1]
        assert isinstance(error_msg, ToolCallResultMessage)
        assert "boom: test" in error_msg.content

    @pytest.mark.asyncio
    async def test_recoverable_error_injects_guidance(
        self, make_agent, mock_provider
    ):
        """LLMRecoverableError sends the tool's message back as context."""

        @tool
        def strict_lookup(user_id: int):
            """Looks up a user by ID."""
            if user_id <= 0:
                raise LLMRecoverableError(
                    "user_id must be positive. Use get_user_id first."
                )
            return f"User {user_id}"

        mock_provider.responses = [
            make_tool_call_response(
                "strict_lookup", {"user_id": -1}
            ),
            make_tool_call_response(
                "strict_lookup", {"user_id": 42}
            ),
            make_text_response("Found the user"),
        ]
        agent = make_agent(tools=[strict_lookup])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Found the user"
        # First call's result is the recovery guidance
        recovery_msg = session.transcript[1]
        assert isinstance(recovery_msg, ToolCallResultMessage)
        assert "user_id must be positive" in recovery_msg.content
        # No "Error calling" prefix — it's the raw guidance
        assert not recovery_msg.content.startswith("Error")
        # Second call succeeded
        success_msg = session.transcript[3]
        assert success_msg.content == "User 42"

    @pytest.mark.asyncio
    async def test_invalid_json_arguments_recorded_in_transcript(
        self, make_agent, mock_provider
    ):
        """Malformed JSON in tool_call arguments is recorded as an error."""
        bad_response = CompletionResult(
            tool_calls=[
                MockToolCall(
                    id="call_1",
                    type="function",
                    function=MockFunction(
                        name="echo",
                        arguments="not valid json{{{",
                    ),
                )
            ]
        )
        mock_provider.responses = [bad_response, make_text_response("OK")]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "OK"
        error_msg = session.transcript[1]
        assert isinstance(error_msg, ToolCallResultMessage)
        assert "invalid arguments" in error_msg.content

    @pytest.mark.asyncio
    async def test_tool_returns_none_serializes_as_null(
        self, make_agent, mock_provider
    ):
        """A tool returning None serializes as 'null' in the transcript."""

        @tool
        def return_nothing():
            """Returns None."""
            return None

        mock_provider.responses = [
            make_tool_call_response("return_nothing", {}),
            make_text_response("Got null"),
        ]
        agent = make_agent(tools=[return_nothing])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Got null"
        tool_result_msg = session.transcript[1]
        assert tool_result_msg.content == "null"


# ---------------------------------------------------------------------------
# Multiple tool calls in a single response
# ---------------------------------------------------------------------------

class TestRunnerMultiToolCall:
    @pytest.mark.asyncio
    async def test_all_tool_calls_dispatched(
        self, make_agent, mock_provider
    ):
        """Multiple tool calls in one response all get executed."""
        mock_provider.responses = [
            make_multi_tool_call_response([
                ("echo", {"text": "first"}, "call_1"),
                ("echo", {"text": "second"}, "call_2"),
            ]),
            make_text_response("Done"),
        ]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Done"
        tool_results = [
            m for m in session.transcript
            if m.role == MessageRole.TOOL
        ]
        assert len(tool_results) == 2
        assert tool_results[0].content == "first"
        assert tool_results[1].content == "second"

    @pytest.mark.asyncio
    async def test_tool_call_ids_match_results(
        self, make_agent, mock_provider
    ):
        """Each tool result references the correct call_id."""
        mock_provider.responses = [
            make_multi_tool_call_response([
                ("echo", {"text": "a"}, "call_aaa"),
                ("echo", {"text": "b"}, "call_bbb"),
            ]),
            make_text_response("Done"),
        ]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        tool_results = [
            m for m in session.transcript
            if isinstance(m, ToolCallResultMessage)
        ]
        assert tool_results[0].tool_call_id == "call_aaa"
        assert tool_results[1].tool_call_id == "call_bbb"

    @pytest.mark.asyncio
    async def test_handoff_in_batch_stops_remaining_calls(
        self, make_agent, mock_provider
    ):
        """A handoff mid-batch prevents later tool calls from running."""
        call_log = []

        @tool
        def track_call(label: str):
            """Tracks that it was called."""
            call_log.append(label)
            return f"tracked: {label}"

        billing_handoff = Handoff(
            tool_name="transfer_to_billing",
            tool_description="Hand off to billing.",
            target_agent="billing",
        )

        mock_provider.responses = [
            make_multi_tool_call_response([
                (
                    "transfer_to_billing",
                    {"message": "handing off"},
                    "call_1",
                ),
                (
                    "track_call",
                    {"label": "should_not_run"},
                    "call_2",
                ),
            ]),
        ]
        agent = make_agent(tools=[track_call])
        session = Session(session_id="s1")

        result = await Runner().run(
            agent, session, State(),
            handoffs=[billing_handoff],
        )

        assert result.hand_off is not None
        assert result.hand_off.target_agent == "billing"
        assert "should_not_run" not in call_log


# ---------------------------------------------------------------------------
# Runner with capability tools (standalone, no Swarm)
# ---------------------------------------------------------------------------

class TestRunnerCapabilityTools:
    @pytest.mark.asyncio
    async def test_runner_dispatches_capability_tool(self, mock_provider):
        @tool
        def raven():
            """Recite the Raven."""
            return "Quoth the Raven, 'Nevermore.'"

        cap = MockCapability(name="raven", tool_list=[raven])
        agent = Agent(
            name="montresor",
            system_prompt="Once upon a midnight dreary.",
            model="mock-model", provider=mock_provider,
            capabilities=[cap],
        )
        mock_provider.responses = [
            make_tool_call_response("raven", {}),
            make_text_response("Darkness there and nothing more"),
        ]
        session = Session(session_id="s1")

        result = await Runner().run(agent, session, State())

        assert result.last_message.content == "Darkness there and nothing more"
        tool_result = [
            m for m in session.transcript
            if m.role == MessageRole.TOOL
        ]
        assert "Nevermore" in tool_result[0].content

    @pytest.mark.asyncio
    async def test_runner_sends_capability_tool_schema(self, mock_provider):
        @tool
        def raven():
            """Recite the Raven."""
            return "Quoth the Raven, 'Nevermore.'"

        cap = MockCapability(name="raven", tool_list=[raven])
        agent = Agent(
            name="montresor",
            system_prompt="Once upon a midnight dreary.",
            model="mock-model", provider=mock_provider,
            capabilities=[cap],
        )
        mock_provider.responses = [make_text_response("Darkness there and nothing more")]
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        sent_tools = mock_provider.call_log[0]["tools"]
        tool_names = [t["function"]["name"] for t in sent_tools]
        assert "raven" in tool_names


# ---------------------------------------------------------------------------
# End-to-end OpenTelemetry integration
# ---------------------------------------------------------------------------

class TestRunnerInstrumentation:
    """Verify that the runner creates the expected OTel spans when
    instrumentation is enabled.

    The mock tracer produces a distinct MagicMock span per
    ``start_as_current_span`` call, stored in ``self.spans``
    keyed by span name.  This lets tests assert attributes
    land on the *right* span.
    """

    @pytest.fixture(autouse=True)
    def _mock_tracer(self):
        """Install a mock tracer and tear it down after each test."""
        self.spans: dict[str, MagicMock] = {}
        self.mock_tracer = MagicMock()

        def _make_cm(name, **kwargs):
            span = MagicMock()
            self.spans[name] = span
            cm = MagicMock()
            cm.__enter__ = MagicMock(return_value=span)
            cm.__exit__ = MagicMock(return_value=False)
            return cm

        self.mock_tracer.start_as_current_span.side_effect = (
            _make_cm
        )
        inst._tracer = self.mock_tracer
        yield
        inst._tracer = None

    def _span_names(self) -> list[str]:
        """Return names of all spans opened by the mock tracer."""
        return [
            call.args[0]
            for call in self.mock_tracer.start_as_current_span.call_args_list
        ]

    @pytest.mark.asyncio
    async def test_simple_response_creates_agent_and_chat_spans(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [make_text_response("Hi")]
        agent = make_agent(name="helper")
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        names = self._span_names()
        assert names[0] == "invoke_agent helper"
        assert names[1] == "chat mock-model"
        assert len(names) == 2

    @pytest.mark.asyncio
    async def test_tool_call_creates_all_three_span_types(
        self, make_agent, mock_provider
    ):
        mock_provider.responses = [
            make_tool_call_response("echo", {"text": "hi"}),
            make_text_response("Done"),
        ]
        agent = make_agent(tools=[echo])
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        names = self._span_names()
        assert "invoke_agent test_agent" in names
        assert any(n.startswith("chat ") for n in names)
        assert "execute_tool echo" in names

    @pytest.mark.asyncio
    async def test_usage_recorded_on_completion_span(
        self, make_agent, mock_provider
    ):
        usage = MagicMock(prompt_tokens=10, completion_tokens=5)
        mock_provider.responses = [
            CompletionResult(
                content="Hi",
                usage=usage,
                response_model="gpt-4o-2024-08-06",
            )
        ]
        agent = make_agent()
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        # Attributes land on the completion span, not the agent span
        c_span = self.spans["chat mock-model"]
        c_span.set_attribute.assert_any_call(
            "gen_ai.usage.input_tokens", 10
        )
        c_span.set_attribute.assert_any_call(
            "gen_ai.usage.output_tokens", 5
        )
        c_span.set_attribute.assert_any_call(
            "gen_ai.response.model", "gpt-4o-2024-08-06"
        )
        # Agent span should NOT have usage attributes
        a_span = self.spans["invoke_agent test_agent"]
        a_span.set_attribute.assert_not_called()

    @pytest.mark.asyncio
    async def test_tool_error_recorded_on_span(
        self, make_agent, mock_provider
    ):
        """When a tool raises, record_error marks the tool span."""

        @tool
        def explode():
            """Always fails."""
            raise RuntimeError("kaboom")

        mock_provider.responses = [
            make_tool_call_response("explode", {}),
            make_text_response("Recovered"),
        ]
        agent = make_agent(tools=[explode])
        session = Session(session_id="s1")

        await Runner().run(agent, session, State())

        t_span = self.spans["execute_tool explode"]
        t_span.set_status.assert_called_once()
        t_span.record_exception.assert_called_once()
        t_span.set_attribute.assert_any_call(
            "error.type", "RuntimeError"
        )
