"""Pydantic models for MCP tool structured outputs."""

from enum import Enum
from typing import List, Optional

from pydantic import BaseModel, Field


class DiagnosticSeverity(str, Enum):
    error = "error"
    warning = "warning"
    info = "info"
    hint = "hint"


class LocalSearchResult(BaseModel):
    name: str = Field(description="Declaration name")
    kind: str = Field(description="Declaration kind (theorem, def, class, etc.)")
    file: str = Field(description="Relative file path")


class LeanSearchResult(BaseModel):
    name: str = Field(description="Full qualified name")
    module_name: str = Field(description="Module where declared")
    kind: Optional[str] = Field(None, description="Declaration kind")
    type: Optional[str] = Field(None, description="Type signature")


class LoogleResult(BaseModel):
    name: str = Field(description="Declaration name")
    type: str = Field(description="Type signature")
    module: str = Field(description="Module where declared")


class LeanFinderResult(BaseModel):
    full_name: str = Field(description="Full qualified name")
    formal_statement: str = Field(description="Lean type signature")
    informal_statement: str = Field(description="Natural language description")


class StateSearchResult(BaseModel):
    name: str = Field(description="Theorem/lemma name")


class PremiseResult(BaseModel):
    name: str = Field(description="Premise name for simp/omega/aesop")


class DiagnosticMessage(BaseModel):
    severity: str = Field(description="error, warning, info, or hint")
    message: str = Field(description="Diagnostic message text")
    line: int = Field(description="Line (1-indexed)")
    column: int = Field(description="Column (1-indexed)")


class GoalState(BaseModel):
    line_context: str = Field(description="Source line where goals were queried")
    goals: Optional[List[str]] = Field(
        None, description="Goal list at specified column position"
    )
    goals_before: Optional[List[str]] = Field(
        None, description="Goals at line start (when column omitted)"
    )
    goals_after: Optional[List[str]] = Field(
        None, description="Goals at line end (when column omitted)"
    )


class CompletionItem(BaseModel):
    label: str = Field(description="Completion text to insert")
    kind: Optional[str] = Field(
        None, description="Completion kind (function, variable, etc.)"
    )
    detail: Optional[str] = Field(None, description="Additional detail")


class HoverInfo(BaseModel):
    symbol: str = Field(description="The symbol being hovered")
    info: str = Field(description="Type signature and documentation")
    diagnostics: List[DiagnosticMessage] = Field(
        default_factory=list, description="Diagnostics at this position"
    )


class TermGoalState(BaseModel):
    line_context: str = Field(description="Source line where term goal was queried")
    expected_type: Optional[str] = Field(
        None, description="Expected type at this position"
    )


class OutlineEntry(BaseModel):
    name: str = Field(description="Declaration name")
    kind: str = Field(description="Declaration kind (Thm, Def, Class, Struct, Ns, Ex)")
    start_line: int = Field(description="Start line (1-indexed)")
    end_line: int = Field(description="End line (1-indexed)")
    type_signature: Optional[str] = Field(
        None, description="Type signature if available"
    )
    children: List["OutlineEntry"] = Field(
        default_factory=list, description="Nested declarations"
    )


class FileOutline(BaseModel):
    imports: List[str] = Field(default_factory=list, description="Import statements")
    declarations: List[OutlineEntry] = Field(
        default_factory=list, description="Top-level declarations"
    )
    total_declarations: Optional[int] = Field(
        None, description="Total count (set when truncated)"
    )


class AttemptResult(BaseModel):
    snippet: str = Field(description="Code snippet that was tried")
    goals: List[str] = Field(
        default_factory=list, description="Goal list after applying snippet"
    )
    diagnostics: List[DiagnosticMessage] = Field(
        default_factory=list, description="Diagnostics for this attempt"
    )


class BuildResult(BaseModel):
    success: bool = Field(description="Whether build succeeded")
    output: str = Field(description="Build output")
    errors: List[str] = Field(default_factory=list, description="Build errors if any")


class RunResult(BaseModel):
    success: bool = Field(description="Whether code compiled successfully")
    diagnostics: List[DiagnosticMessage] = Field(
        default_factory=list, description="Compiler diagnostics"
    )


class DeclarationInfo(BaseModel):
    file_path: str = Field(description="Path to declaration file")
    content: str = Field(description="File content")


# Wrapper models for list-returning tools
# FastMCP flattens bare lists into separate TextContent blocks, causing serialization issues.
# Wrapping in a model ensures proper JSON serialization.


class DiagnosticsResult(BaseModel):
    """Wrapper for diagnostic messages list with build status."""

    success: bool = Field(
        True, description="True if the queried file/range has no errors"
    )
    items: List[DiagnosticMessage] = Field(
        default_factory=list, description="List of diagnostic messages"
    )
    failed_dependencies: List[str] = Field(
        default_factory=list,
        description="File paths of dependencies that failed to build",
    )


class CompletionsResult(BaseModel):
    """Wrapper for completions list."""

    items: List[CompletionItem] = Field(
        default_factory=list, description="List of completion items"
    )


class MultiAttemptResult(BaseModel):
    """Wrapper for multi-attempt results list."""

    items: List[AttemptResult] = Field(
        default_factory=list, description="List of attempt results"
    )


class LocalSearchResults(BaseModel):
    """Wrapper for local search results list."""

    items: List[LocalSearchResult] = Field(
        default_factory=list, description="List of local search results"
    )


class LeanSearchResults(BaseModel):
    """Wrapper for LeanSearch results list."""

    items: List[LeanSearchResult] = Field(
        default_factory=list, description="List of LeanSearch results"
    )


class LoogleResults(BaseModel):
    """Wrapper for Loogle results list."""

    items: List[LoogleResult] = Field(
        default_factory=list, description="List of Loogle results"
    )


class LeanFinderResults(BaseModel):
    """Wrapper for Lean Finder results list."""

    items: List[LeanFinderResult] = Field(
        default_factory=list, description="List of Lean Finder results"
    )


class StateSearchResults(BaseModel):
    """Wrapper for state search results list."""

    items: List[StateSearchResult] = Field(
        default_factory=list, description="List of state search results"
    )


class PremiseResults(BaseModel):
    """Wrapper for premise results list."""

    items: List[PremiseResult] = Field(
        default_factory=list, description="List of premise results"
    )


class WidgetsResult(BaseModel):
    """Wrapper for widget instances at a position."""

    widgets: List[dict] = Field(
        default_factory=list, description="Widget instances (id, name, range, props)"
    )


class InteractiveDiagnosticsResult(BaseModel):
    """Wrapper for interactive diagnostics with embedded widgets."""

    diagnostics: List[dict] = Field(
        default_factory=list,
        description="Interactive diagnostic objects with TaggedText messages",
    )


class WidgetSourceResult(BaseModel):
    """Widget JavaScript source for a given hash."""

    source: dict = Field(description="Widget source data including JavaScript module")


class ReferenceLocation(BaseModel):
    """A single reference location."""

    file_path: str = Field(description="Absolute file path")
    line: int = Field(description="Line (1-indexed)")
    column: int = Field(description="Column (1-indexed)")
    end_line: int = Field(description="End line (1-indexed)")
    end_column: int = Field(description="End column (1-indexed)")


class ReferencesResult(BaseModel):
    """Wrapper for find references results."""

    items: List[ReferenceLocation] = Field(
        default_factory=list, description="List of reference locations"
    )


class LineProfile(BaseModel):
    """Timing for a single source line."""

    line: int = Field(description="Source line number (1-indexed)")
    ms: float = Field(description="Time in milliseconds")
    text: str = Field(description="Source line content (truncated)")


class ProofProfileResult(BaseModel):
    """Profiling result for a theorem."""

    ms: float = Field(description="Total elaboration time in ms")
    lines: List[LineProfile] = Field(
        default_factory=list, description="Time per source line (>1% of total)"
    )
    categories: dict[str, float] = Field(
        default_factory=dict, description="Cumulative time by category in ms"
    )


class CodeActionEdit(BaseModel):
    new_text: str = Field(description="Replacement text")
    start_line: int = Field(description="Start line (1-indexed)")
    start_column: int = Field(description="Start column (1-indexed)")
    end_line: int = Field(description="End line (1-indexed)")
    end_column: int = Field(description="End column (1-indexed)")


class CodeAction(BaseModel):
    title: str = Field(
        description="Code action title (e.g. 'Try this: simp only [...])"
    )
    is_preferred: bool = Field(description="Whether this is the preferred action")
    edits: List[CodeActionEdit] = Field(
        default_factory=list, description="Text edits to apply"
    )


class CodeActionsResult(BaseModel):
    """Wrapper for code actions at a position."""

    actions: List[CodeAction] = Field(
        default_factory=list, description="List of available code actions"
    )


class SourceWarning(BaseModel):
    line: int = Field(description="Line number (1-indexed)")
    pattern: str = Field(description="Matched pattern text")


class VerifyResult(BaseModel):
    axioms: List[str] = Field(
        default_factory=list,
        description="Axioms used. Standard 3: propext, Classical.choice, Quot.sound",
    )
    warnings: List[SourceWarning] = Field(
        default_factory=list,
        description="Suspicious source patterns (if enabled)",
    )
