import json
import logging
from typing import Any, List, Optional, Union

from pydantic import BaseModel, ConfigDict, Field

logger = logging.getLogger(__name__)


class JITOptions(BaseModel):
    """Options for JIT compilation."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    inlining: Optional[bool] = None
    optimization: Optional[bool] = None
    deforming: Optional[bool] = None


class JITTiming(BaseModel):
    """Timing information for JIT compilation."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    generation: Optional[float] = Field(None, alias="Generation")
    inlining: Optional[float] = Field(None, alias="Inlining")
    optimization: Optional[float] = Field(None, alias="Optimization")
    emission: Optional[float] = Field(None, alias="Emission")
    total: Optional[float] = Field(None, alias="Total")


class JITInfo(BaseModel):
    """Just-In-Time (JIT) compilation information."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    functions: Optional[int] = Field(None, alias="Functions")
    options: Optional[JITOptions] = Field(None, alias="Options")
    timing: Optional[JITTiming] = Field(None, alias="Timing")


class Trigger(BaseModel):
    """Represents a trigger executed during the query."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    trigger_name: Optional[str] = Field(None, alias="Trigger Name")
    relation: Optional[str] = Field(None, alias="Relation")
    time: Optional[float] = Field(None, alias="Time")
    calls: Optional[Union[str, int]] = Field(None, alias="Calls")


class ExplainPlan(BaseModel):
    """Represents a single node in a PostgreSQL EXPLAIN plan."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    node_type: Optional[str] = Field(None, alias="Node Type")
    parallel_aware: bool = Field(False, alias="Parallel Aware")
    relation_name: Optional[str] = Field(None, alias="Relation Name")
    alias: Optional[str] = Field(None, alias="Alias")
    filter: Optional[str] = Field(None, alias="Filter")
    index_cond: Optional[str] = Field(None, alias="Index Cond")
    actual_rows: Optional[int] = Field(None, alias="Actual Rows")
    plan_rows: Optional[int] = Field(None, alias="Plan Rows")
    actual_loops: Optional[int] = Field(None, alias="Actual Loops")
    startup_cost: Optional[float] = Field(None, alias="Startup Cost")
    total_cost: Optional[float] = Field(None, alias="Total Cost")
    plan_width: Optional[int] = Field(None, alias="Plan Width")
    actual_startup_time: Optional[float] = Field(None, alias="Actual Startup Time")
    actual_total_time: Optional[float] = Field(None, alias="Actual Total Time")
    plans: Optional[List["ExplainPlan"]] = Field(None, alias="Plans")
    workers_planned: Optional[int] = Field(None, alias="Workers Planned")
    workers_launched: Optional[int] = Field(None, alias="Workers Launched")
    shared_hit_blocks: Optional[int] = Field(None, alias="Shared Hit Blocks")
    shared_read_blocks: Optional[int] = Field(None, alias="Shared Read Blocks")
    shared_dirtied_blocks: Optional[int] = Field(None, alias="Shared Dirtied Blocks")
    shared_written_blocks: Optional[int] = Field(None, alias="Shared Written Blocks")
    local_hit_blocks: Optional[int] = Field(None, alias="Local Hit Blocks")
    local_read_blocks: Optional[int] = Field(None, alias="Local Read Blocks")
    local_dirtied_blocks: Optional[int] = Field(None, alias="Local Dirtied Blocks")
    local_written_blocks: Optional[int] = Field(None, alias="Local Written Blocks")
    temp_read_blocks: Optional[int] = Field(None, alias="Temp Read Blocks")
    temp_written_blocks: Optional[int] = Field(None, alias="Temp Written Blocks")


class ExplainResult(BaseModel):
    """Represents the entire result of a PostgreSQL EXPLAIN command."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    plan: ExplainPlan = Field(..., alias="Plan")
    planning_time: Optional[float] = Field(None, alias="Planning Time")
    execution_time: Optional[float] = Field(None, alias="Execution Time")
    jit: Optional[JITInfo] = Field(None, alias="JIT")
    triggers: List[Trigger] = Field([], alias="Triggers")


class ExplainParser:
    """Parses PostgreSQL EXPLAIN JSON into a structured object."""

    def parse(self, explain_json: Any) -> Optional[ExplainResult]:
        """
        Parses an EXPLAIN JSON string into a structured ExplainResult object.

        Args:
            explain_json: The EXPLAIN plan in JSON format.

        Returns:
            An ExplainResult object, or None if parsing fails.
        """
        if not explain_json:
            return None

        try:
            if isinstance(explain_json, (dict, list)):
                plan_data = explain_json
            elif isinstance(explain_json, str):
                if not explain_json.strip():
                    return None
                plan_data = json.loads(explain_json)
            else:
                return None
            if isinstance(plan_data, list) and plan_data:
                plan_data = plan_data[0]

            if not isinstance(plan_data, dict) or "Plan" not in plan_data:
                return None

            return ExplainResult.model_validate(plan_data)
        except (json.JSONDecodeError, TypeError) as e:
            logger.warning(f"Failed to parse EXPLAIN JSON: {e}")
            return None
