"""Pydantic schemas for the AI auto-fix feature."""
from typing import Any, Literal

from pydantic import BaseModel, Field


# ── Requests ──


class ValidationErrorInput(BaseModel):
    """A single validation error to fix."""

    node_id: str = Field(..., description="ID of the node with the error")
    message: str = Field(..., description="The validation error message")


class AIFixTreeRequest(BaseModel):
    """Request to generate AI fixes for validation errors."""

    tree_structure: dict[str, Any] = Field(..., description="Full tree structure")
    tree_name: str = Field("", max_length=255, description="Name of the flow")
    tree_type: Literal["troubleshooting", "procedural", "maintenance"] = Field(
        "troubleshooting", description="Type of flow"
    )
    validation_errors: list[ValidationErrorInput] = Field(
        ..., min_length=1, max_length=10, description="Errors to fix"
    )


# ── Responses ──


class AIFixProposal(BaseModel):
    """A single proposed fix from the AI."""

    target_node_id: str
    error_message: str
    description: str
    original_node: dict[str, Any]
    fixed_node: dict[str, Any]


class AIFixTokenUsage(BaseModel):
    input: int = 0
    output: int = 0


class AIFixTreeResponse(BaseModel):
    """Response with proposed fixes."""

    fixes: list[AIFixProposal]
    tokens_used: AIFixTokenUsage