# System
import datetime as dt
import pathlib as pl
import typing as ty

# Third Party
import numpy as np
import pydantic as pc
import tomlkit as tk

# Internal
from .orientation import AngleUnits, IndexOrder


class RotationStep(pc.BaseModel):
    """Single rotation step.

    Both angles and axes are stored in the current convention/units.
    - Angles: stored in units specified by parent metadata.angle_units
    - Axes: stored in convention specified by parent metadata.index_order
    """

    axes: ty.Literal["X", "Y", "Z"]
    angle: float = 0.0
    visible: bool = True
    name: str = ""
    name_editable: bool = True
    deletable: bool = True

    @pc.model_validator(mode="before")
    @classmethod
    def handle_legacy_format(cls, data):
        """Handle legacy 'angles' list format and 'axis' field."""
        if isinstance(data, dict):
            if "angles" in data and "angle" not in data:
                angles = data["angles"]
                if isinstance(angles, list) and len(angles) > 0:
                    data["angle"] = angles[0]
                else:
                    data["angle"] = angles
            if "axis" in data and "axes" not in data:
                data["axes"] = data["axis"]
        return data


class RotationMetadata(pc.BaseModel):
    """Metadata for TOML files."""

    coordinate_system: ty.Literal["LPS"] = "LPS"
    index_order: IndexOrder = IndexOrder.ITK
    angle_units: AngleUnits = AngleUnits.RADIANS
    timestamp: str = pc.Field(default_factory=lambda: dt.datetime.now().isoformat())
    volume_label: str = ""


class RotationSequence(pc.BaseModel):
    """Complete rotation sequence.

    Both angles and axes are always stored in the current convention/units:
    - Angles: stored in units specified by metadata.angle_units
    - Axes: stored in convention specified by metadata.index_order

    When convention/units change in the UI, all existing rotations are converted.
    """

    model_config = pc.ConfigDict(frozen=False)

    metadata: RotationMetadata = pc.Field(default_factory=RotationMetadata)
    angles_list: list[RotationStep] = pc.Field(default_factory=list)
    mpr_origin: list[float] = pc.Field(
        default_factory=lambda: [0.0, 0.0, 0.0],
        description="MPR origin position [x, y, z] in LPS coordinates",
    )

    @pc.field_validator("mpr_origin")
    @classmethod
    def validate_mpr_origin(cls, v):
        """Ensure mpr_origin is a 3-element list of floats."""
        if not isinstance(v, list) or len(v) != 3:
            raise ValueError("mpr_origin must be a 3-element list [x, y, z]")
        return [float(x) for x in v]

    @pc.field_validator("angles_list", mode="before")
    @classmethod
    def convert_legacy_list(cls, v):
        """Convert legacy list of dicts to list of RotationStep."""
        if isinstance(v, list) and len(v) > 0 and isinstance(v[0], dict):
            return [RotationStep(**item) for item in v]
        return v

    def to_dict_for_ui(self) -> dict:
        """Convert to UI format.

        Angles are passed through in their current units (metadata.angle_units).
        UI expects 'angles' as a list for backward compatibility.
        """
        return {
            "angles_list": [
                {
                    "axes": step.axes,
                    "angles": [step.angle],
                    "visible": step.visible,
                    "name": step.name,
                    "name_editable": step.name_editable,
                    "deletable": step.deletable,
                }
                for step in self.angles_list
            ],
            "mpr_origin": self.mpr_origin,
        }

    @classmethod
    def from_ui_dict(cls, data: dict, volume_label: str = "") -> "RotationSequence":
        """Create from UI format.

        Angles from UI are stored as-is (in current UI units).
        Caller should set metadata.angle_units to match the UI's current units.
        """
        angles_list = [RotationStep(**step) for step in data.get("angles_list", [])]
        metadata = RotationMetadata(volume_label=volume_label)
        mpr_origin = data.get("mpr_origin", [0.0, 0.0, 0.0])
        return cls(metadata=metadata, angles_list=angles_list, mpr_origin=mpr_origin)

    def to_toml(self) -> str:
        """Serialize to TOML using stored serialization preferences."""
        data = self.model_dump(mode="json")
        return tk.dumps(data)

    @classmethod
    def from_toml(cls, toml_content: str) -> "RotationSequence":
        """Deserialize from TOML (no conversions - loads as-is)."""
        doc = tk.loads(toml_content)
        data = dict(doc)
        return cls(**data)

    @classmethod
    def from_file(cls, path: pl.Path) -> "RotationSequence":
        """Load from TOML file."""
        with open(path, "r") as f:
            return cls.from_toml(f.read())

    def to_file(self, path: pl.Path):
        """Save to TOML file using stored serialization preferences."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            f.write(self.to_toml())
