from __future__ import annotations

from typing import Annotated, Literal, Optional

from pillar._context import ToolContext
from pillar._schema import extract_schema


class TestPrimitives:
    def test_str(self) -> None:
        def fn(name: str) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["name"] == {"type": "string"}
        assert schema["required"] == ["name"]

    def test_int(self) -> None:
        def fn(count: int) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["count"] == {"type": "integer"}

    def test_float(self) -> None:
        def fn(price: float) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["price"] == {"type": "number"}

    def test_bool(self) -> None:
        def fn(active: bool) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["active"] == {"type": "boolean"}


class TestOptional:
    def test_optional_bracket(self) -> None:
        def fn(name: str | None) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["name"] == {"type": "string"}
        assert "required" not in schema or "name" not in schema.get("required", [])

    def test_optional_typing(self) -> None:
        def fn(name: Optional[str]) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["name"] == {"type": "string"}
        assert "name" not in schema.get("required", [])


class TestLiteral:
    def test_string_literal(self) -> None:
        def fn(priority: Literal["low", "medium", "high"]) -> None: ...
        schema = extract_schema(fn)
        prop = schema["properties"]["priority"]
        assert prop["type"] == "string"
        assert prop["enum"] == ["low", "medium", "high"]

    def test_int_literal(self) -> None:
        def fn(level: Literal[1, 2, 3]) -> None: ...
        schema = extract_schema(fn)
        prop = schema["properties"]["level"]
        assert prop["type"] == "integer"
        assert prop["enum"] == [1, 2, 3]


class TestList:
    def test_list_of_str(self) -> None:
        def fn(tags: list[str]) -> None: ...
        schema = extract_schema(fn)
        prop = schema["properties"]["tags"]
        assert prop == {"type": "array", "items": {"type": "string"}}

    def test_list_of_int(self) -> None:
        def fn(ids: list[int]) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["ids"]["items"] == {"type": "integer"}


class TestDict:
    def test_plain_dict(self) -> None:
        def fn(data: dict) -> None: ...
        schema = extract_schema(fn)
        assert schema["properties"]["data"] == {"type": "object"}


class TestAnnotated:
    def test_annotated_with_string_description(self) -> None:
        def fn(email: Annotated[str, "Customer email"]) -> None: ...
        schema = extract_schema(fn)
        prop = schema["properties"]["email"]
        assert prop["type"] == "string"
        assert prop["description"] == "Customer email"


class TestExcludedParams:
    def test_ctx_excluded(self) -> None:
        def fn(email: str, ctx: ToolContext) -> None: ...
        schema = extract_schema(fn)
        assert "ctx" not in schema["properties"]
        assert schema["required"] == ["email"]

    def test_self_excluded(self) -> None:
        def fn(self, name: str) -> None: ...
        schema = extract_schema(fn)
        assert "self" not in schema["properties"]

    def test_context_name_excluded(self) -> None:
        def fn(name: str, context: ToolContext) -> None: ...
        schema = extract_schema(fn)
        assert "context" not in schema["properties"]


class TestDefaults:
    def test_param_with_default_not_required(self) -> None:
        def fn(name: str, priority: str = "medium") -> None: ...
        schema = extract_schema(fn)
        assert schema["required"] == ["name"]
        assert "priority" in schema["properties"]


class TestMultipleParams:
    def test_mixed_params(self) -> None:
        def fn(
            subject: str,
            body: str,
            priority: Literal["low", "medium", "high"] = "medium",
            ctx: ToolContext = None,
        ) -> dict: ...
        schema = extract_schema(fn)
        assert set(schema["properties"]) == {"subject", "body", "priority"}
        assert schema["required"] == ["subject", "body"]
        assert schema["properties"]["priority"]["enum"] == ["low", "medium", "high"]


class TestAsyncFunction:
    def test_async_fn(self) -> None:
        async def fn(email: str, ctx: ToolContext) -> dict: ...
        schema = extract_schema(fn)
        assert schema["properties"]["email"] == {"type": "string"}
        assert "ctx" not in schema["properties"]
