"""Data registry for versioned datasets and configurations.

The registry maintains a single source of truth (registry.yaml) that tracks
all data artifacts with checksums, lineage, and schema information.

Design principles (Jeff Dean/Demis Hassabis style):
- Single manifest file as source of truth
- Content-addressable via SHA256 checksums
- Fail-fast on integrity violations
- Minimal metadata overhead
- Clear lineage tracking
"""

from __future__ import annotations

import hashlib
import os
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

import yaml


# Module directory for relative path resolution
_MODULE_DIR = Path(__file__).resolve().parent.parent
_REGISTRY_PATH = _MODULE_DIR / "data" / "registry.yaml"


@dataclass(frozen=True)
class SchemaDefinition:
    """Schema definition for validating dataset structure.

    Attributes:
        name: Schema identifier (e.g., "router_data_v2").
        required_columns: Columns that must exist.
        optional_columns: Columns that may exist.
        embedding_dim: Expected embedding dimension (if applicable).
    """

    name: str
    required_columns: tuple[str, ...]
    optional_columns: tuple[str, ...] = field(default_factory=tuple)
    embedding_dim: Optional[int] = None

    def validate_columns(self, columns: list[str]) -> list[str]:
        """Validate columns against schema, returning list of errors."""
        errors = []
        column_set = set(columns)
        for col in self.required_columns:
            if col not in column_set:
                errors.append(f"Missing required column: {col}")
        return errors


@dataclass(frozen=True)
class DatasetVersion:
    """Immutable record of a dataset version.

    Attributes:
        version: Version string (date-based, e.g., "2025.12.14").
        file: Relative path to the data file.
        checksum: SHA256 checksum for integrity verification.
        row_count: Number of rows in the dataset.
        created_at: UTC timestamp when version was created.
        source_config: Config that produced this dataset (lineage).
        notes: Optional human-readable notes about this version.
    """

    version: str
    file: str
    checksum: str
    row_count: int
    created_at: datetime
    source_config: Optional[str] = None
    notes: Optional[str] = None

    @property
    def absolute_path(self) -> Path:
        """Resolve to absolute path relative to embergraphrouter/data/."""
        return _MODULE_DIR / "data" / self.file

    def verify_checksum(self) -> bool:
        """Verify file integrity against stored checksum."""
        if not self.absolute_path.exists():
            return False
        actual = compute_file_checksum(self.absolute_path)
        return actual == self.checksum


@dataclass(frozen=True)
class DatasetInfo:
    """Complete dataset information including all versions.

    Attributes:
        name: Dataset identifier (e.g., "router_data").
        description: Human-readable description.
        current_version: The active/default version string.
        schema_name: Name of the schema definition to use.
        versions: Mapping of version strings to DatasetVersion objects.
    """

    name: str
    description: str
    current_version: str
    schema_name: str
    versions: dict[str, DatasetVersion]

    def get_version(self, version: Optional[str] = None) -> DatasetVersion:
        """Get a specific version, defaulting to current."""
        version = version or self.current_version
        if version not in self.versions:
            available = ", ".join(sorted(self.versions.keys()))
            raise KeyError(
                f"Version '{version}' not found for dataset '{self.name}'. "
                f"Available: {available}"
            )
        return self.versions[version]


@dataclass(frozen=True)
class ConfigVersion:
    """Immutable record of a configuration version.

    Attributes:
        version: Version string (e.g., "v7").
        file: Relative path to the config file.
        checksum: SHA256 checksum for integrity verification.
        created_at: UTC timestamp when version was created.
        notes: Optional human-readable notes about this version.
    """

    version: str
    file: str
    checksum: str
    created_at: datetime
    notes: Optional[str] = None

    @property
    def absolute_path(self) -> Path:
        """Resolve to absolute path relative to embergraphrouter/."""
        return _MODULE_DIR / self.file


@dataclass(frozen=True)
class ConfigInfo:
    """Complete configuration information including all versions."""

    name: str
    description: str
    current_version: str
    versions: dict[str, ConfigVersion]

    def get_version(self, version: Optional[str] = None) -> ConfigVersion:
        """Get a specific version, defaulting to current."""
        version = version or self.current_version
        if version not in self.versions:
            available = ", ".join(sorted(self.versions.keys()))
            raise KeyError(
                f"Version '{version}' not found for config '{self.name}'. "
                f"Available: {available}"
            )
        return self.versions[version]


def compute_file_checksum(path: Path, algorithm: str = "sha256") -> str:
    """Compute checksum of a file.

    Args:
        path: Path to the file.
        algorithm: Hash algorithm (default: sha256).

    Returns:
        Checksum string in format "algorithm:hexdigest".
    """
    hasher = hashlib.new(algorithm)
    with open(path, "rb") as f:
        # Read in 64KB chunks for memory efficiency with large files
        for chunk in iter(lambda: f.read(65536), b""):
            hasher.update(chunk)
    return f"{algorithm}:{hasher.hexdigest()}"


def count_csv_rows(path: Path) -> int:
    """Count rows in a CSV file efficiently (excludes header)."""
    with open(path, "rb") as f:
        return sum(1 for _ in f) - 1  # Subtract header


class DataRegistry:
    """Central registry for versioned data artifacts.

    The registry loads from registry.yaml and provides:
    - Dataset and config discovery
    - Version resolution
    - Schema lookup
    - Integrity verification

    Usage:
        registry = DataRegistry.load()
        dataset = registry.get_dataset_info("router_data")
        version = dataset.get_version("2025.12.14")
        if not version.verify_checksum():
            raise DataIntegrityError("Checksum mismatch")
    """

    def __init__(
        self,
        datasets: dict[str, DatasetInfo],
        configs: dict[str, ConfigInfo],
        schemas: dict[str, SchemaDefinition],
        schema_version: str,
        updated_at: datetime,
    ):
        self._datasets = datasets
        self._configs = configs
        self._schemas = schemas
        self._schema_version = schema_version
        self._updated_at = updated_at

    @classmethod
    def load(cls, path: Optional[Path] = None) -> DataRegistry:
        """Load registry from YAML file.

        Args:
            path: Path to registry.yaml (defaults to standard location).

        Returns:
            Loaded DataRegistry instance.

        Raises:
            FileNotFoundError: If registry.yaml doesn't exist.
            ValueError: If registry format is invalid.
        """
        path = path or _REGISTRY_PATH
        if not path.exists():
            raise FileNotFoundError(
                f"Registry not found at {path}. "
                "Run 'python -m embergraphrouter.versioning.cli init' to create."
            )

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

        return cls._from_dict(data)

    @classmethod
    def _from_dict(cls, data: dict[str, Any]) -> DataRegistry:
        """Parse registry from dictionary."""
        # Parse schemas
        schemas = {}
        for name, schema_data in data.get("schemas", {}).items():
            schemas[name] = SchemaDefinition(
                name=name,
                required_columns=tuple(schema_data.get("required_columns", [])),
                optional_columns=tuple(schema_data.get("optional_columns", [])),
                embedding_dim=schema_data.get("embedding_dim"),
            )

        # Parse datasets
        datasets = {}
        for name, ds_data in data.get("datasets", {}).items():
            versions = {}
            for ver_str, ver_data in ds_data.get("versions", {}).items():
                versions[ver_str] = DatasetVersion(
                    version=ver_str,
                    file=ver_data["file"],
                    checksum=ver_data["checksum"],
                    row_count=ver_data["row_count"],
                    created_at=_parse_datetime(ver_data["created_at"]),
                    source_config=ver_data.get("source_config"),
                    notes=ver_data.get("notes"),
                )
            datasets[name] = DatasetInfo(
                name=name,
                description=ds_data.get("description", ""),
                current_version=ds_data["current_version"],
                schema_name=ds_data.get("schema", ""),
                versions=versions,
            )

        # Parse configs
        configs = {}
        for name, cfg_data in data.get("configs", {}).items():
            versions = {}
            for ver_str, ver_data in cfg_data.get("versions", {}).items():
                versions[ver_str] = ConfigVersion(
                    version=ver_str,
                    file=ver_data["file"],
                    checksum=ver_data["checksum"],
                    created_at=_parse_datetime(ver_data["created_at"]),
                    notes=ver_data.get("notes"),
                )
            configs[name] = ConfigInfo(
                name=name,
                description=cfg_data.get("description", ""),
                current_version=cfg_data["current_version"],
                versions=versions,
            )

        return cls(
            datasets=datasets,
            configs=configs,
            schemas=schemas,
            schema_version=data.get("schema_version", "1.0"),
            updated_at=_parse_datetime(data.get("updated_at", datetime.now(timezone.utc).isoformat())),
        )

    def get_dataset_info(self, name: str) -> DatasetInfo:
        """Get dataset info by name."""
        if name not in self._datasets:
            available = ", ".join(sorted(self._datasets.keys()))
            raise KeyError(f"Dataset '{name}' not found. Available: {available}")
        return self._datasets[name]

    def get_config_info(self, name: str) -> ConfigInfo:
        """Get config info by name."""
        if name not in self._configs:
            available = ", ".join(sorted(self._configs.keys()))
            raise KeyError(f"Config '{name}' not found. Available: {available}")
        return self._configs[name]

    def get_schema(self, name: str) -> SchemaDefinition:
        """Get schema definition by name."""
        if name not in self._schemas:
            available = ", ".join(sorted(self._schemas.keys()))
            raise KeyError(f"Schema '{name}' not found. Available: {available}")
        return self._schemas[name]

    def list_datasets(self) -> list[str]:
        """List all registered dataset names."""
        return sorted(self._datasets.keys())

    def list_configs(self) -> list[str]:
        """List all registered config names."""
        return sorted(self._configs.keys())

    def verify_all(self) -> dict[str, list[str]]:
        """Verify integrity of all registered artifacts.

        Returns:
            Dictionary mapping artifact names to lists of errors (empty if valid).
        """
        errors: dict[str, list[str]] = {}

        for name, dataset in self._datasets.items():
            dataset_errors = []
            current = dataset.get_version()
            if not current.absolute_path.exists():
                dataset_errors.append(f"File not found: {current.absolute_path}")
            elif not current.verify_checksum():
                dataset_errors.append("Checksum mismatch (file modified since registration)")
            if dataset_errors:
                errors[f"dataset:{name}"] = dataset_errors

        for name, config in self._configs.items():
            config_errors = []
            current = config.get_version()
            if not current.absolute_path.exists():
                config_errors.append(f"File not found: {current.absolute_path}")
            if config_errors:
                errors[f"config:{name}"] = config_errors

        return errors

    @property
    def schema_version(self) -> str:
        """Registry schema version."""
        return self._schema_version

    @property
    def updated_at(self) -> datetime:
        """Last registry update timestamp."""
        return self._updated_at


def _parse_datetime(value: str | datetime) -> datetime:
    """Parse datetime from string or pass through datetime."""
    if isinstance(value, datetime):
        return value
    # Handle ISO format with or without timezone
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        # Fallback for simpler formats
        return datetime.strptime(value, "%Y-%m-%d").replace(tzinfo=timezone.utc)
