"""Schema unit tests: function_schema and return_schema."""

from aicc.schema import function_schema, return_schema


def test_basic_types():
    def fn(a: int, b: str, c: float = 1.0) -> dict: ...

    s = function_schema(fn)
    assert s["type"] == "object"
    assert s["properties"]["a"]["type"] == "integer"
    assert s["properties"]["b"]["type"] == "string"
    assert s["properties"]["c"]["type"] == "number"
    assert set(s["required"]) == {"a", "b"}


def test_list_and_dict():
    def fn(items: list[int], mapping: dict[str, float]) -> dict: ...

    s = function_schema(fn)
    assert s["properties"]["items"]["type"] == "array"
    assert s["properties"]["items"]["items"]["type"] == "integer"
    assert s["properties"]["mapping"]["type"] == "object"


def test_optional():
    def fn(x: int | None = None) -> dict: ...

    s = function_schema(fn)
    assert s["properties"]["x"]["nullable"] is True
    assert "x" not in s["required"]


def test_pydantic_model_as_root_input():
    from pydantic import BaseModel

    class Input(BaseModel):
        x: int
        y: str

    def fn(data: Input) -> dict: ...

    s = function_schema(fn)
    assert "properties" in s
    assert "x" in s["properties"]
    assert "y" in s["properties"]


def test_return_schema_basic():
    def fn() -> int: ...
    assert return_schema(fn)["type"] == "integer"

    def g() -> list[str]: ...
    assert return_schema(g)["type"] == "array"

    def h() -> dict: ...
    assert return_schema(h)["type"] == "object"