"""Tests for Pydantic v2 schemas."""
import math
import pytest
from llm_guard.schemas import (
    ChainTrustSchema, DirectGenSchema, TrustResultSchema, MultiTurnResultSchema
)


class TestChainTrustSchema:
    def test_valid_instance(self):
        s = ChainTrustSchema(
            risk_score=0.4, confidence_tier="HIGH", needs_alert=False,
            step_count=3, behavioral_score=0.4,
        )
        assert s.risk_score == 0.4

    def test_invalid_risk_score(self):
        with pytest.raises(Exception):
            ChainTrustSchema(
                risk_score=1.5, confidence_tier="HIGH", needs_alert=False,
                step_count=3, behavioral_score=0.4,
            )

    def test_invalid_confidence_tier(self):
        with pytest.raises(Exception):
            ChainTrustSchema(
                risk_score=0.4, confidence_tier="INVALID", needs_alert=False,
                step_count=3, behavioral_score=0.4,
            )

    def test_json_schema_has_title(self):
        schema = ChainTrustSchema.model_json_schema()
        assert "title" in schema
        assert "properties" in schema

    def test_extra_fields_ignored(self):
        s = ChainTrustSchema(
            risk_score=0.4, confidence_tier="HIGH", needs_alert=False,
            step_count=3, behavioral_score=0.4,
            unknown_field="should be ignored",
        )
        assert not hasattr(s, "unknown_field")


class TestDirectGenSchema:
    def test_nan_ptrue_serializes_as_none(self):
        s = DirectGenSchema(
            risk_score=0.5, confidence_tier="MEDIUM",
            ptrue_score=float("nan"), lexical_risk=0.4,
        )
        assert s.ptrue_score is None

    def test_valid_ptrue(self):
        s = DirectGenSchema(
            risk_score=0.3, confidence_tier="HIGH",
            ptrue_score=0.8, lexical_risk=0.3,
        )
        assert s.ptrue_score == 0.8

    def test_json_roundtrip(self):
        s = DirectGenSchema(
            risk_score=0.3, confidence_tier="HIGH",
            ptrue_score=0.8, lexical_risk=0.3,
        )
        json_str = s.model_dump_json()
        assert "risk_score" in json_str

    def test_nan_sc_risk_serializes_as_none(self):
        s = DirectGenSchema(
            risk_score=0.5, confidence_tier="MEDIUM",
            ptrue_score=0.6, lexical_risk=0.4,
            sc_risk=float("nan"),
        )
        assert s.sc_risk is None

    def test_none_ptrue_stays_none(self):
        s = DirectGenSchema(
            risk_score=0.5, confidence_tier="MEDIUM",
            ptrue_score=None, lexical_risk=0.4,
        )
        assert s.ptrue_score is None


class TestTrustResultSchema:
    def test_valid(self):
        s = TrustResultSchema(
            score=72, verdict="LIKELY_RELIABLE",
            explanation="Good chain.", raw_risk=0.3,
        )
        assert s.score == 72

    def test_invalid_verdict(self):
        with pytest.raises(Exception):
            TrustResultSchema(
                score=72, verdict="GREAT", explanation=".", raw_risk=0.3,
            )

    def test_score_out_of_range(self):
        with pytest.raises(Exception):
            TrustResultSchema(score=150, verdict="LIKELY_RELIABLE",
                              explanation=".", raw_risk=0.3)


class TestSchemaJsonSchemaExport:
    def test_all_schemas_exportable(self):
        for SchemaClass in [ChainTrustSchema, DirectGenSchema,
                             TrustResultSchema, MultiTurnResultSchema]:
            schema = SchemaClass.model_json_schema()
            assert isinstance(schema, dict)
            assert "properties" in schema

    def test_exported_from_llm_guard(self):
        from llm_guard import ChainTrustSchema as C, TrustResultSchema as T
        assert C is ChainTrustSchema
        assert T is TrustResultSchema
