# standard
from typing import Annotated, Literal

# third party
# custom
from sunwaee.modules.gen.engine.types import (
    Cost,
    Error,
    Message,
    Performance,
    Response,
    Role,
    StopReason,
    ToolCall,
    Usage,
    make_non_stream_performance,
    resolve_tokens,
)
from sunwaee.core.tools import tool

### TYPES


def test_message_defaults():
    m = Message(role=Role.USER, content="hello")
    assert m.reasoning_content is None
    assert m.reasoning_signature is None
    assert m.tool_call_id is None
    assert m.tool_calls is None


def test_message_with_reasoning():
    m = Message(
        role=Role.ASSISTANT,
        content="answer",
        reasoning_content="I think...",
        reasoning_signature="sig-abc",
    )
    assert m.reasoning_content == "I think..."
    assert m.reasoning_signature == "sig-abc"


def test_response_defaults():
    r = Response(
        content="hi",
        tool_calls=None,
        usage=Usage(input_tokens=10, output_tokens=5),
        stop_reason=StopReason.END_TURN,
    )
    assert r.reasoning_content is None
    assert r.reasoning_signature is None


def test_response_with_reasoning():
    r = Response(
        content=None,
        tool_calls=None,
        usage=Usage(input_tokens=10, output_tokens=5),
        stop_reason=StopReason.END_TURN,
        reasoning_content="chain of thought",
        reasoning_signature="sig-xyz",
    )
    assert r.reasoning_content == "chain of thought"
    assert r.reasoning_signature == "sig-xyz"


def test_response_identity_defaults():
    r = Response()
    assert r.provider == ""
    assert r.model == ""
    assert r.streaming is False
    assert r.error is None
    assert r.cost is None
    assert r.performance is None
    assert r.usage is None


def test_response_streaming_flag():
    r = Response(provider="anthropic", model="claude-x", streaming=True)
    assert r.streaming is True
    assert r.provider == "anthropic"
    assert r.model == "claude-x"


def test_error_defaults():
    e = Error()
    assert e.message == ""
    assert e.status_code == 0


def test_performance_defaults():
    p = Performance()
    assert p.latency == 0.0
    assert p.reasoning_duration == 0.0
    assert p.content_duration == 0.0
    assert p.total_duration == 0.0
    assert p.throughput == 0


def test_cost_defaults():
    c = Cost()
    assert c.input == 0.0
    assert c.output == 0.0
    assert c.total == 0.0


def test_usage_defaults():
    u = Usage()
    assert u.input_tokens == 0
    assert u.output_tokens == 0
    assert u.total_tokens == 0


def test_make_non_stream_performance_split():
    # 100 reasoning chars, 300 content chars → 25% / 75% split
    r = Response(
        reasoning_content="a" * 100, content="b" * 300, usage=Usage(output_tokens=40)
    )
    perf = make_non_stream_performance(r, total_duration=2.0)
    assert perf.total_duration == 2.0
    assert perf.latency == 2.0
    assert abs(perf.reasoning_duration - 0.5) < 1e-9  # 100/400 * 2.0
    assert abs(perf.content_duration - 1.5) < 1e-9  # 300/400 * 2.0
    assert perf.throughput == 20  # 40 tokens / 2.0 s


def test_make_non_stream_performance_no_content():
    # No chars → reasoning_duration and content_duration both 0
    r = Response(usage=Usage(output_tokens=10))
    perf = make_non_stream_performance(r, total_duration=1.0)
    assert perf.reasoning_duration == 0.0
    assert perf.content_duration == 0.0
    assert perf.throughput == 10


def test_make_non_stream_performance_reasoning_only():
    r = Response(reasoning_content="thinking...", usage=Usage(output_tokens=5))
    perf = make_non_stream_performance(r, total_duration=1.0)
    assert perf.reasoning_duration == 1.0
    assert perf.content_duration == 0.0


def test_resolve_tokens_no_total():
    # total==0: computed as input + output
    inp, out, tot = resolve_tokens(100, 50, 0)
    assert inp == 100
    assert out == 50
    assert tot == 150


def test_resolve_tokens_consistent():
    # input + output already equals total: no change
    inp, out, tot = resolve_tokens(100, 50, 150)
    assert inp == 100
    assert out == 50
    assert tot == 150


def test_resolve_tokens_reasoning_gap():
    # total > input + output: reasoning tokens fill the gap, output is corrected
    inp, out, tot = resolve_tokens(360, 26, 625)
    assert inp == 360
    assert tot == 625
    assert out == 625 - 360  # 265 — reasoning tokens absorbed into output


def test_tool_call():
    tc = ToolCall(id="call-1", name="tasks", arguments={"action": "list"})
    assert tc.id == "call-1"
    assert tc.name == "tasks"
    assert tc.arguments == {"action": "list"}


### TOOL DECORATOR


def test_tool_required_param():
    @tool("A simple tool.")
    def my_tool(name: str) -> str:
        return name

    assert my_tool._tool.name == "my_tool"
    assert my_tool._tool.description == "A simple tool."
    assert my_tool._tool.parameters["required"] == ["name"]
    assert my_tool._tool.parameters["properties"]["name"] == {"type": "string"}


def test_tool_optional_param_not_in_required():
    @tool("Optional param tool.")
    def my_tool(name: str | None = None) -> str:
        return name or ""

    assert "name" not in my_tool._tool.parameters["required"]


def test_tool_annotated_description():
    @tool("Annotated tool.")
    def my_tool(
        query: Annotated[str, "Search query"],
    ) -> str:
        return query

    prop = my_tool._tool.parameters["properties"]["query"]
    assert prop["type"] == "string"
    assert prop["description"] == "Search query"


def test_tool_literal_becomes_enum():
    @tool("Enum tool.")
    def my_tool(
        action: Annotated[Literal["create", "list", "delete"], "Action to perform"],
    ) -> str:
        return action

    prop = my_tool._tool.parameters["properties"]["action"]
    assert prop["type"] == "string"
    assert prop["enum"] == ["create", "list", "delete"]
    assert prop["description"] == "Action to perform"
    assert my_tool._tool.parameters["required"] == ["action"]


def test_tool_list_param():
    @tool("List param tool.")
    def my_tool(tags: Annotated[list[str] | None, "Tag list"] = None) -> str:
        return ""

    prop = my_tool._tool.parameters["properties"]["tags"]
    assert prop["type"] == "array"
    assert prop["items"] == {"type": "string"}
    assert "tags" not in my_tool._tool.parameters["required"]


def test_tool_primitive_types():
    @tool("Primitives.")
    def my_tool(count: int, score: float, active: bool) -> str:
        return ""

    props = my_tool._tool.parameters["properties"]
    assert props["count"]["type"] == "integer"
    assert props["score"]["type"] == "number"
    assert props["active"]["type"] == "boolean"
    assert my_tool._tool.parameters["required"] == ["count", "score", "active"]


def test_tool_fn_is_callable():
    @tool("Callable tool.")
    def my_tool(x: str) -> str:
        return x.upper()

    assert my_tool._tool.fn is not None
    assert my_tool._tool.fn("hello") == "HELLO"


def test_tool_optional_annotated_not_required():
    @tool("Optional annotated.")
    def my_tool(
        query: Annotated[str | None, "Optional query"] = None,
    ) -> str:
        return ""

    assert "query" not in my_tool._tool.parameters["required"]
    prop = my_tool._tool.parameters["properties"]["query"]
    assert prop["description"] == "Optional query"
