"""Pydantic models for storage and agent responses."""

from datetime import datetime
from typing import Any, Literal
from pydantic import BaseModel, Field


# ============================================================================
# Storage Models (for database persistence)
# ============================================================================


class DatabaseInfo(BaseModel):
    """Information about a registered database."""

    id: int | None = None
    name: str = Field(description="User-friendly name for the database")
    path: str = Field(description="Absolute path to the SQLite database file")
    database_context: dict[str, Any] | None = Field(
        default=None, description="Persisted holistic database context from SchemaContextAgent"
    )
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    last_synced_at: str = Field(default_factory=lambda: datetime.now().isoformat())


class SchemaInfo(BaseModel):
    """Schema metadata for a table."""

    id: int | None = None
    database_id: int = Field(description="Foreign key to databases table")
    table_name: str = Field(description="Name of the table")
    original_metadata: dict[str, Any] = Field(
        description="Raw metadata from dq_db_manager"
    )
    enriched_metadata: dict[str, Any] | None = Field(
        default=None, description="LLM-enriched schema description"
    )
    sample_data: list[dict[str, Any]] = Field(
        default_factory=list, description="Sample rows from the table"
    )
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())


class SettingInfo(BaseModel):
    """Application setting."""

    key: str = Field(description="Setting key")
    value: str = Field(description="Setting value")


# ============================================================================
# Agent Response Models (for Pydantic AI structured outputs)
# ============================================================================


class ColumnDescription(BaseModel):
    """Description of a single column."""

    column_name: str = Field(description="Exact column name from the table")
    description: str = Field(description="Human-readable description of what this column stores")


class EnrichedSchema(BaseModel):
    """Enriched schema description generated by LLM."""

    table_name: str = Field(description="Name of the table")
    table_purpose: str = Field(
        description="High-level purpose of this table (what it stores, its role)"
    )
    column_descriptions: list[ColumnDescription] = Field(
        description="Description for each column in the table"
    )
    relationships: list[str] = Field(
        default_factory=list,
        description="Relationships to other tables (foreign keys, logical connections)",
    )
    business_context: str | None = Field(
        default=None,
        description="Additional business context inferred from sample data",
    )

    def to_storage_dict(self) -> dict:
        """Convert to storage format (column_descriptions as dict)."""
        d = self.model_dump()
        d["column_descriptions"] = {c.column_name: c.description for c in self.column_descriptions}
        return d


class ConversationalResponse(BaseModel):
    """Simple conversational message response."""

    message: str = Field(description="Natural language response to user")


class QuestionAnalysis(BaseModel):
    """Analysis of user's question."""

    intent: str = Field(
        description="What the user wants to know (summarize in 1-2 sentences)"
    )
    relevant_tables: list[str] = Field(
        description="List of table names relevant to answering this question"
    )
    requires_join: bool = Field(
        description="Whether the query requires joining multiple tables"
    )
    complexity: str = Field(
        description="Query complexity: 'simple', 'medium', or 'complex'"
    )


class QueryPlan(BaseModel):
    """Generated SQL query plan."""

    sql_query: str = Field(description="SQL query to execute")
    explanation: str = Field(description="Explanation of what the query does")
    tables_used: list[str] = Field(description="List of tables involved in the query")
    estimated_rows: int | None = Field(
        default=None,
        description="Estimated number of rows the query will return (if known)",
    )


class SuccessAnswer(BaseModel):
    """Successful query result with answer."""

    success: bool = Field(default=True, description="Always True for success")
    answer: str = Field(
        description="Natural language answer based on query results"
    )
    confidence: float = Field(
        ge=0.0,
        le=1.0,
        description="Confidence in the answer (0.0 to 1.0)",
    )


class RetryWithFeedback(BaseModel):
    """Indicates query needs to be retried with modifications."""

    success: bool = Field(default=False, description="Always False for retry")
    retry_reason: str = Field(
        description="Why the query failed or produced incorrect results"
    )
    suggested_fix: str = Field(
        description="Suggestion for how to modify the query"
    )


# ============================================================================
# Query Planning Models
# ============================================================================


class QueryComplexity(BaseModel):
    """Result of complexity classification."""
    is_complex: bool = Field(description="True if question needs multi-step decomposition")
    reasoning: str = Field(description="Why this is simple or complex")
    estimated_steps: int = Field(
        default=1,
        description="Estimated number of sub-problems needed"
    )


class SubProblem(BaseModel):
    """A single sub-problem in a query execution plan."""
    id: str = Field(description="Unique step ID, e.g. 'step_1'")
    title: str = Field(description="Short title, e.g. 'Explore sales table structure'")
    purpose: str = Field(
        description="What this step establishes, e.g. 'Verify column names and join keys for store_sales'"
    )
    approach: str = Field(
        description="SQL approach, e.g. 'SELECT column names from pragma_table_info or sample rows'"
    )
    tables_needed: list[str] = Field(description="Tables this sub-problem queries")
    depends_on: list[str] = Field(
        default_factory=list,
        description="IDs of sub-problems whose results are needed first"
    )


class ExecutionPlan(BaseModel):
    """Multi-step execution plan for a complex question."""
    question: str = Field(description="Original user question")
    approach_summary: str = Field(
        description="High-level approach in 1-2 sentences"
    )
    sub_problems: list[SubProblem] = Field(
        description="Ordered list of exploratory sub-problems"
    )
    final_assembly_strategy: str = Field(
        description="How sub-problem results will inform the final query (CTE strategy, join order, etc.)"
    )


class SubProblemResult(BaseModel):
    """Result of executing a single sub-problem."""
    sub_problem_id: str
    sql_query: str = ""
    success: bool = True
    data_summary: str = Field(
        default="",
        description="Natural language summary of what was found"
    )
    row_count: int = 0
    key_findings: list[str] = Field(
        default_factory=list,
        description="Specific verified facts from query results"
    )
    verified_columns: dict[str, str] = Field(
        default_factory=dict,
        description="Map of {concept: real_column_name} verified in this step, e.g. {'revenue': 'ss_net_profit'}"
    )
    verified_joins: list[str] = Field(
        default_factory=list,
        description="Successful JOIN clauses verified, e.g. 'store_sales.ss_item_sk = item.i_item_sk'"
    )
    sample_data: list[dict[str, Any]] | None = None
    error: str | None = None


class PlanRevision(BaseModel):
    """Revised plan after a sub-problem failure or unexpected result."""
    should_revise: bool = Field(description="Whether the plan needs changes")
    reason: str = Field(description="Why the plan is being revised")
    updated_sub_problems: list[SubProblem] = Field(
        default_factory=list,
        description="New/modified sub-problems to replace remaining ones"
    )
    drop_ids: list[str] = Field(
        default_factory=list,
        description="IDs of sub-problems to remove from the plan"
    )


# ============================================================================
# Data Profiling Models
# ============================================================================


class ColumnProfile(BaseModel):
    """Statistical profile for a single column."""

    column_name: str = Field(description="Name of the column")
    data_type: str = Field(description="SQL data type")
    row_count: int = Field(description="Total number of rows")
    null_count: int = Field(description="Number of NULL values")
    null_percentage: float = Field(description="Percentage of NULL values")
    distinct_count: int = Field(description="Number of distinct values")
    distinct_percentage: float = Field(description="Percentage of distinct values")

    # Numeric statistics (None if not numeric)
    min_value: float | None = None
    max_value: float | None = None
    mean: float | None = None
    median: float | None = None
    std_dev: float | None = None
    q1: float | None = None
    q3: float | None = None
    skewness: float | None = None
    kurtosis: float | None = None

    # String statistics (None if not string)
    min_length: int | None = None
    max_length: int | None = None
    avg_length: float | None = None

    # Value distribution (top 10 most common values)
    top_values: list[tuple[Any, int]] = Field(
        default_factory=list, description="[(value, count), ...] top 10"
    )

    # Semantic role classification
    inferred_role: str | None = Field(
        default=None,
        description="Inferred semantic role: id, name, email, date, category, etc.",
    )
    role_confidence: float | None = Field(
        default=None, description="Confidence in role classification (0.0-1.0)"
    )


class TableProfile(BaseModel):
    """Complete profile for a table."""

    id: int | None = None
    database_id: int = Field(description="Foreign key to databases table")
    table_name: str = Field(description="Name of the table")
    row_count: int = Field(description="Total number of rows")
    column_count: int = Field(description="Total number of columns")
    columns: list[ColumnProfile] = Field(description="Column-level profiles")
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())


# ============================================================================
# Data Quality Models
# ============================================================================


class QualityIssue(BaseModel):
    """A single data quality issue detected."""

    severity: Literal["critical", "warning", "info"] = Field(
        description="Severity level"
    )
    category: str = Field(
        description="Issue category: completeness, uniqueness, validity, consistency"
    )
    column_name: str | None = Field(default=None, description="Affected column")
    issue: str = Field(description="Description of the issue")
    affected_rows: int | None = Field(
        default=None, description="Number of rows affected"
    )
    recommendation: str = Field(description="Suggested action to fix")


class QualityReport(BaseModel):
    """Data quality assessment report for a table."""

    id: int | None = None
    database_id: int = Field(description="Foreign key to databases table")
    table_name: str = Field(description="Name of the table")
    overall_score: float = Field(
        ge=0.0, le=100.0, description="Overall quality score (0-100)"
    )
    completeness_score: float = Field(
        ge=0.0, le=100.0, description="Completeness score"
    )
    uniqueness_score: float = Field(ge=0.0, le=100.0, description="Uniqueness score")
    validity_score: float = Field(ge=0.0, le=100.0, description="Validity score")
    consistency_score: float = Field(ge=0.0, le=100.0, description="Consistency score")
    issues: list[QualityIssue] = Field(
        default_factory=list, description="List of detected issues"
    )
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())


# ============================================================================
# Insights Models
# ============================================================================


class Insight(BaseModel):
    """A single actionable insight about the data."""

    category: str = Field(
        description="Insight category: outlier, correlation, distribution, trend, etc."
    )
    title: str = Field(description="Short title for the insight")
    description: str = Field(description="Detailed explanation generated by LLM")
    severity: Literal["high", "medium", "low"] = Field(description="Importance level")
    affected_columns: list[str] = Field(
        default_factory=list, description="Columns involved"
    )
    supporting_data: dict[str, Any] = Field(
        default_factory=dict, description="Evidence/statistics supporting the insight"
    )


class InsightsReport(BaseModel):
    """Collection of insights for a table."""

    id: int | None = None
    database_id: int = Field(description="Foreign key to databases table")
    table_name: str = Field(description="Name of the table")
    insights: list[Insight] = Field(description="List of generated insights")
    summary: str = Field(description="Natural language summary of key findings")
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())


# ============================================================================
# Token Usage Tracking
# ============================================================================


class UsageStats(BaseModel):
    """Accumulated token usage across agent calls for a single question."""
    input_tokens: int = 0
    output_tokens: int = 0
    requests: int = 0

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens

    def add(self, usage) -> None:
        """Add usage from a pydantic_ai RunUsage object."""
        if usage:
            self.input_tokens += getattr(usage, 'input_tokens', 0) or 0
            self.output_tokens += getattr(usage, 'output_tokens', 0) or 0
            self.requests += getattr(usage, 'requests', 0) or 0
