"""Pydantic models for structured LLM output.

Used with Instructor to get type-safe, validated responses
from LLMs for fix generation.
"""

from __future__ import annotations

from enum import Enum
from typing import List

from pydantic import BaseModel, Field, field_validator


class FixStrategy(str, Enum):
    """Available fix strategies."""

    # Data handling
    USE_GET_DEFAULT = "use_get_default"
    ADD_VALIDATION = "add_validation"
    ADD_TYPE_CHECK = "add_type_check"
    ADD_NONE_CHECK = "add_none_check"

    # Logic fixes
    FIX_LOGIC = "fix_logic"
    FIX_COMPARISON = "fix_comparison"
    FIX_CALCULATION = "fix_calculation"
    FIX_INDEX = "fix_index"

    # Error handling
    HANDLE_ERROR = "handle_error"
    ADD_TRY_EXCEPT = "add_try_except"
    ADD_BOUNDS_CHECK = "add_bounds_check"

    # Network/IO
    ADD_RETRY = "add_retry"
    ADD_TIMEOUT = "add_timeout"
    FIX_URL = "fix_url"
    FIX_REQUEST = "fix_request"

    # Auth/Security
    FIX_AUTH = "fix_auth"
    ADD_PERMISSION_CHECK = "add_permission_check"

    # Database
    FIX_QUERY = "fix_query"
    ADD_TRANSACTION = "add_transaction"

    # Imports
    ADD_IMPORT = "add_import"
    FIX_IMPORT = "fix_import"

    # Test fixes
    FIX_ASSERTION = "fix_assertion"
    FIX_MOCK = "fix_mock"
    UPDATE_EXPECTED = "update_expected"

    # Generic
    FIX_THE_CODE = "fix_the_code"
    FIX_THE_TEST = "fix_the_test"


class Analysis(BaseModel):
    """Analysis of the error root cause."""

    root_cause: str = Field(
        description="One sentence describing the actual root cause of the error"
    )
    reasoning: str = Field(
        description="2-3 sentences explaining how you identified the root cause by analyzing the code and error"
    )
    fix_location: str = Field(
        description="The specific location where the fix should be applied (e.g., 'file.py:42' or 'function process_data')"
    )


class Fix(BaseModel):
    """The proposed fix with unified diff."""

    strategy: FixStrategy = Field(
        description="The fix strategy being applied from the available strategies"
    )
    unified_diff: str = Field(
        description="Complete unified diff format patch that can be applied with 'patch -p1'"
    )
    explanation: str = Field(
        description="2-3 sentences explaining what the fix does and why it solves the problem"
    )
    confidence: float = Field(
        ge=0.0,
        le=1.0,
        description="Confidence score from 0.0 to 1.0 based on how certain you are this fix is correct"
    )

    @field_validator("unified_diff")
    @classmethod
    def validate_diff(cls, v: str) -> str:
        """Ensure diff is valid unified diff format."""
        if not v or not v.strip():
            raise ValueError("Diff cannot be empty")

        lines = v.strip().split("\n")

        # Check for required headers
        has_minus = any(l.startswith("---") for l in lines)
        has_plus = any(l.startswith("+++") for l in lines)
        has_hunk = any(l.startswith("@@") for l in lines)

        if not has_minus:
            raise ValueError("Diff must include '--- a/filename' header line")
        if not has_plus:
            raise ValueError("Diff must include '+++ b/filename' header line")
        if not has_hunk:
            raise ValueError("Diff must include '@@ -line,count +line,count @@' hunk header")

        return v


class Verification(BaseModel):
    """Self-verification of the fix."""

    addresses_root_cause: bool = Field(
        description="Does this fix directly address the stated root cause?"
    )
    preserves_valid_behavior: bool = Field(
        description="Does the fix preserve correct behavior for valid inputs?"
    )
    potential_side_effects: List[str] = Field(
        default_factory=list,
        description="List of any potential side effects or edge cases that might be affected"
    )


class HypothesisResponse(BaseModel):
    """Complete hypothesis response from LLM.

    This is the structured output format that Instructor will
    enforce on LLM responses.
    """

    analysis: Analysis = Field(
        description="Analysis of the error root cause"
    )
    fix: Fix = Field(
        description="The proposed fix with unified diff"
    )
    verification: Verification = Field(
        description="Self-verification that the fix is correct"
    )

    def to_hypothesis_dict(self) -> dict:
        """Convert to dictionary format compatible with Hypothesis class."""
        return {
            "root_cause": self.analysis.root_cause,
            "confidence": self.fix.confidence,
            "patch": self.fix.unified_diff,
            "explanation": self.fix.explanation,
            "strategy": self.fix.strategy.value,
            "reasoning": self.analysis.reasoning,
            "fix_location": self.analysis.fix_location,
            "verification": self.verification.model_dump(),
        }


class MultiHypothesisResponse(BaseModel):
    """Multiple hypotheses response for diverse fix generation."""

    hypotheses: List[HypothesisResponse] = Field(
        min_length=1,
        max_length=5,
        description="List of 1-5 fix hypotheses, ordered by confidence"
    )

    def to_hypothesis_dicts(self) -> List[dict]:
        """Convert all hypotheses to dictionary format."""
        return [h.to_hypothesis_dict() for h in self.hypotheses]
