# standard
import json

# third party
import pytest

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

### MOCK


class MockStreamEngine(BaseEngine):
    """Yields pre-configured Response chunk sequences, one list per stream() call."""

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

    async def chat(self, messages, tools=None) -> Response:
        raise NotImplementedError

    async def stream(self, messages, tools=None):
        self.calls.append((list(messages), tools))
        for chunk in self._sequences.pop(0):
            yield chunk


### HELPERS


def _final(
    stop_reason: StopReason = StopReason.END_TURN,
    tool_calls: list[ToolCall] | None = None,
) -> Response:
    """Final chunk — carries stop_reason, usage, performance."""
    return Response(
        streaming=True,
        stop_reason=stop_reason,
        tool_calls=tool_calls,
        usage=Usage(input_tokens=5, output_tokens=10, total_tokens=15),
        performance=Performance(latency=0.1, total_duration=0.5),
    )


def _text(t: str) -> Response:
    return Response(streaming=True, content=t)


def _reasoning(t: str) -> Response:
    return Response(streaming=True, reasoning_content=t)


def _sentinel(t: str) -> Response:
    return Response(streaming=True, synthetic=True, reasoning_content=t)


def _signature(sig: str) -> Response:
    return Response(streaming=True, reasoning_signature=sig)


def _tool_chunk(tc: ToolCall) -> Response:
    return Response(streaming=True, tool_calls=[tc])


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


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


### TESTS


@pytest.mark.asyncio
async def test_simple_stream_no_tools():
    engine = MockStreamEngine([[_text("Hello"), _text(", world!"), _final()]])
    chunks = []
    async for chunk in stream_run([Message(role=Role.USER, content="Hi")], [], engine):
        chunks.append(chunk)

    assert len(engine.calls) == 1
    text_chunks = [c for c in chunks if c.content]
    assert "".join(c.content for c in text_chunks) == "Hello, world!"


@pytest.mark.asyncio
async def test_new_messages_populated():
    engine = MockStreamEngine([[_text("Answer"), _final()]])
    new_msgs: list[Message] = []
    async for _ in stream_run(
        [Message(role=Role.USER, content="Q")], [], engine, new_messages=new_msgs
    ):
        pass

    assert len(new_msgs) == 1
    assert new_msgs[0].role == Role.ASSISTANT
    assert new_msgs[0].content == "Answer"


@pytest.mark.asyncio
async def test_single_tool_call():
    tc = ToolCall(id="c1", name="echo", arguments={"message": "hi"})
    engine = MockStreamEngine(
        [
            [_tool_chunk(tc), _final(StopReason.TOOL_USE)],
            [_text("Done!"), _final()],
        ]
    )
    new_msgs: list[Message] = []
    chunks = []
    async for chunk in stream_run(
        [Message(role=Role.USER, content="Call echo")],
        [ECHO],
        engine,
        new_messages=new_msgs,
    ):
        chunks.append(chunk)

    assert len(engine.calls) == 2

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

    tool_msg = next(m for m in new_msgs if m.role == Role.TOOL)
    assert tool_msg.tool_call_id == "c1"
    assert tool_msg.content is not None
    assert json.loads(tool_msg.content)["data"] == "hi"

    text_chunks = [c for c in chunks if c.content]
    assert "".join(c.content for c in text_chunks) == "Done!"


@pytest.mark.asyncio
async def test_multiple_tool_iterations():
    tc1 = ToolCall(id="c1", name="echo", arguments={"message": "first"})
    tc2 = ToolCall(id="c2", name="echo", arguments={"message": "second"})
    engine = MockStreamEngine(
        [
            [_tool_chunk(tc1), _final(StopReason.TOOL_USE)],
            [_tool_chunk(tc2), _final(StopReason.TOOL_USE)],
            [_text("All done"), _final()],
        ]
    )
    new_msgs: list[Message] = []
    async for _ in stream_run(
        [Message(role=Role.USER, content="Go")],
        [ECHO],
        engine,
        new_messages=new_msgs,
    ):
        pass

    assert len(engine.calls) == 3
    tool_msgs = [m for m in new_msgs if m.role == Role.TOOL]
    assert len(tool_msgs) == 2


@pytest.mark.asyncio
async def test_concurrent_tool_calls():
    second = Tool(
        name="second",
        description=".",
        parameters={},
        fn=lambda: json.dumps({"ok": True, "data": "s"}),
    )
    tc1 = ToolCall(id="c1", name="echo", arguments={"message": "a"})
    tc2 = ToolCall(id="c2", name="second", arguments={})
    engine = MockStreamEngine(
        [
            [_tool_chunk(tc1), _tool_chunk(tc2), _final(StopReason.TOOL_USE)],
            [_text("Done"), _final()],
        ]
    )
    new_msgs: list[Message] = []
    async for _ in stream_run(
        [Message(role=Role.USER, content="Both")],
        [ECHO, second],
        engine,
        new_messages=new_msgs,
    ):
        pass

    tool_msgs = [m for m in new_msgs if m.role == Role.TOOL]
    assert len(tool_msgs) == 2
    assert {m.tool_call_id for m in tool_msgs} == {"c1", "c2"}


@pytest.mark.asyncio
async def test_reasoning_chunks_forwarded():
    engine = MockStreamEngine(
        [[_reasoning("think..."), _signature("sig-1"), _text("Answer"), _final()]]
    )
    chunks = []
    async for chunk in stream_run([Message(role=Role.USER, content="Q")], [], engine):
        chunks.append(chunk)

    reasoning = "".join(c.reasoning_content for c in chunks if c.reasoning_content)
    assert reasoning == "think..."
    sig_chunks = [c for c in chunks if c.reasoning_signature]
    assert sig_chunks[-1].reasoning_signature == "sig-1"


@pytest.mark.asyncio
async def test_reasoning_stored_in_new_messages():
    engine = MockStreamEngine(
        [[_reasoning("chain"), _signature("sig-x"), _text("Result"), _final()]]
    )
    new_msgs: list[Message] = []
    async for _ in stream_run(
        [Message(role=Role.USER, content="Q")], [], engine, new_messages=new_msgs
    ):
        pass

    assistant = new_msgs[0]
    assert assistant.reasoning_content == "chain"
    assert assistant.reasoning_signature == "sig-x"


@pytest.mark.asyncio
async def test_unknown_tool_returns_error():
    tc = ToolCall(id="c1", name="ghost", arguments={})
    engine = MockStreamEngine(
        [
            [_tool_chunk(tc), _final(StopReason.TOOL_USE)],
            [_text("Ok"), _final()],
        ]
    )
    new_msgs: list[Message] = []
    async for _ in stream_run(
        [Message(role=Role.USER, content="Hi")], [ECHO], engine, new_messages=new_msgs
    ):
        pass

    tool_msg = next(m for m in new_msgs if m.role == Role.TOOL)
    assert tool_msg.content is not None
    data = json.loads(tool_msg.content)
    assert data["ok"] is False
    assert "ghost" in data["error"]


@pytest.mark.asyncio
async def test_max_iterations():
    tc = ToolCall(id="c1", name="echo", arguments={"message": "x"})
    engine = MockStreamEngine([[_tool_chunk(tc), _final(StopReason.TOOL_USE)]] * 3)
    chunks = []
    async for chunk in stream_run(
        [Message(role=Role.USER, content="Hi")],
        [ECHO],
        engine,
        max_iterations=2,
    ):
        chunks.append(chunk)

    assert len(engine.calls) == 2


@pytest.mark.asyncio
async def test_history_not_mutated():
    engine = MockStreamEngine([[_text("Hi"), _final()]])
    original = [Message(role=Role.USER, content="Hello")]
    original_len = len(original)

    async for _ in stream_run(original, [], engine):
        pass

    assert len(original) == original_len


@pytest.mark.asyncio
async def test_new_messages_excludes_input():
    engine = MockStreamEngine([[_text("Response"), _final()]])
    history = [
        Message(role=Role.USER, content="Old"),
        Message(role=Role.ASSISTANT, content="Old reply"),
        Message(role=Role.USER, content="New"),
    ]
    new_msgs: list[Message] = []
    async for _ in stream_run(history, [], engine, new_messages=new_msgs):
        pass

    assert len(new_msgs) == 1
    assert new_msgs[0].content == "Response"


@pytest.mark.asyncio
async def test_no_new_messages_arg_works():
    """stream_run works fine without passing new_messages."""
    engine = MockStreamEngine([[_text("Ok"), _final()]])
    chunks = []
    async for chunk in stream_run([Message(role=Role.USER, content="Hi")], [], engine):
        chunks.append(chunk)
    assert any(c.content for c in chunks)


@pytest.mark.asyncio
async def test_synthetic_sentinel_yielded_but_not_stored():
    """Synthetic sentinel chunks are forwarded to the caller but not accumulated
    into the assistant message, so they are never sent back to the provider API."""
    engine = MockStreamEngine(
        [
            [
                _sentinel(
                    "Reasoning in progress - tokens not available for this model"
                ),
                _text("Answer"),
                _final(),
            ]
        ]
    )
    new_msgs: list[Message] = []
    chunks = []
    async for chunk in stream_run(
        [Message(role=Role.USER, content="Q")], [], engine, new_messages=new_msgs
    ):
        chunks.append(chunk)

    # sentinel was yielded to the caller
    sentinel_chunks = [c for c in chunks if c.synthetic]
    assert len(sentinel_chunks) == 1

    # but not stored in the assistant message history
    assistant = new_msgs[0]
    assert assistant.reasoning_content is None
