"""
Pydantic v2 schemas for llm-guard-kit result types.
=====================================================

These mirror the dataclass result types but add validation, JSON serialization,
and JSON Schema export. Use these when building API clients or validating
guard outputs against a schema.

Quick start::

    from llm_guard.pipeline.schemas import ChainTrustSchema, DirectGenSchema, TrustResultSchema
    import json

    # Validate and serialize a ChainTrustResult
    schema = ChainTrustSchema.model_validate(result.__dict__)
    print(schema.model_dump_json())

    # Get JSON Schema for client code generation
    print(json.dumps(ChainTrustSchema.model_json_schema(), indent=2))
Purpose:    Mirrors dataclass result types with Pydantic validation, JSON serialization, and JSON Schema export. Use for API client code generation and guard output validation.
Exports:    ChainTrustSchema, DirectGenSchema, TrustResultSchema
Added:      v0.1.0
Status:     Stable
"""

from __future__ import annotations

import math
from typing import Dict, List, Optional

from pydantic import BaseModel, Field, field_validator


class ChainTrustSchema(BaseModel):
    """
    Pydantic v2 schema for ChainTrustResult.
    Maps to AgentGuard.score_chain() output.
    """
    model_config = {"extra": "ignore"}

    risk_score: float = Field(
        ge=0.0, le=1.0,
        description="Composite risk score [0, 1]. Higher = more likely wrong.",
    )
    confidence_tier: str = Field(
        description="HIGH (risk<0.50) / MEDIUM (0.50<=risk<0.70) / LOW (risk>=0.70).",
        pattern="^(HIGH|MEDIUM|LOW)$",
    )
    needs_alert: bool = Field(
        description="True if risk_score >= alert_threshold."
    )
    failure_mode: Optional[str] = Field(
        default=None,
        description="Detected failure pattern (repeated_query, empty_answer, "
                    "confident_wrong, long_chain) or null.",
    )
    judge_label: Optional[str] = Field(
        default=None,
        description="GOOD / BORDERLINE / POOR / STRUCTURAL / LOCAL / null.",
    )
    step_count: int = Field(
        ge=0,
        description="Number of Search steps in the chain.",
    )
    behavioral_score: float = Field(
        ge=0.0, le=1.0,
        description="SC_OLD behavioral risk before judge blending.",
    )
    behavioral_components: Dict[str, float] = Field(
        default_factory=dict,
        description="Per-signal behavioral component scores.",
    )
    latency_ms: float = Field(
        default=0.0, ge=0.0,
        description="Wall-clock scoring latency in milliseconds.",
    )


class DirectGenSchema(BaseModel):
    """
    Pydantic v2 schema for DirectGenResult.
    Maps to DirectGenerationGuard.score() output.
    """
    model_config = {"extra": "ignore"}

    risk_score: float = Field(ge=0.0, le=1.0,
                               description="Risk score [0,1]. Higher = more likely wrong.")
    confidence_tier: str = Field(pattern="^(HIGH|MEDIUM|LOW)$")
    ptrue_score: Optional[float] = Field(
        default=None,
        description="Raw P(correct) from Haiku judge. null if use_ptrue=False.",
    )
    lexical_risk: float = Field(ge=0.0, le=1.0,
                                 description="Zero-cost lexical risk component.")
    signals: Dict[str, float] = Field(default_factory=dict)
    sc_risk: Optional[float] = Field(
        default=None,
        description="Self-consistency risk. null if score_with_selfconsistency() not called.",
    )

    @field_validator("ptrue_score", "sc_risk", mode="before")
    @classmethod
    def nan_to_none(cls, v):
        """Convert NaN floats to None for clean JSON output."""
        if v is None:
            return None
        try:
            if math.isnan(float(v)):
                return None
        except (TypeError, ValueError):
            pass
        return v


class TrustResultSchema(BaseModel):
    """
    Pydantic v2 schema for TrustResult.
    Maps to TrustScorer.score() output.
    """
    model_config = {"extra": "ignore"}

    score: int = Field(ge=0, le=100,
                        description="Composite trust score 0-100. Higher = more trustworthy.")
    verdict: str = Field(
        pattern="^(HIGHLY_RELIABLE|LIKELY_RELIABLE|UNCERTAIN|LIKELY_UNRELIABLE|UNRELIABLE)$",
        description="Human-readable reliability verdict.",
    )
    explanation: str = Field(description="Natural-language explanation of the trust score.")
    dimension_scores: Dict[str, int] = Field(
        default_factory=dict,
        description="Per-dimension trust scores (0-100 each).",
    )
    raw_risk: float = Field(
        ge=0.0, le=1.0,
        description="Original risk_score from the underlying ChainTrustResult.",
    )


class MultiTurnResultSchema(BaseModel):
    """
    Pydantic v2 schema for MultiTurnResult.
    Maps to MultiTurnGuard.score_turn() output.
    """
    model_config = {"extra": "ignore"}

    risk_score: float = Field(ge=0.0, le=1.0,
                               description="Overall conversation risk score.")
    confidence_tier: str = Field(pattern="^(HIGH|MEDIUM|LOW)$")
    turn_scores: List = Field(default_factory=list,
                               description="Per-turn quality scores.")
    n_turns_scored: int = Field(ge=0)
    trend: str = Field(
        description="Score trend: STABLE / DEGRADING / RECOVERING.",
        pattern="^(STABLE|DEGRADING|RECOVERING)$",
    )
