"""Dataset registry for RDF knowledge graphs.

This module provides a comprehensive registry of public RDF datasets
that can be downloaded and used for GISM training. The registry is loaded
from a JSON configuration file for easy maintenance.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass


@dataclass
class DatasetInfo:
    """Information about a downloadable RDF dataset."""

    id: str
    name: str
    description: str
    url: str
    format: str
    size_gb: float
    compressed_size_gb: float | None
    entities: str | None
    triples: str | None
    category: str
    license: str | None
    recommended_for: str
    extraction_cmd: str | None = None
    notes: str | None = None
    available: bool = True  # Flag to mark if dataset is currently accessible
    huggingface_repo: str | None = None  # HuggingFace repository name
    short_name: str | None = None  # Short name for the dataset


def load_dataset_config() -> dict:
    """Load dataset configuration from JSON file.

    Returns:
        Dictionary containing dataset configuration
    """
    config_path = Path(__file__).parent / "dataset_config.json"
    if not config_path.exists():
        raise FileNotFoundError(f"Dataset configuration file not found: {config_path}")

    with open(config_path) as f:
        return json.load(f)


def build_dataset_registry() -> dict[str, DatasetInfo]:
    """Build dataset registry from JSON configuration.

    Returns:
        Dictionary mapping dataset IDs to DatasetInfo objects
    """
    config = load_dataset_config()
    registry = {}

    for dataset_id, dataset_config in config["datasets"].items():
        # Extract download and metadata information
        download = dataset_config["download"]
        metadata = dataset_config["metadata"]
        huggingface = dataset_config.get("huggingface", {})

        # Handle unavailable datasets
        available = True
        if metadata.get("notes") and any(
            phrase in metadata["notes"].lower()
            for phrase in [
                "requires registration",
                "multi-file ftp",
                "hdt format",
                "multiple files on figshare",
                "databus collection",
            ]
        ):
            available = False

        # Special cases for specific datasets
        if dataset_id in [
            "babelnet",
            "pubchem",
            "bio2rdf",
            "lod-a-lot",
            "opencitations",
            "dbpedia-latest",
        ]:
            available = False
        elif dataset_id in ["conceptnet", "framenet", "opencyc"]:
            available = True

        registry[dataset_id] = DatasetInfo(
            id=dataset_id,
            name=dataset_config["full_name"],
            description=dataset_config["description"],
            url=download["url"],
            format=download["format"],
            size_gb=download["size_gb"],
            compressed_size_gb=download.get("compressed_size_gb"),
            entities=metadata.get("entities"),
            triples=metadata.get("triples"),
            category=dataset_config["category"],
            license=metadata.get("license"),
            recommended_for=metadata.get("recommended_for", ""),
            extraction_cmd=download.get("extraction_command"),
            notes=metadata.get("notes"),
            available=available,
            huggingface_repo=huggingface.get("repository"),
            short_name=dataset_config.get("short_name"),
        )

    return registry


# Build the registry on module import
try:
    DATASET_REGISTRY = build_dataset_registry()
except FileNotFoundError:
    # Fallback to empty registry if config file is missing
    print("Warning: dataset_config.json not found, using empty registry")
    DATASET_REGISTRY = {}


def get_dataset_info(dataset_id: str) -> DatasetInfo | None:
    """Get information about a dataset by ID.

    Args:
        dataset_id: Dataset identifier

    Returns:
        DatasetInfo object or None if not found
    """
    return DATASET_REGISTRY.get(dataset_id)


def list_datasets(category: str | None = None) -> list[DatasetInfo]:
    """List all available datasets, optionally filtered by category.

    Args:
        category: Filter by category (small, medium, large, xlarge)

    Returns:
        List of DatasetInfo objects
    """
    datasets = list(DATASET_REGISTRY.values())
    if category:
        datasets = [d for d in datasets if d.category == category]
    return datasets


def get_dataset_categories() -> list[str]:
    """Get list of all dataset categories.

    Returns:
        List of category names
    """
    return sorted({d.category for d in DATASET_REGISTRY.values()})


def get_dataset_config() -> dict:
    """Get the full dataset configuration.

    Returns:
        Dictionary containing the full dataset configuration
    """
    return load_dataset_config()