from unittest.mock import MagicMock, patch
from pydantic import BaseModel
from app.agents.cerebras_client import CerebrasClient
from app.agents.cerebras_errors import CerebrasErrorKind, CerebrasError


class FakeOutput(BaseModel):
    answer: str
    confidence: int


def _make_mock_response(content: str):
    msg = MagicMock()
    msg.content = content
    choice = MagicMock()
    choice.message = msg
    resp = MagicMock()
    resp.choices = [choice]
    # Real SDK responses always carry a numeric usage -> the speed-log line divides
    # by it, so a bare MagicMock here would crash. Give it a realistic int.
    resp.usage.completion_tokens = 12
    return resp


def test_structured_complete_parses_json():
    client = CerebrasClient.__new__(CerebrasClient)
    client._health = {"status": "ok"}
    client._rate_limit_until = 0.0
    mock_sdk = MagicMock()
    client._client = mock_sdk
    mock_sdk.chat.completions.create.return_value = _make_mock_response(
        '{"answer": "entropy is disorder", "confidence": 90}'
    )
    result = client.structured_complete(
        [{"role": "user", "content": "define entropy"}], FakeOutput
    )
    assert isinstance(result, FakeOutput)
    assert result.confidence == 90


def test_schema_build_rejects_open_ended_dict_field():
    client = CerebrasClient.__new__(CerebrasClient)
    client._health = {"status": "ok"}
    client._rate_limit_until = 0.0
    client._client = MagicMock()

    class HasOpenDict(BaseModel):
        panel_ids_by_requirement: dict[str, list[str]]

    try:
        client._build_schema(HasOpenDict)
        assert False, "Should have raised ValueError for an open-ended dict field"
    except ValueError as err:
        assert "open-ended dict field" in str(err)


def test_schema_strips_defs_and_sets_additional_properties():
    client = CerebrasClient.__new__(CerebrasClient)
    client._health = {"status": "ok"}
    client._rate_limit_until = 0.0
    client._client = MagicMock()

    class Nested(BaseModel):
        value: int

    class Outer(BaseModel):
        nested: Nested

    schema = client._build_schema(Outer)
    assert "$defs" not in schema
    assert schema["additionalProperties"] is False


def test_rate_limit_short_circuit():
    import time
    client = CerebrasClient.__new__(CerebrasClient)
    client._health = {"status": "ok"}
    client._rate_limit_until = time.time() + 60
    client._client = MagicMock()

    try:
        client.structured_complete([{"role": "user", "content": "test"}], FakeOutput)
        assert False, "Should have raised CerebrasError"
    except CerebrasError as err:
        assert err.kind == CerebrasErrorKind.RATE_LIMITED


def test_complete_with_tools_returns_message_with_tool_calls():
    client = CerebrasClient.__new__(CerebrasClient)
    client._health = {"status": "ok"}
    client._rate_limit_until = 0.0
    mock_sdk = MagicMock()
    client._client = mock_sdk

    tool_call = MagicMock()
    tool_call.id = "call_1"
    tool_call.function.name = "web_search"
    tool_call.function.arguments = '{"query": "VAE"}'
    message = MagicMock()
    message.tool_calls = [tool_call]
    message.content = ""
    resp = MagicMock()
    resp.choices = [MagicMock(message=message)]
    mock_sdk.chat.completions.create.return_value = resp

    tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
    result = client.complete_with_tools([{"role": "user", "content": "what is a VAE"}], tools)

    assert result.tool_calls[0].function.name == "web_search"
    # tools + tool_choice forwarded to the SDK
    _, kwargs = mock_sdk.chat.completions.create.call_args
    assert kwargs["tools"] == tools
    assert kwargs["tool_choice"] == "auto"


def test_complete_with_tools_passes_through_no_tool_answer():
    client = CerebrasClient.__new__(CerebrasClient)
    client._health = {"status": "ok"}
    client._rate_limit_until = 0.0
    mock_sdk = MagicMock()
    client._client = mock_sdk

    message = MagicMock()
    message.tool_calls = None
    message.content = "A VAE is a generative model."
    resp = MagicMock()
    resp.choices = [MagicMock(message=message)]
    mock_sdk.chat.completions.create.return_value = resp

    result = client.complete_with_tools([{"role": "user", "content": "vae?"}], [])
    assert result.tool_calls is None
    assert "generative model" in result.content