"""Versioned data loading with integrity validation.

Provides the main API for loading datasets and configurations with automatic
checksum verification and schema validation.

Usage:
    from embergraphrouter.versioning import load_dataset, load_config

    # Load current version (auto-selects Parquet if available)
    df = load_dataset("router_data")

    # Load specific version with validation
    df = load_dataset("router_data", version="2025.12.07", validate=True)

    # Force CSV format
    df = load_dataset("router_data", prefer_parquet=False)

    # Load config
    config = load_config("training")
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Optional

import pandas as pd
import yaml

from embergraphrouter.versioning.registry import DataRegistry

logger = logging.getLogger(__name__)


class DataIntegrityError(Exception):
    """Raised when data fails integrity checks."""


class SchemaValidationError(Exception):
    """Raised when data fails schema validation."""


# Mapping from CSV dataset names to their Parquet equivalents
_PARQUET_EQUIVALENTS: dict[str, str] = {
    "router_data": "router_data_parquet",
    "router_data_enriched": "router_data_enriched_parquet",
    "unified_qa_data": "unified_qa_data_parquet",
}


def load_dataset(
    name: str,
    version: Optional[str] = None,
    validate_checksum: bool = True,
    validate_schema: bool = True,
    prefer_parquet: bool = True,
    registry: Optional[DataRegistry] = None,
) -> pd.DataFrame:
    """Load a versioned dataset with optional validation.

    Args:
        name: Dataset name (e.g., "router_data", "unified_qa_data").
        version: Specific version to load (default: current version).
        validate_checksum: Verify file integrity before loading.
        validate_schema: Verify required columns exist.
        prefer_parquet: Auto-select Parquet equivalent if available (10x faster).
        registry: Pre-loaded registry (loads fresh if not provided).

    Returns:
        Loaded DataFrame.

    Raises:
        KeyError: If dataset or version not found in registry.
        DataIntegrityError: If checksum validation fails.
        SchemaValidationError: If schema validation fails.
        FileNotFoundError: If data file doesn't exist.
    """
    registry = registry or DataRegistry.load()

    # Auto-select Parquet if available and preferred
    # BUT: don't auto-select if a specific version is requested (versions may differ)
    actual_name = name
    if prefer_parquet and version is None and name in _PARQUET_EQUIVALENTS:
        parquet_name = _PARQUET_EQUIVALENTS[name]
        try:
            registry.get_dataset_info(parquet_name)
            actual_name = parquet_name
            logger.info("Auto-selected Parquet: %s -> %s", name, parquet_name)
        except KeyError:
            pass  # Parquet not registered, use original

    dataset_info = registry.get_dataset_info(actual_name)
    version_info = dataset_info.get_version(version)

    path = version_info.absolute_path
    if not path.exists():
        raise FileNotFoundError(f"Dataset file not found: {path}")

    # Integrity check
    if validate_checksum:
        logger.debug("Verifying checksum for %s@%s", name, version_info.version)
        if not version_info.verify_checksum():
            raise DataIntegrityError(
                f"Checksum mismatch for {name}@{version_info.version}. "
                "File has been modified since registration. "
                "Re-register with 'versioning register' or restore from backup."
            )

    # Load data
    logger.info("Loading %s@%s from %s", name, version_info.version, path)
    if path.suffix == ".parquet":
        df = pd.read_parquet(path)
    elif path.suffix == ".csv":
        df = pd.read_csv(path)
    else:
        raise ValueError(f"Unsupported file format: {path.suffix}")

    # Schema validation
    if validate_schema and dataset_info.schema_name:
        schema = registry.get_schema(dataset_info.schema_name)
        errors = schema.validate_columns(list(df.columns))
        if errors:
            raise SchemaValidationError(
                f"Schema validation failed for {name}@{version_info.version}: "
                + "; ".join(errors)
            )

    # Row count sanity check
    if len(df) != version_info.row_count:
        logger.warning(
            "Row count mismatch for %s@%s: expected %d, got %d",
            name,
            version_info.version,
            version_info.row_count,
            len(df),
        )

    return df


def load_config(
    name: str,
    version: Optional[str] = None,
    registry: Optional[DataRegistry] = None,
) -> dict[str, Any]:
    """Load a versioned configuration.

    Args:
        name: Config name (e.g., "training").
        version: Specific version to load (default: current version).
        registry: Pre-loaded registry (loads fresh if not provided).

    Returns:
        Loaded configuration dictionary.

    Raises:
        KeyError: If config or version not found in registry.
        FileNotFoundError: If config file doesn't exist.
    """
    registry = registry or DataRegistry.load()
    config_info = registry.get_config_info(name)
    version_info = config_info.get_version(version)

    path = version_info.absolute_path
    if not path.exists():
        raise FileNotFoundError(f"Config file not found: {path}")

    logger.info("Loading config %s@%s from %s", name, version_info.version, path)
    with open(path) as f:
        return yaml.safe_load(f)


def load_dataset_lazy(
    name: str,
    version: Optional[str] = None,
    registry: Optional[DataRegistry] = None,
) -> "LazyDataset":
    """Create a lazy reference to a dataset (does not load immediately).

    Useful when you need dataset metadata without loading the full data.

    Args:
        name: Dataset name.
        version: Specific version (default: current version).
        registry: Pre-loaded registry.

    Returns:
        LazyDataset wrapper that loads on first access.
    """
    return LazyDataset(name, version, registry)


class LazyDataset:
    """Lazy wrapper for versioned datasets.

    Provides metadata access without loading the full dataset. Loads on
    first DataFrame access.

    Attributes:
        name: Dataset name.
        version: Version string.
        path: Absolute path to data file.
        row_count: Expected row count.
        checksum: Expected checksum.
    """

    def __init__(
        self,
        name: str,
        version: Optional[str] = None,
        registry: Optional[DataRegistry] = None,
    ):
        self._registry = registry or DataRegistry.load()
        self._dataset_info = self._registry.get_dataset_info(name)
        self._version_info = self._dataset_info.get_version(version)
        self._df: Optional[pd.DataFrame] = None

    @property
    def name(self) -> str:
        return self._dataset_info.name

    @property
    def version(self) -> str:
        return self._version_info.version

    @property
    def path(self) -> Path:
        return self._version_info.absolute_path

    @property
    def row_count(self) -> int:
        return self._version_info.row_count

    @property
    def checksum(self) -> str:
        return self._version_info.checksum

    @property
    def schema_name(self) -> str:
        return self._dataset_info.schema_name

    def load(self, validate: bool = True) -> pd.DataFrame:
        """Load the dataset (cached after first call)."""
        if self._df is None:
            self._df = load_dataset(
                self.name,
                self.version,
                validate_checksum=validate,
                validate_schema=validate,
                registry=self._registry,
            )
        return self._df

    def __repr__(self) -> str:
        return f"LazyDataset({self.name!r}, version={self.version!r}, rows={self.row_count})"
