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

# third party
import pytest

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

ENGINE = OpenAIEngine(model="gpt-4o", api_key="test-key")
DEEPSEEK = OpenAIEngine(
    model="deepseek-reasoner",
    api_key="test-key",
    base_url="https://api.deepseek.com/v1",
)

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


### MESSAGES


def test_build_messages_user():
    result = ENGINE._build_messages([Message(role=Role.USER, content="Hi")])
    assert result == [{"role": "user", "content": "Hi"}]


def test_build_messages_system():
    result = ENGINE._build_messages([Message(role=Role.SYSTEM, content="You are Sun.")])
    assert result == [{"role": "system", "content": "You are Sun."}]


def test_build_messages_assistant_text():
    result = ENGINE._build_messages([Message(role=Role.ASSISTANT, content="Hello")])
    assert result == [{"role": "assistant", "content": "Hello"}]


def test_build_messages_assistant_with_tool_calls():
    tc = ToolCall(id="call-1", name="tasks", arguments={"action": "list"})
    result = ENGINE._build_messages(
        [Message(role=Role.ASSISTANT, content=None, tool_calls=[tc])]
    )
    msg = result[0]
    assert msg["role"] == "assistant"
    assert msg["tool_calls"][0]["id"] == "call-1"
    assert msg["tool_calls"][0]["type"] == "function"
    assert msg["tool_calls"][0]["function"]["name"] == "tasks"
    assert json.loads(msg["tool_calls"][0]["function"]["arguments"]) == {
        "action": "list"
    }


def test_build_messages_tool_result():
    result = ENGINE._build_messages(
        [Message(role=Role.TOOL, content='{"ok": true}', tool_call_id="call-1")]
    )
    assert result == [
        {"role": "tool", "tool_call_id": "call-1", "content": '{"ok": true}'}
    ]


def test_build_messages_reasoning_content_echoed_on_last_turn():
    result = ENGINE._build_messages(
        [
            Message(
                role=Role.ASSISTANT, content="Hi", reasoning_content="internal thought"
            )
        ]
    )
    assert result[0]["reasoning_content"] == "internal thought"


def test_build_messages_reasoning_content_not_echoed_on_intermediate_turns():
    msgs = [
        Message(role=Role.USER, content="Q1"),
        Message(role=Role.ASSISTANT, content="A1", reasoning_content="thought1"),
        Message(role=Role.USER, content="Q2"),
        Message(role=Role.ASSISTANT, content="A2", reasoning_content="thought2"),
    ]
    result = ENGINE._build_messages(msgs)
    assistant_msgs = [m for m in result if m.get("role") == "assistant"]
    assert "reasoning_content" not in assistant_msgs[0]
    assert assistant_msgs[1]["reasoning_content"] == "thought2"


def test_build_messages_reasoning_content_echoed_with_tool_calls_on_last_turn():
    tc = ToolCall(id="call-1", name="tasks", arguments={"action": "list"})
    result = ENGINE._build_messages(
        [
            Message(
                role=Role.ASSISTANT,
                content=None,
                reasoning_content="thinking...",
                tool_calls=[tc],
            )
        ]
    )
    assert result[0]["reasoning_content"] == "thinking..."
    assert result[0]["tool_calls"][0]["id"] == "call-1"


def test_build_messages_no_reasoning_content_key_when_none():
    result = ENGINE._build_messages([Message(role=Role.ASSISTANT, content="Hi")])
    assert "reasoning_content" not in result[0]


### TOOLS


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


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]["function"]["name"] == "tasks"


def test_build_payload_no_reasoning_effort_by_default():
    msgs = [Message(role=Role.USER, content="Hi")]
    payload = ENGINE._build_payload(msgs, tools=None)
    assert "reasoning_effort" not in payload


def test_build_payload_reasoning_effort_included_when_set():
    engine = OpenAIEngine(model="grok-3-mini", api_key="k", reasoning_effort="high")
    payload = engine._build_payload([Message(role=Role.USER, content="Hi")], tools=None)
    assert payload["reasoning_effort"] == "high"


### RESPONSE PARSING


def _raw(
    content=None,
    tool_calls=None,
    finish_reason="stop",
    reasoning_content=None,
    prompt_tokens=10,
    completion_tokens=20,
):
    msg = {"content": content}
    if reasoning_content is not None:
        msg["reasoning_content"] = reasoning_content
    if tool_calls:
        msg["tool_calls"] = tool_calls
    return {
        "choices": [{"message": msg, "finish_reason": finish_reason}],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
        },
    }


def test_parse_response_text():
    r = ENGINE._parse_response(_raw(content="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 == ENGINE.provider
    assert r.model == ENGINE.model


def test_parse_response_tool_calls():
    raw_tool_calls = [
        {
            "id": "call-1",
            "function": {"name": "tasks", "arguments": '{"action": "list"}'},
        }
    ]
    r = ENGINE._parse_response(
        _raw(tool_calls=raw_tool_calls, finish_reason="tool_calls")
    )
    assert r.stop_reason == StopReason.TOOL_USE
    assert r.tool_calls
    assert len(r.tool_calls) == 1
    assert r.tool_calls[0].id == "call-1"
    assert r.tool_calls[0].name == "tasks"
    assert r.tool_calls[0].arguments == {"action": "list"}


def test_parse_response_deepseek_reasoning():
    r = DEEPSEEK._parse_response(
        _raw(
            content="The answer is 42.",
            reasoning_content="Let me think step by step...",
        )
    )
    assert r.content == "The answer is 42."
    assert r.reasoning_content == "Let me think step by step..."
    assert r.reasoning_signature is None  # OpenAI-compat has no signature


def test_parse_response_stop_reason_mapping():
    assert (
        ENGINE._parse_response(_raw(finish_reason="stop")).stop_reason
        == StopReason.END_TURN
    )
    assert (
        ENGINE._parse_response(_raw(finish_reason="tool_calls")).stop_reason
        == StopReason.TOOL_USE
    )
    assert (
        ENGINE._parse_response(_raw(finish_reason="length")).stop_reason
        == StopReason.MAX_TOKENS
    )


def test_parse_response_multiple_tool_calls():
    raw_calls = [
        {"id": "c1", "function": {"name": "tasks", "arguments": '{"action": "list"}'}},
        {"id": "c2", "function": {"name": "notes", "arguments": '{"action": "list"}'}},
    ]
    r = ENGINE._parse_response(_raw(tool_calls=raw_calls, finish_reason="tool_calls"))
    assert r.tool_calls
    assert len(r.tool_calls) == 2


### CHAT - chat()


@pytest.mark.asyncio
async def test_chat_returns_response():
    mock_resp = MagicMock()
    mock_resp.is_success = True
    mock_resp.json.return_value = _raw(content="Hello from OpenAI")

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

    assert response.content == "Hello from OpenAI"


@pytest.mark.asyncio
async def test_chat_raises_on_error():
    mock_resp = MagicMock()
    mock_resp.is_success = False
    mock_resp.status_code = 429
    mock_resp.text = "Rate limit"

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


### STREAM - stream()


@pytest.mark.asyncio
async def test_stream_yields_text_chunks():
    lines = [
        'data: {"choices": [{"delta": {"content": "Hello"}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {"content": " world"}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}',
        'data: {"choices": [], "usage": {"prompt_tokens": 8, "completion_tokens": 2}}',
        "data: [DONE]",
    ]

    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")])
        ]

    text_chunks = [c for c in chunks if c.content]
    assert [c.content for c in text_chunks] == ["Hello", " world"]
    summary = chunks[-1]
    assert summary.content is None
    assert summary.stop_reason == StopReason.END_TURN
    assert summary.usage
    assert summary.usage.input_tokens == 8
    assert summary.usage.output_tokens == 2


@pytest.mark.asyncio
async def test_stream_yields_reasoning_content():
    lines = [
        'data: {"choices": [{"delta": {"reasoning_content": "thinking..."}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {"content": "Answer"}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}',
        "data: [DONE]",
    ]

    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 chunks[0].reasoning_content == "thinking..."
    assert chunks[0].stop_reason is None
    assert chunks[1].content == "Answer"
    assert chunks[1].stop_reason is None


@pytest.mark.asyncio
async def test_stream_yields_tool_call():
    lines = [
        'data: {"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "tc_1", "function": {"name": "update_chat", "arguments": ""}}]}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": "{\\"title\\": \\"Test\\"}"}}]}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {}, "finish_reason": "tool_calls"}]}',
        "data: [DONE]",
    ]

    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")])
        ]

    tc_chunk = next(c for c in chunks if c.tool_calls)
    assert tc_chunk.tool_calls
    assert tc_chunk.tool_calls[0].id == "tc_1"
    assert tc_chunk.tool_calls[0].name == "update_chat"
    assert tc_chunk.tool_calls[0].arguments == {"title": "Test"}
    assert chunks[-1].stop_reason == StopReason.TOOL_USE


@pytest.mark.asyncio
async def test_stream_only_done_yields_summary_chunk():
    lines = ["data: [DONE]"]

    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) == 1
    assert chunks[0].content is None


@pytest.mark.asyncio
async def test_stream_hidden_reasoning_sentinel_emitted_immediately():
    """Engine with reasoning_effort set yields sentinel before any content chunks."""
    engine = OpenAIEngine(model="gpt-5", api_key="k", reasoning_effort="medium")
    lines = [
        'data: {"choices": [{"delta": {"content": "Hello"}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}',
        "data: [DONE]",
    ]

    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")])
        ]

    # sentinel is first
    assert (
        chunks[0].reasoning_content
        == "Reasoning in progress - tokens not available for this model"
    )
    assert chunks[0].content is None
    # content follows
    assert chunks[1].content == "Hello"
    # summary carries performance with near-zero latency
    summary = chunks[-1]
    assert summary.performance is not None
    assert summary.performance.latency < 0.1


@pytest.mark.asyncio
async def test_stream_no_sentinel_without_reasoning_effort():
    """Engine without reasoning_effort does not emit the hidden-reasoning sentinel."""
    lines = [
        'data: {"choices": [{"delta": {"content": "Hi"}, "finish_reason": null}]}',
        'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}',
        "data: [DONE]",
    ]

    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 all(c.reasoning_content is None for c in chunks)
    assert chunks[0].content == "Hi"


### HELPERS


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