from pydantic import BaseModel, Field, RootModel, field_validator
from typing import Dict, List, Optional, Union

class PC(BaseModel):
    p_pc: float
    d_pc: float
    headon: float = 4.9E-5
    overtaking: float = 1.1E-4
    crossing: float = 1.3E-4
    bend: float = 1.3E-4
    grounding: float = 1.6E-4
    allision: float = 1.9E-4

class Rose(RootModel[Dict[str, float]]):
    pass  # keys like "0", "45", etc.

class Repair(BaseModel):
    func: str
    std: float
    loc: float
    scale: float
    use_lognormal: bool

class Drift(BaseModel):
    drift_p: int
    anchor_p: float
    anchor_d: int
    speed: float
    rose: Rose
    repair: Repair

class TrafficDirectionData(BaseModel):
    Frequency_ships_per_year: List[List[float]] = Field(alias="Frequency (ships/year)")
    Speed_knots: List[List[float]] = Field(alias="Speed (knots)")
    Draught_meters: List[List[float]] = Field(alias="Draught (meters)")
    Ship_heights_meters: List[List[float]] = Field(alias="Ship heights (meters)")
    Ship_Beam_meters: List[List[float]] = Field(alias="Ship Beam (meters)")

class TrafficLeg(RootModel[Dict[str, TrafficDirectionData]]):
    """A leg's traffic data, keyed by direction label.

    Every leg has exactly two directions (forward + reverse), but the
    actual labels are derived from the leg's compass bearing in
    ``HandleQGISIface.update_segment_data`` and vary per-leg:
    ``(N going, S going)``, ``(E going, W going)``, ``(S going,
    N going)``, or ``(W going, E going)``.  An earlier schema only
    accepted the ``East going`` / ``West going`` pair as fixed alias
    keys and silently failed to load any project drawn with a
    north-south leg -- the RootModel accepts whatever directions the
    bearing logic emits, while the validator below preserves the
    "exactly two" constraint that the rest of the calculation relies
    on.
    """
    @field_validator('root')
    @classmethod
    def _exactly_two_directions(cls, v):
        if not isinstance(v, dict) or len(v) != 2:
            raise ValueError(
                "Traffic leg must have exactly two direction entries "
                f"(got {len(v) if isinstance(v, dict) else type(v).__name__})"
            )
        return v

class TrafficData(RootModel[Dict[str, TrafficLeg]]):
    pass

class Segment(BaseModel):
    Start_Point: str
    End_Point: str
    Dirs: List[str]
    Width: int
    line_length: float
    Route_Id: int
    Leg_name: str
    Segment_Id: str
    mean1_1: float
    std1_1: float
    mean2_1: float
    std2_1: float
    weight1_1: float
    weight2_1: float
    mean1_2: float
    mean1_3: float
    std1_2: float
    std1_3: float
    mean2_2: float
    mean2_3: float
    std2_2: float
    std2_3: float
    weight1_2: float
    weight1_3: float
    weight2_2: float
    weight2_3: float
    u_min1: float
    u_max1: float
    u_p1: int
    ai1: float
    u_min2: float
    u_max2: float
    u_p2: int
    ai2: float

class SegmentData(RootModel[Dict[str, Segment]]):
    pass

class PolygonEntry(RootModel[List[str]]):
    pass  # ["1", "0.0-3.0", "MultiPolygon (...)"]

class Depths(RootModel[List[PolygonEntry]]):
    pass

class Objects(RootModel[List[PolygonEntry]]):
    pass

class LengthInterval(BaseModel):
    min: Union[float, str]
    max: Union[float, str]
    label: str

class ShipCategoriesModel(BaseModel):
    types: List[str]
    length_intervals: List[LengthInterval]
    selection_mode: Optional[str] = None

class CatastropheLevel(BaseModel):
    name: str
    quantity: float

class ConsequenceModel(BaseModel):
    """Project-level oil-spill consequence inputs.

    All four matrices are sized to match the current ship-category
    dimensions or accident / spill-level taxonomy; loaders reshape them on
    read so older projects still validate after the user adds rows.
    """
    oil_onboard: List[List[float]] = Field(default_factory=list)
    spill_probability: List[List[float]] = Field(default_factory=list)
    spill_fraction: List[List[float]] = Field(default_factory=list)
    catastrophe_levels: List[CatastropheLevel] = Field(default_factory=list)


class JunctionModel(BaseModel):
    """A single junction (place where two or more legs meet).

    ``transitions[in_leg][out_leg]`` stores the share of traffic that,
    after arriving on ``in_leg``, continues onto ``out_leg``.  Rows
    must sum to ~1.0 (validated separately via
    :func:`geometries.junctions.validate_junctions`).
    """
    point: List[float] = Field(default_factory=list)  # [lon, lat]
    legs: Dict[str, str] = Field(default_factory=dict)  # leg_id -> "start"|"end"
    transitions: Dict[str, Dict[str, float]] = Field(default_factory=dict)
    source: str = "geometry"  # "geometry" | "ais" | "user"


class JunctionsModel(RootModel[Dict[str, JunctionModel]]):
    pass


class ImportedSegmentEndpoints(BaseModel):
    """Write-once snapshot of a leg's geometry as it was first created
    (IWRAP import / first canvas digitize).  Diffed against the live
    ``segment_data`` by the audit report to expose silent waypoint moves.
    """
    Start_Point: str
    End_Point: str


class SegmentsImported(RootModel[Dict[str, ImportedSegmentEndpoints]]):
    pass


class RootModelSchema(BaseModel):
    pc: PC
    drift: Drift
    traffic_data: TrafficData
    segment_data: SegmentData
    depths: Depths
    objects: Objects
    ship_categories: Optional[ShipCategoriesModel] = None
    consequence: Optional[ConsequenceModel] = None
    junctions: Optional[JunctionsModel] = None
    segments_imported: Optional[SegmentsImported] = None