from typing import Any, Dict, List, Optional

import pandas as pd


_datasets: Dict[str, Dict[str, Any]] = {}


def store_dataset(
    name: str,
    df: pd.DataFrame,
    dims: List[str],
    total_name: str,
    size_name: Optional[str] = None,
) -> None:
    _datasets[name] = {
        "df": df,
        "dims": dims,
        "total_name": total_name,
        "size_name": size_name,
    }


def get_dataset(name: str) -> Dict[str, Any]:
    if name not in _datasets:
        raise KeyError(f"Dataset '{name}' not found. Available: {list(_datasets.keys())}")
    return _datasets[name]


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


def remove_dataset(name: str) -> None:
    if name not in _datasets:
        raise KeyError(f"Dataset '{name}' not found. Available: {list(_datasets.keys())}")
    del _datasets[name]


def clear_all() -> None:
    _datasets.clear()
