"""Tests for structured output parsing."""

from __future__ import annotations

from typing import Literal

import pytest
from pydantic import BaseModel

from koan.output.parser import OutputParseError, parse_response

# -- Test models -----------------------------------------------------------

class SimpleResult(BaseModel):
    answer: str
    confidence: float


class NestedResult(BaseModel):
    class Finding(BaseModel):
        title: str
        score: float

    summary: str
    findings: list[Finding]
    sentiment: Literal["positive", "neutral", "negative"]


# -- Tests: Successful parsing ---------------------------------------------

class TestSuccessfulParsing:
    def test_pure_json(self) -> None:
        text = '{"answer": "42", "confidence": 0.95}'
        result = parse_response(text, SimpleResult)
        assert result.answer == "42"
        assert result.confidence == 0.95

    def test_json_in_code_fence(self) -> None:
        text = '```json\n{"answer": "hello", "confidence": 0.8}\n```'
        result = parse_response(text, SimpleResult)
        assert result.answer == "hello"

    def test_json_in_bare_code_fence(self) -> None:
        text = '```\n{"answer": "hello", "confidence": 0.8}\n```'
        result = parse_response(text, SimpleResult)
        assert result.answer == "hello"

    def test_json_with_surrounding_prose(self) -> None:
        text = 'Here is the result:\n{"answer": "test", "confidence": 0.7}\nHope that helps!'
        result = parse_response(text, SimpleResult)
        assert result.answer == "test"

    def test_nested_model(self) -> None:
        text = """{
            "summary": "Good quarter",
            "findings": [
                {"title": "Revenue up", "score": 0.9},
                {"title": "Costs stable", "score": 0.85}
            ],
            "sentiment": "positive"
        }"""
        result = parse_response(text, NestedResult)
        assert result.summary == "Good quarter"
        assert len(result.findings) == 2
        assert result.findings[0].title == "Revenue up"
        assert result.sentiment == "positive"

    def test_json_array_extraction(self) -> None:
        """Arrays should also be extractable (for list models)."""
        from koan.output.parser import _extract_json

        text = 'Sure! [{"name": "a"}, {"name": "b"}]'
        extracted = _extract_json(text)
        assert extracted.startswith("[")

    def test_whitespace_tolerance(self) -> None:
        text = '  \n  {"answer": "ws", "confidence": 0.5}  \n  '
        result = parse_response(text, SimpleResult)
        assert result.answer == "ws"


# -- Tests: Error handling -------------------------------------------------

class TestErrorHandling:
    def test_malformed_json(self) -> None:
        text = '{"answer": "oops", confidence: bad}'
        with pytest.raises(OutputParseError, match="Failed to parse JSON") as exc_info:
            parse_response(text, SimpleResult)
        assert exc_info.value.raw_text == text
        assert exc_info.value.cause is not None

    def test_valid_json_bad_schema(self) -> None:
        text = '{"wrong_field": "value"}'
        with pytest.raises(OutputParseError, match="failed.*SimpleResult.*validation"):
            parse_response(text, SimpleResult)

    def test_invalid_enum_value(self) -> None:
        text = """{
            "summary": "Bad",
            "findings": [],
            "sentiment": "angry"
        }"""
        with pytest.raises(OutputParseError, match="failed.*NestedResult.*validation"):
            parse_response(text, NestedResult)

    def test_empty_string(self) -> None:
        with pytest.raises(OutputParseError, match="Failed to parse JSON"):
            parse_response("", SimpleResult)

    def test_no_json_at_all(self) -> None:
        with pytest.raises(OutputParseError, match="Failed to parse JSON"):
            parse_response("No JSON here, just plain text.", SimpleResult)

    def test_error_preserves_raw_text(self) -> None:
        raw = "not json at all"
        with pytest.raises(OutputParseError) as exc_info:
            parse_response(raw, SimpleResult)
        assert exc_info.value.raw_text == raw
