# standard
from unittest.mock import AsyncMock, MagicMock, patch

# third party
import pytest

# custom
from sunwaee.modules.gen.engine.providers.google import GoogleEngine
from sunwaee.modules.gen.engine.types import Message, Role, StopReason, Tool, ToolCall

ENGINE = GoogleEngine(model="gemini-2.0-flash", api_key="test-key")

DUMMY_TOOL = Tool(
    name="tasks",
    description="Manage tasks.",
    parameters={"type": "object", "properties": {}, "required": []},
)


### CONTENTS


def test_build_contents_extracts_system():
    msgs = [
        Message(role=Role.SYSTEM, content="You are Sun."),
        Message(role=Role.USER, content="Hi"),
    ]
    system, contents = ENGINE._build_contents(msgs)
    assert system == "You are Sun."
    assert len(contents) == 1
    assert contents[0]["role"] == "user"


def test_build_contents_system_only_promotes_to_user():
    system, contents = ENGINE._build_contents(
        [Message(role=Role.SYSTEM, content="Respond with ok.")]
    )
    assert system is None
    assert len(contents) == 1
    assert contents[0]["role"] == "user"
    assert contents[0]["parts"][0]["text"] == "Respond with ok."


def test_build_contents_user_and_assistant():
    msgs = [
        Message(role=Role.USER, content="Hello"),
        Message(role=Role.ASSISTANT, content="Hi there"),
    ]
    _, contents = ENGINE._build_contents(msgs)
    assert contents[0] == {"role": "user", "parts": [{"text": "Hello"}]}
    assert contents[1]["role"] == "model"
    assert contents[1]["parts"][0] == {"text": "Hi there"}


def test_build_contents_tool_result():
    msgs = [Message(role=Role.TOOL, content='{"ok": true}', tool_call_id="my_tool")]
    _, contents = ENGINE._build_contents(msgs)
    assert contents[0]["role"] == "user"
    assert contents[0]["parts"][0]["functionResponse"]["name"] == "my_tool"


def test_build_contents_assistant_with_tool_calls():
    tc = ToolCall(id="tasks", name="tasks", arguments={"action": "list"})
    msgs = [Message(role=Role.ASSISTANT, tool_calls=[tc])]
    _, contents = ENGINE._build_contents(msgs)
    parts = contents[0]["parts"]
    assert parts[0]["functionCall"]["name"] == "tasks"
    assert parts[0]["functionCall"]["args"] == {"action": "list"}


def test_build_contents_reasoning_in_assistant():
    """reasoning_content is not echoed in requests; only thoughtSignature is."""
    msgs = [
        Message(role=Role.USER, content="Question"),
        Message(
            role=Role.ASSISTANT,
            content="Answer",
            reasoning_content="I think...",
            reasoning_signature="sig",
        ),
    ]
    _, contents = ENGINE._build_contents(msgs)
    model_parts = contents[1]["parts"]
    # only the text part — no separate thought block
    assert model_parts[0] == {"text": "Answer", "thoughtSignature": "sig"}


def test_build_contents_thought_signature_only_on_last_turn():
    """thoughtSignature only on the last assistant turn (safe for cross-engine history)."""
    msgs = [
        Message(role=Role.USER, content="Q1"),
        Message(
            role=Role.ASSISTANT,
            content="A1",
            reasoning_content="t1",
            reasoning_signature="s1",
        ),
        Message(role=Role.USER, content="Q2"),
        Message(
            role=Role.ASSISTANT,
            content="A2",
            reasoning_content="t2",
            reasoning_signature="s2",
        ),
    ]
    _, contents = ENGINE._build_contents(msgs)

    # first assistant turn: no signature
    assert "thoughtSignature" not in contents[1]["parts"][0]
    # last assistant turn: signature in the text part
    assert contents[3]["parts"][0].get("thoughtSignature") == "s2"


def test_build_contents_thought_signature_on_last_tool_call_turn():
    """Signature included when the last assistant turn has tool calls."""
    tc = ToolCall(id="update_chat", name="update_chat", arguments={"title": "Test"})
    msgs = [
        Message(role=Role.USER, content="Update chat"),
        Message(
            role=Role.ASSISTANT,
            content="",
            reasoning_signature="sig-abc",
            tool_calls=[tc],
        ),
    ]
    _, contents = ENGINE._build_contents(msgs)

    part = contents[1]["parts"][0]
    assert part["functionCall"]["name"] == "update_chat"
    assert part["thoughtSignature"] == "sig-abc"


def test_build_contents_thought_signature_not_on_intermediate_tool_call_turn():
    """Signature NOT included on an intermediate tool-call turn (not the last assistant turn)."""
    tc = ToolCall(id="fn", name="fn", arguments={})
    msgs = [
        Message(role=Role.USER, content="Go"),
        Message(
            role=Role.ASSISTANT,
            content="",
            reasoning_signature="sig-intermediate",
            tool_calls=[tc],
        ),
        Message(role=Role.TOOL, content='{"ok": true}', tool_call_id="fn"),
        Message(role=Role.ASSISTANT, content="Done", reasoning_signature="sig-last"),
    ]
    _, contents = ENGINE._build_contents(msgs)

    # intermediate assistant (tool call): no signature
    assert "thoughtSignature" not in contents[1]["parts"][0]
    # last assistant turn: signature present
    assert contents[3]["parts"][0].get("thoughtSignature") == "sig-last"


def test_build_contents_no_signature_without_reasoning():
    msgs = [
        Message(role=Role.USER, content="Q"),
        Message(role=Role.ASSISTANT, content="A"),
    ]
    _, contents = ENGINE._build_contents(msgs)
    parts = contents[1]["parts"]
    assert all("thoughtSignature" not in p for p in parts)


### TOOLS


def test_build_tools():
    tools = ENGINE._build_tools([DUMMY_TOOL])
    assert tools == [
        {
            "functionDeclarations": [
                {
                    "name": "tasks",
                    "description": "Manage tasks.",
                    "parameters": {"type": "object", "properties": {}, "required": []},
                }
            ]
        }
    ]


def test_build_payload_with_system():
    msgs = [
        Message(role=Role.SYSTEM, content="You are Sun."),
        Message(role=Role.USER, content="Hi"),
    ]
    payload = ENGINE._build_payload(msgs, tools=None)
    assert payload["systemInstruction"] == {"parts": [{"text": "You are Sun."}]}


def test_build_payload_with_tools():
    msgs = [Message(role=Role.USER, content="Hi")]
    payload = ENGINE._build_payload(msgs, tools=[DUMMY_TOOL])
    assert "tools" in payload
    assert payload["tools"][0]["functionDeclarations"][0]["name"] == "tasks"


def test_build_payload_thinking_enabled_by_default():
    msgs = [Message(role=Role.USER, content="Hi")]
    payload = ENGINE._build_payload(msgs, tools=None)
    assert payload["generationConfig"]["thinkingConfig"]["includeThoughts"] is True
    assert "thinkingBudget" not in payload["generationConfig"]["thinkingConfig"]


def test_build_payload_thinking_budget_custom():
    engine = GoogleEngine(model="gemini-2.0-flash", api_key="k", thinking_budget=5000)
    payload = engine._build_payload([Message(role=Role.USER, content="Hi")], tools=None)
    cfg = payload["generationConfig"]["thinkingConfig"]
    assert cfg["includeThoughts"] is True
    assert cfg["thinkingBudget"] == 5000


def test_build_payload_thinking_disabled():
    engine = GoogleEngine(model="gemini-2.0-flash", api_key="k", thinking_budget=0)
    payload = engine._build_payload([Message(role=Role.USER, content="Hi")], tools=None)
    cfg = payload["generationConfig"]["thinkingConfig"]
    assert cfg["includeThoughts"] is True
    assert "thinkingBudget" not in cfg


### RESPONSE PARSING


def _raw(parts, finish_reason="STOP", prompt_tokens=10, candidate_tokens=20):
    return {
        "candidates": [
            {
                "content": {"parts": parts},
                "finishReason": finish_reason,
            }
        ],
        "usageMetadata": {
            "promptTokenCount": prompt_tokens,
            "candidatesTokenCount": candidate_tokens,
        },
    }


def test_parse_response_text():
    r = ENGINE._parse_response(_raw([{"text": "Hello!"}]))
    assert r.content == "Hello!"
    assert r.tool_calls is None
    assert r.stop_reason == StopReason.END_TURN
    assert r.usage
    assert r.usage.input_tokens == 10
    assert r.usage.output_tokens == 20
    assert r.reasoning_content is None
    assert r.provider == "google"
    assert r.model == ENGINE.model


def test_parse_response_thought_part():
    r = ENGINE._parse_response(
        _raw(
            [
                {"thought": True, "text": "I think...", "thoughtSignature": "sig-xyz"},
                {"text": "My answer"},
            ]
        )
    )
    assert r.reasoning_content == "I think..."
    assert r.reasoning_signature == "sig-xyz"
    assert r.content == "My answer"


def test_parse_response_function_call():
    r = ENGINE._parse_response(
        _raw(
            [
                {"functionCall": {"name": "tasks", "args": {"action": "list"}}},
            ]
        )
    )
    assert r.stop_reason == StopReason.TOOL_USE
    assert r.tool_calls
    assert len(r.tool_calls) == 1
    # NOTE gemini uses function name as id
    assert r.tool_calls[0].id == "tasks"
    assert r.tool_calls[0].name == "tasks"
    assert r.tool_calls[0].arguments == {"action": "list"}


def test_parse_response_thought_and_tool_call():
    r = ENGINE._parse_response(
        _raw(
            [
                {"thought": True, "text": "Should call tool", "thoughtSignature": "s"},
                {"functionCall": {"name": "notes", "args": {"action": "list"}}},
            ]
        )
    )
    assert r.reasoning_content == "Should call tool"
    assert r.tool_calls
    assert r.tool_calls[0].name == "notes"
    assert r.stop_reason == StopReason.TOOL_USE


def test_parse_response_max_tokens():
    r = ENGINE._parse_response(_raw([{"text": "partial"}], finish_reason="MAX_TOKENS"))
    assert r.stop_reason == StopReason.MAX_TOKENS


def test_parse_response_multiple_function_calls():
    r = ENGINE._parse_response(
        _raw(
            [
                {"functionCall": {"name": "tasks", "args": {}}},
                {"functionCall": {"name": "notes", "args": {}}},
            ]
        )
    )
    assert r.tool_calls
    assert len(r.tool_calls) == 2


### CHAT - chat()


@pytest.mark.asyncio
async def test_chat_returns_response():
    raw = _raw([{"text": "Hi from Gemini"}])
    mock_resp = MagicMock()
    mock_resp.is_success = True
    mock_resp.json.return_value = raw

    with patch(
        "httpx.AsyncClient.post", new_callable=AsyncMock, return_value=mock_resp
    ):
        response = await ENGINE.chat([Message(role=Role.USER, content="Hello")])

    assert response.content == "Hi from Gemini"


@pytest.mark.asyncio
async def test_chat_raises_on_error():
    mock_resp = MagicMock()
    mock_resp.is_success = False
    mock_resp.status_code = 403
    mock_resp.text = "Forbidden"

    with patch(
        "httpx.AsyncClient.post", new_callable=AsyncMock, return_value=mock_resp
    ):
        with pytest.raises(RuntimeError, match="403"):
            await ENGINE.chat([Message(role=Role.USER, content="Hello")])


# STREAM - stream()


@pytest.mark.asyncio
async def test_stream_yields_text_parts():
    lines = [
        'data: {"candidates": [{"content": {"parts": [{"text": "Hello"}]}, "finishReason": null}], "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1}}',
        'data: {"candidates": [{"content": {"parts": [{"text": " Gemini"}]}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2}}',
    ]

    mock_resp = AsyncMock()
    mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
    mock_resp.__aexit__ = AsyncMock(return_value=False)
    mock_resp.raise_for_status = MagicMock()
    mock_resp.aiter_lines = MagicMock(return_value=aiter(lines))

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)
    mock_client.stream = MagicMock(return_value=mock_resp)

    with patch("httpx.AsyncClient", return_value=mock_client):
        chunks = [
            c async for c in ENGINE.stream([Message(role=Role.USER, content="Hi")])
        ]

    # 2 content chunks + 1 final summary chunk (tokens + stop_reason + performance)
    assert len(chunks) == 3
    assert chunks[0].content == "Hello"
    assert chunks[0].stop_reason is None
    assert chunks[0].usage is None
    assert chunks[1].content == " Gemini"
    assert chunks[1].stop_reason is None
    assert chunks[1].usage is None
    # final summary
    assert chunks[2].content is None
    assert chunks[2].stop_reason == StopReason.END_TURN
    assert chunks[2].usage
    assert chunks[2].usage.input_tokens == 3
    assert chunks[2].usage.output_tokens == 2
    assert chunks[2].streaming is True
    assert chunks[2].performance is not None


@pytest.mark.asyncio
async def test_stream_yields_reasoning_and_signature():
    lines = [
        'data: {"candidates": [{"content": {"parts": [{"thought": true, "text": "I think...", "thoughtSignature": "sig-abc"}]}, "finishReason": null}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 0}}',
        'data: {"candidates": [{"content": {"parts": [{"text": "My answer"}]}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 3}}',
    ]

    mock_resp = AsyncMock()
    mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
    mock_resp.__aexit__ = AsyncMock(return_value=False)
    mock_resp.raise_for_status = MagicMock()
    mock_resp.aiter_lines = MagicMock(return_value=aiter(lines))

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)
    mock_client.stream = MagicMock(return_value=mock_resp)

    with patch("httpx.AsyncClient", return_value=mock_client):
        chunks = [
            c async for c in ENGINE.stream([Message(role=Role.USER, content="Hi")])
        ]

    # reasoning_content chunk + reasoning_signature chunk + text chunk + summary
    assert len(chunks) == 4
    assert chunks[0].reasoning_content == "I think..."
    assert chunks[0].content is None
    assert chunks[1].reasoning_signature == "sig-abc"
    assert chunks[1].content is None
    assert chunks[2].content == "My answer"
    assert chunks[3].stop_reason == StopReason.END_TURN
    assert chunks[3].usage
    assert chunks[3].usage.input_tokens == 5
    assert chunks[3].usage.output_tokens == 3


@pytest.mark.asyncio
async def test_stream_yields_tool_call():
    lines = [
        'data: {"candidates": [{"content": {"parts": [{"functionCall": {"name": "tasks", "args": {"action": "list"}}}]}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 2}}',
    ]

    mock_resp = AsyncMock()
    mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
    mock_resp.__aexit__ = AsyncMock(return_value=False)
    mock_resp.raise_for_status = MagicMock()
    mock_resp.aiter_lines = MagicMock(return_value=aiter(lines))

    mock_client = AsyncMock()
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)
    mock_client.stream = MagicMock(return_value=mock_resp)

    with patch("httpx.AsyncClient", return_value=mock_client):
        chunks = [
            c async for c in ENGINE.stream([Message(role=Role.USER, content="Hi")])
        ]

    assert len(chunks) == 2
    assert chunks[0].tool_calls is not None
    assert chunks[0].tool_calls[0].name == "tasks"
    assert chunks[0].tool_calls[0].id == "tasks"
    assert chunks[0].tool_calls[0].arguments == {"action": "list"}
    assert chunks[1].stop_reason == StopReason.TOOL_USE
    assert chunks[1].usage
    assert chunks[1].usage.input_tokens == 4


### HELPERS


async def aiter(items):
    for item in items:
        yield item
