from pydantic import BaseModel

from srs2std.llms.llm import FinishReason, Response, fallback_parse


class ExampleSchema(BaseModel):
    integer: int
    string: str | None


class NestedSchema(BaseModel):
    boolean: bool
    example: ExampleSchema


def test_fallback_parse_flat():
    text_content = """\
{
    "integer": 1,
    "string": "hello"
}
"""
    expected = Response(
        content=ExampleSchema(integer=1, string="hello"),
        finish_reason=FinishReason.STOP,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected

    text_content = """\
```json
{
    "integer": 1,
    "string": "hello"
}
```
"""
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_flat_following_other_text():
    text_content = """\
The following json represents an example.
These texts should be ignored before parsing.

```json
{
    "integer": 1,
    "string": "hello"
}
```
"""
    expected = Response(
        content=ExampleSchema(integer=1, string="hello"),
        finish_reason=FinishReason.STOP,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_flat_leading_other_text():
    text_content = """\
```
{
    "integer": 1,
    "string": "hello"
}
```

The leading json represents an example.
These texts should be ignored before parsing.
"""
    expected = Response(
        content=ExampleSchema(integer=1, string="hello"),
        finish_reason=FinishReason.STOP,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_flat_with_none():
    text_content = """\
{
    "integer": 1,
    "string": None
}
"""
    expected = Response(
        content=ExampleSchema(integer=1, string=None),
        finish_reason=FinishReason.STOP,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_flat_with_null():
    text_content = """\
{
    "integer": 1,
    "string": null
}
"""
    expected = Response(
        content=ExampleSchema(integer=1, string=None),
        finish_reason=FinishReason.STOP,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_nested():
    text_content = """\
{
    "boolean": True,
    "example": {
        "integer": 1,
        "string": "hello"
    }
}
"""
    expected = Response(
        content=NestedSchema(
            boolean=True, example=ExampleSchema(integer=1, string="hello")
        ),
        finish_reason=FinishReason.STOP,
    )
    actual = fallback_parse(text_content, NestedSchema)
    assert actual == expected


def test_fallback_parse_fail():
    text_content = """\
{
    "integer": 1,
    "string": 1
}
"""
    expected = Response(
        content=None,
        finish_reason=FinishReason.PARSE_FAILED,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_fail_for_not_json():
    text_content = """\
integer: 1
string: "hello"
"""
    expected = Response(
        content=None,
        finish_reason=FinishReason.PARSE_FAILED,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_fail_for_invalid_key():
    text_content = """\
{
    "not_integer": 1,
    "string": "hello"
}
"""
    expected = Response(
        content=None,
        finish_reason=FinishReason.PARSE_FAILED,
    )
    actual = fallback_parse(text_content, ExampleSchema)
    assert actual == expected


def test_fallback_parse_fail_nested_for_invalid_key():
    text_content = """\
{
    "boolean": True,
    "example": {
        "not_integer": 1,
        "string": "hello"
    }
}
"""
    expected = Response(
        content=None,
        finish_reason=FinishReason.PARSE_FAILED,
    )
    actual = fallback_parse(text_content, NestedSchema)
    assert actual == expected

    text_content = """\
{
    "boolean": True,
    "not_example": {
        "integer": 1,
        "string": "hello"
    }
}
"""
    actual = fallback_parse(text_content, NestedSchema)
    assert actual == expected
