"""Tests for structured output (output_schema)."""
import pytest
from pydantic import BaseModel

from logos_adk.agent import Agent
from logos_adk.providers.mock import MockProvider
from logos_adk.types import Response, StructuredResponse, Usage


class Sentiment(BaseModel):
    sentiment: str
    confidence: float


class Person(BaseModel):
    name: str
    age: int


class TestOutputSchemaBuilder:
    def test_builder_sets_schema(self):
        provider = MockProvider()
        agent = Agent("a", model=provider).output_schema(Sentiment)
        assert agent._config.output_schema is Sentiment

    def test_builder_returns_self(self):
        provider = MockProvider()
        agent = Agent("a", model=provider)
        result = agent.output_schema(Sentiment)
        assert result is agent


class TestStructuredOutputRun:
    @pytest.mark.asyncio
    async def test_parses_json_response(self):
        provider = MockProvider(responses=[
            Response(
                content='{"sentiment": "positive", "confidence": 0.95}',
                tool_calls=[],
                usage=Usage(10, 20),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider).output_schema(Sentiment)
        response = await agent.run("Analyze: great!")

        assert isinstance(response, StructuredResponse)
        assert response.parsed is not None
        assert response.parsed.sentiment == "positive"
        assert response.parsed.confidence == 0.95

    @pytest.mark.asyncio
    async def test_preserves_original_fields(self):
        provider = MockProvider(responses=[
            Response(
                content='{"sentiment": "negative", "confidence": 0.8}',
                tool_calls=[],
                usage=Usage(5, 10),
                raw={"model": "test"},
            ),
        ])
        agent = Agent("a", model=provider).output_schema(Sentiment)
        response = await agent.run("Analyze: bad!")

        assert response.content == '{"sentiment": "negative", "confidence": 0.8}'
        assert response.usage.input_tokens == 5
        assert response.usage.output_tokens == 10
        assert response.raw == {"model": "test"}

    @pytest.mark.asyncio
    async def test_invalid_json_returns_none_parsed(self):
        provider = MockProvider(responses=[
            Response(
                content="I think the sentiment is positive",
                tool_calls=[],
                usage=Usage(5, 10),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider).output_schema(Sentiment)
        response = await agent.run("Analyze")

        assert isinstance(response, StructuredResponse)
        assert response.parsed is None
        assert response.content == "I think the sentiment is positive"

    @pytest.mark.asyncio
    async def test_schema_validation_error_returns_none_parsed(self):
        provider = MockProvider(responses=[
            Response(
                content='{"wrong_field": "value"}',
                tool_calls=[],
                usage=Usage(5, 10),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider).output_schema(Sentiment)
        response = await agent.run("Analyze")

        assert isinstance(response, StructuredResponse)
        assert response.parsed is None

    @pytest.mark.asyncio
    async def test_markdown_code_block_json(self):
        provider = MockProvider(responses=[
            Response(
                content='```json\n{"sentiment": "neutral", "confidence": 0.5}\n```',
                tool_calls=[],
                usage=Usage(5, 10),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider).output_schema(Sentiment)
        response = await agent.run("Analyze")

        assert isinstance(response, StructuredResponse)
        assert response.parsed is not None
        assert response.parsed.sentiment == "neutral"

    @pytest.mark.asyncio
    async def test_per_call_output_schema(self):
        """output_schema can be passed to run() directly."""
        provider = MockProvider(responses=[
            Response(
                content='{"name": "Alice", "age": 30}',
                tool_calls=[],
                usage=Usage(5, 10),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider)
        response = await agent.run("Who?", output_schema=Person)

        assert isinstance(response, StructuredResponse)
        assert response.parsed is not None
        assert response.parsed.name == "Alice"
        assert response.parsed.age == 30

    @pytest.mark.asyncio
    async def test_no_schema_returns_plain_response(self):
        provider = MockProvider(responses=[
            Response(
                content="just text",
                tool_calls=[],
                usage=Usage(5, 10),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider)
        response = await agent.run("Hello")

        assert isinstance(response, Response)
        assert not isinstance(response, StructuredResponse)

    @pytest.mark.asyncio
    async def test_schema_forwarded_to_provider(self):
        provider = MockProvider(responses=[
            Response(
                content='{"sentiment": "positive", "confidence": 0.9}',
                tool_calls=[],
                usage=Usage(5, 10),
                raw={},
            ),
        ])
        agent = Agent("a", model=provider).output_schema(Sentiment)
        await agent.run("Analyze")

        assert provider.call_history[0]["output_schema"] is Sentiment
