# standard
import json

# third party
import pytest

# custom
from sunwaee.modules.gen.agent import run
from sunwaee.modules.gen.engine.base import BaseEngine
from sunwaee.modules.gen.engine.types import (
    Message,
    Response,
    Role,
    StopReason,
    Tool,
    ToolCall,
    Usage,
)

### MOCK


class MockEngine(BaseEngine):
    """Returns pre-configured Response objects in sequence."""

    def __init__(self, responses: list[Response]):
        self._responses = list(responses)
        self.calls: list[tuple[list[Message], list[Tool] | None]] = []

    async def chat(
        self, messages: list[Message], tools: list[Tool] | None = None
    ) -> Response:
        self.calls.append((list(messages), tools))
        return self._responses.pop(0)

    def stream(self, messages, tools=None):
        async def _empty():
            return
            yield  # pragma: no cover

        return _empty()


def _end(
    content: str, reasoning: str | None = None, signature: str | None = None
) -> Response:
    return Response(
        content=content,
        tool_calls=None,
        usage=Usage(input_tokens=5, output_tokens=10, total_tokens=15),
        stop_reason=StopReason.END_TURN,
        reasoning_content=reasoning,
        reasoning_signature=signature,
    )


def _tool_call(name: str, arguments: dict, call_id: str = "call-1") -> Response:
    return Response(
        content=None,
        tool_calls=[ToolCall(id=call_id, name=name, arguments=arguments)],
        usage=Usage(input_tokens=5, output_tokens=10, total_tokens=15),
        stop_reason=StopReason.TOOL_USE,
    )


def _max_tokens(content: str) -> Response:
    return Response(
        content=content,
        tool_calls=None,
        usage=Usage(input_tokens=5, output_tokens=10, total_tokens=15),
        stop_reason=StopReason.MAX_TOKENS,
    )


### DUMMY


def echo_tool(message: str = "pong") -> str:
    return json.dumps({"ok": True, "data": message})


ECHO = Tool(name="echo", description="Echo a message.", parameters={}, fn=echo_tool)


### TESTS


@pytest.mark.asyncio
async def test_simple_response_no_tools():
    engine = MockEngine([_end("Hello, user!")])
    messages = [Message(role=Role.USER, content="Hi")]

    text, new_msgs = await run(messages, [], engine)

    assert text == "Hello, user!"
    assert len(engine.calls) == 1
    # new_msgs = [assistant message]
    assert len(new_msgs) == 1
    assert new_msgs[0].role == Role.ASSISTANT
    assert new_msgs[0].content == "Hello, user!"


@pytest.mark.asyncio
async def test_reasoning_stored_in_message():
    engine = MockEngine(
        [_end("Answer", reasoning="my chain of thought", signature="sig-abc")]
    )
    text, new_msgs = await run([Message(role=Role.USER, content="Q")], [], engine)

    assistant_msg = new_msgs[0]
    assert assistant_msg.reasoning_content == "my chain of thought"
    assert assistant_msg.reasoning_signature == "sig-abc"


@pytest.mark.asyncio
async def test_single_tool_call():
    engine = MockEngine(
        [
            _tool_call("echo", {"message": "hello"}),
            _end("Done!"),
        ]
    )
    text, new_msgs = await run(
        [Message(role=Role.USER, content="Call echo")], [ECHO], engine
    )

    assert text == "Done!"
    assert len(engine.calls) == 2

    roles = [m.role.value for m in new_msgs]
    assert roles == ["assistant", "tool", "assistant"]

    tool_result = new_msgs[1]
    assert tool_result.role == Role.TOOL
    assert tool_result.tool_call_id == "call-1"
    assert tool_result.content
    result_data = json.loads(tool_result.content)
    assert result_data["data"] == "hello"


@pytest.mark.asyncio
async def test_multiple_concurrent_tool_calls():
    second_tool = Tool(
        name="second",
        description="Second tool.",
        parameters={},
        fn=lambda: json.dumps({"ok": True, "data": "second result"}),
    )
    engine = MockEngine(
        [
            Response(
                content=None,
                tool_calls=[
                    ToolCall(id="c1", name="echo", arguments={"message": "first"}),
                    ToolCall(id="c2", name="second", arguments={}),
                ],
                usage=Usage(input_tokens=5, output_tokens=10, total_tokens=15),
                stop_reason=StopReason.TOOL_USE,
            ),
            _end("All done"),
        ]
    )

    text, new_msgs = await run(
        [Message(role=Role.USER, content="Call both")],
        [ECHO, second_tool],
        engine,
    )

    assert text == "All done"
    tool_results = [m for m in new_msgs if m.role.value == "tool"]
    assert len(tool_results) == 2
    ids = {m.tool_call_id for m in tool_results}
    assert ids == {"c1", "c2"}


@pytest.mark.asyncio
async def test_unknown_tool_returns_error():
    engine = MockEngine(
        [
            _tool_call("nonexistent", {}),
            _end("Sorry"),
        ]
    )
    _, new_msgs = await run([Message(role=Role.USER, content="Hi")], [ECHO], engine)

    tool_result = next(m for m in new_msgs if m.role.value == "tool")
    assert tool_result.content
    data = json.loads(tool_result.content)
    assert data["ok"] is False
    assert "nonexistent" in data["error"]


@pytest.mark.asyncio
async def test_max_iterations_returns_last_content():
    # NOTE only tool_use responses — never returns end_turn within max_iterations
    engine = MockEngine([_tool_call("echo", {"message": "x"})] * 3)
    # patch the echo tool to avoid real execution errors being the issue
    text, new_msgs = await run(
        [Message(role=Role.USER, content="Hi")],
        [ECHO],
        engine,
        max_iterations=2,
    )
    # should return the warning message since no end_turn happened
    assert "Max iterations" in text or isinstance(text, str)


@pytest.mark.asyncio
async def test_tool_exception_returns_error_message():
    def broken_tool() -> str:
        raise ValueError("Something went wrong")

    bad_tool = Tool(name="broken", description="Breaks.", parameters={}, fn=broken_tool)
    engine = MockEngine(
        [
            _tool_call("broken", {}),
            _end("Recovered"),
        ]
    )

    _, new_msgs = await run([Message(role=Role.USER, content="Hi")], [bad_tool], engine)
    tool_result = next(m for m in new_msgs if m.role.value == "tool")
    assert tool_result.content
    data = json.loads(tool_result.content)
    assert data["ok"] is False
    assert "Something went wrong" in data["error"]


@pytest.mark.asyncio
async def test_async_tool_is_awaited():
    async def async_echo(message: str = "async") -> str:
        return json.dumps({"ok": True, "data": f"async:{message}"})

    async_tool = Tool(
        name="async_echo", description="Async echo.", parameters={}, fn=async_echo
    )
    engine = MockEngine(
        [
            _tool_call("async_echo", {"message": "hi"}),
            _end("Done"),
        ]
    )

    _, new_msgs = await run(
        [Message(role=Role.USER, content="Go")], [async_tool], engine
    )
    tool_result = next(m for m in new_msgs if m.role.value == "tool")
    assert tool_result.content
    data = json.loads(tool_result.content)
    assert data["data"] == "async:hi"


@pytest.mark.asyncio
async def test_toolcall_results_populated():
    engine = MockEngine(
        [
            _tool_call("echo", {"message": "hello"}),
            _end("Done"),
        ]
    )
    _, new_msgs = await run([Message(role=Role.USER, content="Go")], [ECHO], engine)

    assistant_msg = new_msgs[0]
    assert assistant_msg.tool_calls is not None
    tc = assistant_msg.tool_calls[0]
    assert tc.results == [{"ok": True, "data": "hello"}]
    assert tc.error is None
    assert tc.duration >= 0.0


@pytest.mark.asyncio
async def test_toolcall_error_populated_on_unknown_tool():
    engine = MockEngine([_tool_call("ghost", {}), _end("Ok")])
    _, new_msgs = await run([Message(role=Role.USER, content="Hi")], [ECHO], engine)

    assert new_msgs[0].tool_calls is not None
    tc = new_msgs[0].tool_calls[0]
    assert tc.error is not None
    assert "ghost" in tc.error
    assert tc.results == []
    assert tc.duration == 0.0


@pytest.mark.asyncio
async def test_toolcall_error_populated_on_exception():
    def broken() -> str:
        raise RuntimeError("boom")

    bad = Tool(name="bad", description=".", parameters={}, fn=broken)
    engine = MockEngine([_tool_call("bad", {}), _end("Ok")])
    _, new_msgs = await run([Message(role=Role.USER, content="Hi")], [bad], engine)

    assert new_msgs[0].tool_calls is not None
    tc = new_msgs[0].tool_calls[0]
    assert tc.error == "boom"
    assert tc.duration >= 0.0


@pytest.mark.asyncio
async def test_history_not_mutated():
    """Input messages list must not be modified by the agent."""
    engine = MockEngine([_end("Hello")])
    original = [Message(role=Role.USER, content="Hi")]
    original_len = len(original)

    await run(original, [], engine)
    assert len(original) == original_len


@pytest.mark.asyncio
async def test_new_messages_excludes_input():
    """new_messages should only contain messages added during THIS run."""
    engine = MockEngine([_end("Response")])
    history = [
        Message(role=Role.USER, content="Old message"),
        Message(role=Role.ASSISTANT, content="Old response"),
        Message(role=Role.USER, content="New question"),
    ]

    _, new_msgs = await run(history, [], engine)
    # only the assistant response from this run
    assert len(new_msgs) == 1
    assert new_msgs[0].content == "Response"
