"""Pydantic models for registry validation.

Provides strict typing and validation for registry data structures,
eliminating YAML's type ambiguity issues.

Usage:
    from embergraphrouter.versioning.models import RegistryModel

    # Load and validate
    registry = RegistryModel.from_yaml("registry.yaml")

    # Access with full type safety
    dataset = registry.datasets["router_data"]
    print(dataset.current_version)
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path
from typing import Optional

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator


class SchemaModel(BaseModel):
    """Schema definition with strict validation."""

    model_config = ConfigDict(frozen=True)

    required_columns: list[str] = Field(min_length=1)
    optional_columns: list[str] = Field(default_factory=list)
    embedding_dim: Optional[int] = Field(default=None, ge=1)


class DatasetVersionModel(BaseModel):
    """Single dataset version with validation."""

    model_config = ConfigDict(frozen=True)

    file: str = Field(min_length=1)
    checksum: str = Field(pattern=r"^sha256:[a-f0-9]{64}$")
    row_count: int = Field(ge=0)
    created_at: datetime
    source_config: Optional[str] = None
    notes: Optional[str] = None

    @field_validator("checksum")
    @classmethod
    def validate_checksum_format(cls, v: str) -> str:
        """Ensure checksum is properly formatted SHA256."""
        if not v.startswith("sha256:"):
            raise ValueError("Checksum must start with 'sha256:'")
        hex_part = v[7:]
        if len(hex_part) != 64:
            raise ValueError("SHA256 checksum must be 64 hex characters")
        return v


class DatasetModel(BaseModel):
    """Dataset definition with version history."""

    model_config = ConfigDict(frozen=True, populate_by_name=True)

    description: str = Field(default="")
    current_version: str = Field(min_length=1)
    schema_name: str = Field(
        default="", validation_alias="schema", serialization_alias="schema"
    )
    versions: dict[str, DatasetVersionModel]

    @model_validator(mode="after")
    def validate_current_version_exists(self) -> "DatasetModel":
        """Ensure current_version references an existing version."""
        if self.current_version not in self.versions:
            raise ValueError(
                f"current_version '{self.current_version}' not found in versions"
            )
        return self


class ConfigVersionModel(BaseModel):
    """Single config version with validation."""

    model_config = ConfigDict(frozen=True)

    file: str = Field(min_length=1)
    checksum: str = Field(pattern=r"^sha256:[a-f0-9]{64}$")
    created_at: datetime
    notes: Optional[str] = None


class ConfigModel(BaseModel):
    """Config definition with version history."""

    model_config = ConfigDict(frozen=True)

    description: str = Field(default="")
    current_version: str = Field(min_length=1)
    versions: dict[str, ConfigVersionModel]

    @model_validator(mode="after")
    def validate_current_version_exists(self) -> "ConfigModel":
        """Ensure current_version references an existing version."""
        if self.current_version not in self.versions:
            raise ValueError(
                f"current_version '{self.current_version}' not found in versions"
            )
        return self


class RegistryModel(BaseModel):
    """Root registry model with full validation.

    Validates:
    - Schema version format
    - All datasets reference valid schemas
    - All current_versions exist in versions dict
    - Checksum format correctness
    """

    model_config = ConfigDict(frozen=True)

    schema_version: str = Field(pattern=r"^\d+\.\d+$")
    updated_at: datetime
    schemas: dict[str, SchemaModel] = Field(default_factory=dict)
    datasets: dict[str, DatasetModel] = Field(default_factory=dict)
    configs: dict[str, ConfigModel] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_schema_references(self) -> "RegistryModel":
        """Ensure all dataset schema references exist."""
        for name, dataset in self.datasets.items():
            if dataset.schema_name and dataset.schema_name not in self.schemas:
                raise ValueError(
                    f"Dataset '{name}' references unknown schema '{dataset.schema_name}'"
                )
        return self

    @classmethod
    def from_yaml(cls, path: Path) -> "RegistryModel":
        """Load and validate registry from YAML file."""
        import yaml

        with open(path) as f:
            data = yaml.safe_load(f)
        return cls.model_validate(data)

    @classmethod
    def from_json(cls, path: Path) -> "RegistryModel":
        """Load and validate registry from JSON file."""
        import json

        with open(path) as f:
            data = json.load(f)
        return cls.model_validate(data)

    def to_yaml(self, path: Path) -> None:
        """Write registry to YAML file."""
        import yaml

        with open(path, "w") as f:
            yaml.dump(
                self.model_dump(mode="json", by_alias=True),
                f,
                default_flow_style=False,
                sort_keys=False,
            )

    def to_json(self, path: Path, indent: int = 2) -> None:
        """Write registry to JSON file."""
        with open(path, "w") as f:
            f.write(self.model_dump_json(indent=indent, by_alias=True))
