import pandas as pd
import pytest

from wise_pizza_mcp.store import (
    clear_all,
    get_dataset,
    list_datasets,
    remove_dataset,
    store_dataset,
)


def test_store_and_retrieve(sample_df):
    store_dataset("test1", sample_df, ["region", "product"], "total", "size")
    ds = get_dataset("test1")
    assert ds["df"] is sample_df
    assert ds["dims"] == ["region", "product"]
    assert ds["total_name"] == "total"
    assert ds["size_name"] == "size"


def test_missing_dataset_raises():
    with pytest.raises(KeyError, match="not found"):
        get_dataset("nonexistent")


def test_list_datasets(sample_df):
    assert list_datasets() == []
    store_dataset("ds1", sample_df, ["region"], "total")
    store_dataset("ds2", sample_df, ["product"], "total")
    assert sorted(list_datasets()) == ["ds1", "ds2"]


def test_remove_dataset(sample_df):
    store_dataset("ds1", sample_df, ["region"], "total")
    assert "ds1" in list_datasets()
    remove_dataset("ds1")
    assert "ds1" not in list_datasets()


def test_remove_missing_raises():
    with pytest.raises(KeyError, match="not found"):
        remove_dataset("nonexistent")


def test_clear_all(sample_df):
    store_dataset("ds1", sample_df, ["region"], "total")
    store_dataset("ds2", sample_df, ["product"], "total")
    clear_all()
    assert list_datasets() == []
