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

# third party
import pytest

# custom
from sunwaee.modules.gen.engine.providers.anthropic import AnthropicEngine
from sunwaee.modules.gen.engine.types import (
    Message,
    Role,
    StopReason,
    Tool,
    ToolCall,
)

ENGINE = AnthropicEngine(model="claude-3-5-haiku", api_key="test-key")

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


### MESSAGES


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


def test_build_messages_assistant_text():
    msgs = [Message(role=Role.ASSISTANT, content="Hi there")]
    result = ENGINE._build_messages(msgs)
    assert result == [
        {"role": "assistant", "content": [{"type": "text", "text": "Hi there"}]}
    ]


def test_build_messages_assistant_with_reasoning():
    msgs = [
        Message(
            role=Role.ASSISTANT,
            content="Answer",
            reasoning_content="I think...",
            reasoning_signature="sig-123",
        )
    ]
    result = ENGINE._build_messages(msgs)
    assert len(result) == 1
    parts = result[0]["content"]
    assert parts[0] == {
        "type": "thinking",
        "thinking": "I think...",
        "signature": "sig-123",
    }
    assert parts[1] == {"type": "text", "text": "Answer"}


def test_build_messages_assistant_reasoning_no_content():
    msgs = [
        Message(
            role=Role.ASSISTANT,
            reasoning_content="thinking only",
            reasoning_signature="s",
        )
    ]
    result = ENGINE._build_messages(msgs)
    parts = result[0]["content"]
    assert parts[0]["type"] == "thinking"
    assert len(parts) == 1  # no text block added when content is None


def test_build_messages_assistant_with_tool_calls():
    tc = ToolCall(id="call-1", name="tasks", arguments={"action": "list"})
    msgs = [Message(role=Role.ASSISTANT, content="Using tool", tool_calls=[tc])]
    result = ENGINE._build_messages(msgs)
    parts = result[0]["content"]
    assert parts[0] == {"type": "text", "text": "Using tool"}
    assert parts[1] == {
        "type": "tool_use",
        "id": "call-1",
        "name": "tasks",
        "input": {"action": "list"},
    }


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


def test_build_messages_system_excluded():
    """System messages are stripped by _split_system before _build_messages is called."""
    msgs = [Message(role=Role.USER, content="Hi")]
    system, rest = ENGINE._split_system(
        [Message(role=Role.SYSTEM, content="You are Sun.")] + msgs
    )
    assert system == "You are Sun."
    built = ENGINE._build_messages(rest)
    assert all(m["role"] != "system" for m in built)


def test_split_system_only_promotes_to_user():
    """System-only input: no top-level system field, content becomes a user message."""
    system, rest = ENGINE._split_system(
        [Message(role=Role.SYSTEM, content="Respond with ok.")]
    )
    assert system is None
    assert len(rest) == 1
    assert rest[0].role == Role.USER
    assert rest[0].content == "Respond with ok."


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["system"] == "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]["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["thinking"]["type"] == "enabled"
    assert payload["thinking"]["budget_tokens"] >= 1024


def test_build_payload_thinking_budget_custom():
    engine = AnthropicEngine(
        model="claude-3-5-haiku", api_key="k", max_tokens=4096, thinking_budget=2000
    )
    payload = engine._build_payload([Message(role=Role.USER, content="Hi")], tools=None)
    assert payload["thinking"]["budget_tokens"] == 2000


def test_build_payload_thinking_budget_default_from_max_tokens():
    engine = AnthropicEngine(model="m", api_key="k", max_tokens=2048)
    assert engine.thinking_budget == max(1024, 2048 - 1024)


def test_build_messages_reasoning_only_on_last_turn():
    """Reasoning block only on the last assistant turn (safe for cross-engine history)."""
    msgs = [
        Message(role=Role.USER, content="First question"),
        Message(
            role=Role.ASSISTANT,
            content="First answer",
            reasoning_content="thinking1",
            reasoning_signature="sig1",
        ),
        Message(role=Role.USER, content="Second question"),
        Message(
            role=Role.ASSISTANT,
            content="Second answer",
            reasoning_content="thinking2",
            reasoning_signature="sig2",
        ),
    ]
    result = ENGINE._build_messages(msgs)
    assistant_msgs = [m for m in result if m["role"] == "assistant"]
    assert len(assistant_msgs) == 2
    # First turn: no thinking block
    assert assistant_msgs[0]["content"][0]["type"] == "text"
    # Last turn: thinking block present
    assert assistant_msgs[1]["content"][0]["type"] == "thinking"
    assert assistant_msgs[1]["content"][0]["thinking"] == "thinking2"


### TOOLS


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


### RESPONSE PARSING


def _raw(content_blocks, stop_reason="end_turn", input_tokens=10, output_tokens=20):
    return {
        "content": content_blocks,
        "stop_reason": stop_reason,
        "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
    }


def test_parse_response_text():
    r = ENGINE._parse_response(_raw([{"type": "text", "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 == "anthropic"
    assert r.model == ENGINE.model


def test_parse_response_thinking():
    r = ENGINE._parse_response(
        _raw(
            [
                {"type": "thinking", "thinking": "I think...", "signature": "sig-abc"},
                {"type": "text", "text": "Answer"},
            ]
        )
    )
    assert r.content == "Answer"
    assert r.reasoning_content == "I think..."
    assert r.reasoning_signature == "sig-abc"


def test_parse_response_tool_use():
    r = ENGINE._parse_response(
        _raw(
            [
                {
                    "type": "tool_use",
                    "id": "call-1",
                    "name": "tasks",
                    "input": {"action": "list"},
                },
            ],
            stop_reason="tool_use",
        )
    )
    assert r.content is None
    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_thinking_and_tool_use():
    r = ENGINE._parse_response(
        _raw(
            [
                {
                    "type": "thinking",
                    "thinking": "Should I use a tool?",
                    "signature": "s",
                },
                {
                    "type": "tool_use",
                    "id": "c1",
                    "name": "notes",
                    "input": {"action": "list"},
                },
            ],
            stop_reason="tool_use",
        )
    )
    assert r.reasoning_content == "Should I use a tool?"
    assert r.tool_calls
    assert r.tool_calls[0].name == "notes"


def test_parse_response_multiple_tool_calls():
    r = ENGINE._parse_response(
        _raw(
            [
                {
                    "type": "tool_use",
                    "id": "c1",
                    "name": "tasks",
                    "input": {"action": "list"},
                },
                {
                    "type": "tool_use",
                    "id": "c2",
                    "name": "notes",
                    "input": {"action": "list"},
                },
            ],
            stop_reason="tool_use",
        )
    )
    assert r.tool_calls
    assert len(r.tool_calls) == 2


def test_parse_response_stop_reason_mapping():
    assert (
        ENGINE._parse_response(_raw([], stop_reason="end_turn")).stop_reason
        == StopReason.END_TURN
    )
    assert (
        ENGINE._parse_response(_raw([], stop_reason="tool_use")).stop_reason
        == StopReason.TOOL_USE
    )
    assert (
        ENGINE._parse_response(_raw([], stop_reason="max_tokens")).stop_reason
        == StopReason.MAX_TOKENS
    )


### CHAT - chat()


@pytest.mark.asyncio
async def test_chat_returns_response():
    raw = _raw([{"type": "text", "text": "Hi from Anthropic"}])
    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 Anthropic"


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

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


### STREAM - stream()


@pytest.mark.asyncio
async def test_stream_yields_reasoning_content_and_text():
    lines = [
        'data: {"type": "message_start", "message": {"usage": {"input_tokens": 5}}}',
        'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}',
        'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "I think"}}',
        'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "sig-abc"}}',
        'data: {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}',
        'data: {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Hello"}}',
        'data: {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " world"}}',
        'data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}}',
        'data: {"type": "message_stop"}',
    ]

    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 chunk, signature chunk, 2 text chunks, summary
    assert chunks[0].reasoning_content == "I think"
    assert chunks[0].stop_reason is None
    assert chunks[1].reasoning_signature == "sig-abc"
    assert chunks[1].stop_reason is None
    assert chunks[2].content == "Hello"
    assert chunks[2].stop_reason is None
    assert chunks[3].content == " world"
    assert chunks[3].stop_reason is None
    summary = chunks[-1]
    assert summary.content is None
    assert summary.usage
    assert summary.usage.input_tokens == 5
    assert summary.usage.output_tokens == 2
    assert summary.stop_reason == StopReason.END_TURN


@pytest.mark.asyncio
async def test_stream_yields_tool_call():
    lines = [
        'data: {"type": "message_start", "message": {"usage": {"input_tokens": 10}}}',
        'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "tc_1", "name": "update_chat", "input": {}}}',
        'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\\"title\\":"}}',
        'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": " \\"Test\\"}"}}',
        'data: {"type": "content_block_stop", "index": 0}',
        'data: {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 20}}',
        'data: {"type": "message_stop"}',
    ]

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

    # tool call chunk + summary
    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 tc_chunk.stop_reason is None
    summary = chunks[-1]
    assert summary.stop_reason == StopReason.TOOL_USE
    assert summary.usage
    assert summary.usage.output_tokens == 20


### HELPERS


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