"""Tests for native structured output support in PydanticAIProvider."""

from __future__ import annotations

import asyncio
import json
from unittest.mock import MagicMock

import pytest
from pydantic import BaseModel

pydantic_ai = pytest.importorskip("pydantic_ai")

from parsantic.extract.options import ExtractOptions  # noqa: E402
from parsantic.extract.providers.base import InferenceRequest  # noqa: E402
from parsantic.extract.providers.pydantic_ai_provider import (  # noqa: E402
    PydanticAIProvider,
    _extract_raw_json_from_messages,
    _parse_model_spec,
    _requires_explicit_api_key,
    _result_to_json_string,
)

# ── helpers ──────────────────────────────────────────────────────────────


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


def _make_provider(*, supports_native: bool = False) -> PydanticAIProvider:
    """Create a PydanticAIProvider without hitting the network."""
    provider = object.__new__(PydanticAIProvider)
    provider.model_id = "openai:gpt-4o-mini"
    provider.api_key = None
    provider.base_url = None
    provider.project_id = None
    provider.region = None
    provider.service_account_file = None
    provider.max_concurrency = 8
    provider._agent = MagicMock()
    provider._supports_native = supports_native
    return provider


class _FakeSyncStreamResult:
    def __init__(
        self,
        *,
        outputs: list[object] | None = None,
        texts: list[str] | None = None,
        final_output: object | None = None,
    ) -> None:
        self._outputs = outputs or []
        self._texts = texts or []
        self._final_output = final_output

    def stream_output(self, *, debounce_by: float | None = 0.1):
        del debounce_by
        yield from self._outputs

    def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1):
        del delta, debounce_by
        yield from self._texts

    def get_output(self):
        return self._final_output


class _FakeAsyncStreamResult:
    def __init__(
        self,
        *,
        outputs: list[object] | None = None,
        texts: list[str] | None = None,
        final_output: object | None = None,
    ) -> None:
        self._outputs = outputs or []
        self._texts = texts or []
        self._final_output = final_output

    async def stream_output(self, *, debounce_by: float | None = 0.1):
        del debounce_by
        for output in self._outputs:
            yield output

    async def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1):
        del delta, debounce_by
        for text in self._texts:
            yield text

    async def get_output(self):
        return self._final_output


class _FakeAsyncStreamContext:
    def __init__(self, result: _FakeAsyncStreamResult) -> None:
        self._result = result

    async def __aenter__(self) -> _FakeAsyncStreamResult:
        return self._result

    async def __aexit__(self, exc_type, exc, tb) -> bool:
        del exc_type, exc, tb
        return False


# ── _parse_model_spec ────────────────────────────────────────────────────


@pytest.mark.parametrize(
    "spec,expected",
    [
        ("openai:gpt-4o", ("openai", "gpt-4o")),
        ("gpt-4o", ("openai", "gpt-4o")),
        ("claude-3-opus", ("anthropic", "claude-3-opus")),
        ("gemini-2.0-flash", ("gemini", "gemini-2.0-flash")),
        ("custom-model", ("openai", "custom-model")),
    ],
    ids=["explicit", "bare-gpt", "bare-claude", "bare-gemini", "unknown-bare"],
)
def test_parse_model_spec(spec, expected):
    assert _parse_model_spec(spec) == expected


@pytest.mark.parametrize(
    ("spec", "expected"),
    [
        ("gemini:gemini-2.5-flash", True),
        ("gemini-2.0-flash", True),
        ("openai:gpt-4o-mini", False),
        ("gpt-4o-mini", False),
    ],
)
def test_requires_explicit_api_key(spec, expected):
    assert _requires_explicit_api_key(spec) is expected


def test_gemini_provider_without_api_key_raises_clear_error():
    with pytest.raises(
        ValueError,
        match="Gemini models require GEMINI_API_KEY or an explicit api_key provider kwarg",
    ):
        PydanticAIProvider(model_id="gemini:gemini-2.5-flash")


# ── _result_to_json_string ───────────────────────────────────────────────


@pytest.mark.parametrize(
    "input_val,check",
    [
        ('{"a":1}', lambda r: r == '{"a":1}'),
        ({"name": "Alice", "age": 30}, lambda r: json.loads(r) == {"name": "Alice", "age": 30}),
        ([1, 2, 3], lambda r: json.loads(r) == [1, 2, 3]),
        (_SampleModel(name="Bob", age=25), lambda r: json.loads(r) == {"name": "Bob", "age": 25}),
        ({"text": "日本語"}, lambda r: "日本語" in r),
    ],
    ids=["string-passthrough", "dict", "list", "pydantic-model", "unicode"],
)
def test_result_to_json_string(input_val, check):
    assert check(_result_to_json_string(input_val))


def test_result_to_json_string_non_native_values():
    from datetime import datetime
    from decimal import Decimal

    result = _result_to_json_string({"ts": datetime(2025, 1, 1), "val": Decimal("3.14")})
    parsed = json.loads(result)
    assert "2025" in parsed["ts"]
    assert float(parsed["val"]) == pytest.approx(3.14)


# ── _extract_raw_json_from_messages ──────────────────────────────────────


class TestExtractRawJsonFromMessages:
    def test_empty_messages(self):
        assert _extract_raw_json_from_messages([]) is None

    def test_extracts_tool_call_args(self):
        from pydantic_ai.messages import ModelResponse, ToolCallPart

        tool_part = ToolCallPart(
            tool_name="final_result",
            args='{"name":"Alice","age":30}',
            tool_call_id="tc_1",
        )
        msg = ModelResponse(parts=[tool_part])
        result = _extract_raw_json_from_messages([msg])
        assert result == '{"name":"Alice","age":30}'

    def test_extracts_text_content(self):
        from pydantic_ai.messages import ModelResponse, TextPart

        text_part = TextPart(content='{"name":"Bob"}')
        msg = ModelResponse(parts=[text_part])
        result = _extract_raw_json_from_messages([msg])
        assert result == '{"name":"Bob"}'

    def test_skips_empty_text(self):
        from pydantic_ai.messages import ModelResponse, TextPart

        text_part = TextPart(content="   ")
        msg = ModelResponse(parts=[text_part])
        assert _extract_raw_json_from_messages([msg]) is None

    def test_tool_call_prioritized_over_text(self):
        """ToolCallPart should be preferred even when TextPart comes first."""
        from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart

        # TextPart first, ToolCallPart second — tool call should win
        parts = [
            TextPart(content="some reasoning"),
            ToolCallPart(
                tool_name="final_result",
                args='{"from":"tool"}',
                tool_call_id="tc_2",
            ),
        ]
        msg = ModelResponse(parts=parts)
        result = _extract_raw_json_from_messages([msg])
        assert result == '{"from":"tool"}'

    def test_scans_from_last_message(self):
        from pydantic_ai.messages import ModelResponse, TextPart

        msg1 = ModelResponse(parts=[TextPart(content='{"old":true}')])
        msg2 = ModelResponse(parts=[TextPart(content='{"new":true}')])
        result = _extract_raw_json_from_messages([msg1, msg2])
        assert result == '{"new":true}'


# ── supports_native_structured_output ────────────────────────────────────


def test_supports_native_structured_output():
    assert _make_provider(supports_native=False).supports_native_structured_output() is False
    assert _make_provider(supports_native=True).supports_native_structured_output() is True


# ── _resolve_output_type ─────────────────────────────────────────────────


class TestResolveOutputType:
    def test_returns_none_when_appropriate(self):
        provider = _make_provider(supports_native=True)
        assert provider._resolve_output_type(_SampleModel, "prompt") is None
        assert provider._resolve_output_type(None, "native") is None
        assert (
            _make_provider(supports_native=False)._resolve_output_type(_SampleModel, "auto") is None
        )

    def test_returns_native_output_when_appropriate(self):
        from pydantic_ai import NativeOutput

        assert isinstance(
            _make_provider(supports_native=False)._resolve_output_type(_SampleModel, "native"),
            NativeOutput,
        )
        assert isinstance(
            _make_provider(supports_native=True)._resolve_output_type(_SampleModel, "auto"),
            NativeOutput,
        )


# ── infer with native structured output ──────────────────────────────────


class TestInferWithNativeOutput:
    def _make_async_mock_result(self, output):
        """Create a coroutine-returning mock for agent.run."""

        mock_result = MagicMock()
        mock_result.output = output

        async def _fake_run(*args, **kwargs):
            return mock_result

        return mock_result, _fake_run

    def test_infer_passes_kwargs_to_run(self):
        provider = _make_provider(supports_native=False)
        mock_result, fake_run = self._make_async_mock_result('{"name":"Alice","age":30}')
        provider._agent.run = MagicMock(side_effect=fake_run)

        results = provider.infer(
            ["extract name"],
            target_type=_SampleModel,
            structured_output="prompt",
        )
        assert results == ['{"name":"Alice","age":30}']
        provider._agent.run.assert_called_once()

    def test_infer_with_native_mode_uses_output_type_override(self):
        provider = _make_provider(supports_native=False)
        mock_result, fake_run = self._make_async_mock_result(_SampleModel(name="Alice", age=30))
        provider._agent.run = MagicMock(side_effect=fake_run)

        results = provider.infer(
            ["extract name"],
            target_type=_SampleModel,
            structured_output="native",
        )

        # Should serialize the pydantic model back to JSON
        parsed = json.loads(results[0])
        assert parsed == {"name": "Alice", "age": 30}

        # run should have been called with output_type override
        call_kwargs = provider._agent.run.call_args
        assert call_kwargs.kwargs.get("output_type") is not None

    def test_infer_native_failure_reraises_when_no_raw_json(self):
        """When native fails with no message history, should re-raise."""
        provider = _make_provider(supports_native=False)

        async def _fail(*args, **kwargs):
            raise RuntimeError("native failed")

        provider._agent.run = MagicMock(side_effect=_fail)

        with pytest.raises(RuntimeError, match="native failed"):
            provider.infer(
                ["extract name"],
                target_type=_SampleModel,
                structured_output="native",
            )

    def test_infer_native_fallback_extracts_raw_json_from_captured_messages(self):
        """When native fails, extract raw JSON from capture_run_messages."""
        from unittest.mock import patch

        from pydantic_ai.messages import ModelResponse, ToolCallPart

        provider = _make_provider(supports_native=False)

        # Mock capture_run_messages to return messages with tool call
        captured_messages = [
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name="final_result",
                        args='{"name":"Alice","age":"thirty"}',
                        tool_call_id="tc_1",
                    )
                ]
            )
        ]

        # agent.run raises on native call
        async def _fail(*args, **kwargs):
            raise RuntimeError("validation failed")

        provider._agent.run = MagicMock(side_effect=_fail)

        with patch(
            "parsantic.extract.providers.pydantic_ai_provider.capture_run_messages"
        ) as mock_capture:
            # Make the context manager yield our captured messages list
            mock_ctx = MagicMock()
            mock_ctx.__enter__ = MagicMock(return_value=captured_messages)
            mock_ctx.__exit__ = MagicMock(return_value=False)
            mock_capture.return_value = mock_ctx

            results = provider.infer(
                ["extract name"],
                target_type=_SampleModel,
                structured_output="native",
            )
            assert results == ['{"name":"Alice","age":"thirty"}']


# ── ExtractOptions structured_output field ───────────────────────────────


def test_extract_options_structured_output():
    assert ExtractOptions().structured_output == "auto"
    for val in ("auto", "native", "prompt"):
        assert ExtractOptions(structured_output=val).structured_output == val


# ── streaming inference ─────────────────────────────────────────────────


class TestStreamingInfer:
    def test_infer_stream_uses_stream_text_for_prompt_mode(self):
        provider = _make_provider(supports_native=True)
        provider._agent.run_stream_sync.return_value = _FakeSyncStreamResult(
            texts=['{"name":"A"', '{"name":"Alice","age":30}'],
            final_output='{"name":"Alice","age":30}',
        )

        snapshots = list(provider.infer_stream("extract name", structured_output="prompt"))

        assert snapshots == ['{"name":"A"', '{"name":"Alice","age":30}']
        provider._agent.run_stream_sync.assert_called_once()

    def test_infer_media_stream_serializes_native_output_snapshots(self):
        from parsantic.extract.media.attachments import Attachment, AttachmentKind

        provider = _make_provider(supports_native=True)
        provider._agent.run_stream_sync.return_value = _FakeSyncStreamResult(
            outputs=[
                {"name": "Ali"},
                _SampleModel(name="Alice", age=30),
            ],
            final_output=_SampleModel(name="Alice", age=30),
        )
        request = InferenceRequest(
            prompt="extract name",
            attachments=(Attachment(kind=AttachmentKind.PDF, source=b"%PDF-1.7"),),
        )

        snapshots = list(
            provider.infer_media_stream(
                request,
                target_type=_SampleModel,
                structured_output="native",
            )
        )

        assert json.loads(snapshots[0]) == {"name": "Ali"}
        assert json.loads(snapshots[-1]) == {"name": "Alice", "age": 30}

    def test_ainfer_stream_uses_async_streaming_api(self):
        provider = _make_provider(supports_native=True)
        provider._agent.run_stream.return_value = _FakeAsyncStreamContext(
            _FakeAsyncStreamResult(
                outputs=[
                    {"name": "Ali"},
                    _SampleModel(name="Alice", age=30),
                ],
                final_output=_SampleModel(name="Alice", age=30),
            )
        )

        async def _run() -> list[str]:
            return [
                snapshot
                async for snapshot in provider.ainfer_stream(
                    "extract name",
                    target_type=_SampleModel,
                    structured_output="native",
                )
            ]

        snapshots = asyncio.run(_run())

        assert json.loads(snapshots[0]) == {"name": "Ali"}
        assert json.loads(snapshots[-1]) == {"name": "Alice", "age": 30}
