"""Dataset registry — register, retrieve, search, and list datasets.

Simple in-memory catalog backed by a dict. Provides tag-based search
and summary statistics.
"""

from __future__ import annotations

import time

from orxaq.fabric.connectors import Dataset

_CATALOG_COUNTER = 0


def _make_id(name: str) -> str:
    """Generate a unique dataset ID from name and counter."""
    global _CATALOG_COUNTER  # noqa: PLW0603
    _CATALOG_COUNTER += 1
    ts = int(time.time() * 1000)
    return f"{name}_{ts}_{_CATALOG_COUNTER}"


class DataCatalog:
    """In-memory dataset registry."""

    def __init__(self) -> None:
        self._datasets: dict[str, Dataset] = {}
        self._tags: dict[str, set[str]] = {}  # dataset_id -> tags
        self._tag_index: dict[str, set[str]] = {}  # tag -> dataset_ids

    def register(self, dataset: Dataset, tags: list[str] | None = None) -> str:
        """Register a dataset and return its ID."""
        dataset_id = _make_id(dataset.source)
        self._datasets[dataset_id] = dataset
        tag_set = set(tags) if tags else set()
        self._tags[dataset_id] = tag_set

        for tag in tag_set:
            if tag not in self._tag_index:
                self._tag_index[tag] = set()
            self._tag_index[tag].add(dataset_id)

        return dataset_id

    def get(self, dataset_id: str) -> Dataset | None:
        """Retrieve a dataset by ID."""
        return self._datasets.get(dataset_id)

    def search(self, tag: str) -> list[str]:
        """Search for datasets by tag. Returns dataset IDs."""
        return sorted(self._tag_index.get(tag, set()))

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

    def summary(self) -> dict:
        """Return catalog summary statistics."""
        total_rows = sum(ds.row_count for ds in self._datasets.values())
        all_tags = set()
        for tag_set in self._tags.values():
            all_tags.update(tag_set)

        return {
            "total_datasets": len(self._datasets),
            "total_rows": total_rows,
            "tags": sorted(all_tags),
        }
