"""Track which datasets are cached locally on this node."""
import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


class DatasetTracker:
    def __init__(self, cache_dir: str = ""):
        self.cache_dir = cache_dir
        self._registry_path = Path(cache_dir or "/tmp") / "mesh_dataset_registry.json"
        self._datasets: Dict[str, dict] = {}
        self._load()

    def _load(self):
        if self._registry_path.exists():
            try:
                self._datasets = json.loads(self._registry_path.read_text())
            except Exception as e:
                logger.warning("Failed to load dataset registry: %s", e)
                self._datasets = {}

    def _save(self):
        try:
            self._registry_path.parent.mkdir(parents=True, exist_ok=True)
            self._registry_path.write_text(json.dumps(self._datasets, indent=2))
        except Exception as e:
            logger.warning("Failed to save dataset registry: %s", e)

    def register(self, dataset_id: str, path: str, size_mb: float = 0.0):
        self._datasets[dataset_id] = {
            "path": path,
            "size_mb": size_mb,
            "registered_at": datetime.now(timezone.utc).isoformat(),
        }
        self._save()
        logger.info("Dataset registered: %s at %s (%.1f MB)", dataset_id, path, size_mb)

    def unregister(self, dataset_id: str):
        if dataset_id in self._datasets:
            del self._datasets[dataset_id]
            self._save()

    def list_ids(self) -> List[str]:
        return list(self._datasets.keys())

    def get(self, dataset_id: str) -> Optional[dict]:
        return self._datasets.get(dataset_id)

    def list_all(self) -> Dict[str, dict]:
        return dict(self._datasets)
