import json
import os
import tempfile

import pandas as pd
import pytest

from wise_pizza_mcp.server import (
    explain_changes_in_average,
    explain_changes_in_totals,
    explain_levels,
    explain_timeseries,
    list_loaded_datasets,
    load_data,
    remove_loaded_dataset,
)
from wise_pizza_mcp.store import store_dataset


def test_load_data(tmp_csv):
    result = json.loads(load_data(
        file_path=tmp_csv,
        dataset_name="test",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    ))
    assert result["dataset_name"] == "test"
    assert result["num_rows"] == 90
    assert result["num_columns"] == 5
    assert "region" in result["columns"]
    assert result["dims"] == ["region", "product", "channel"]
    assert result["total_name"] == "total"
    assert result["size_name"] == "size"


def test_load_data_missing_file():
    with pytest.raises(FileNotFoundError):
        load_data(
            file_path="/nonexistent/path.csv",
            dataset_name="test",
            dims=["a"],
            total_name="b",
        )


def test_load_data_missing_column(tmp_csv):
    with pytest.raises(ValueError, match="Columns not found"):
        load_data(
            file_path=tmp_csv,
            dataset_name="test",
            dims=["nonexistent_column"],
            total_name="total",
        )


def test_list_and_remove(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="ds1",
        dims=["region"],
        total_name="total",
    )
    load_data(
        file_path=tmp_csv,
        dataset_name="ds2",
        dims=["region"],
        total_name="total",
    )
    datasets = eval(list_loaded_datasets())
    assert "ds1" in datasets
    assert "ds2" in datasets

    result = remove_loaded_dataset("ds1")
    assert "ds1" in result

    datasets = eval(list_loaded_datasets())
    assert "ds1" not in datasets
    assert "ds2" in datasets


def test_explain_levels(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="test",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_levels(dataset_name="test"))
    assert result["task"] == "levels"
    assert "segments" in result
    assert len(result["segments"]) > 0
    for seg in result["segments"]:
        assert "segment" in seg
        assert "total" in seg
        assert "seg_size" in seg
        assert "naive_avg" in seg


def test_explain_levels_with_image(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="test",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_levels(
        dataset_name="test",
        image_size=[800, 400],
    ))
    assert result["task"] == "levels"
    # image_path should be set (or null if kaleido not available)
    if result["image_path"] is not None:
        assert os.path.exists(result["image_path"])
        assert result["image_path"].endswith(".png")


def test_explain_levels_max_segments(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="test",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_levels(
        dataset_name="test",
        max_segments=3,
    ))
    assert len(result["segments"]) <= 3


def test_explain_changes_in_totals(tmp_csv_pair):
    path1, path2 = tmp_csv_pair
    load_data(
        file_path=path1,
        dataset_name="pre",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    load_data(
        file_path=path2,
        dataset_name="post",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_changes_in_totals(
        dataset_name_1="pre",
        dataset_name_2="post",
    ))
    assert result["task"] == "changes in totals"
    assert "segments" in result
    assert len(result["segments"]) > 0


def test_explain_changes_in_totals_split_fits(tmp_csv_pair):
    path1, path2 = tmp_csv_pair
    load_data(
        file_path=path1,
        dataset_name="pre",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    load_data(
        file_path=path2,
        dataset_name="post",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_changes_in_totals(
        dataset_name_1="pre",
        dataset_name_2="post",
        how="split_fits",
    ))
    assert "size_analysis" in result
    assert "average_analysis" in result
    assert len(result["size_analysis"]["segments"]) > 0
    assert len(result["average_analysis"]["segments"]) > 0


def test_explain_changes_in_average(tmp_csv_pair):
    path1, path2 = tmp_csv_pair
    load_data(
        file_path=path1,
        dataset_name="pre",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    load_data(
        file_path=path2,
        dataset_name="post",
        dims=["region", "product", "channel"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_changes_in_average(
        dataset_name_1="pre",
        dataset_name_2="post",
    ))
    assert result["task"] == "changes in average"
    assert len(result["segments"]) > 0


def test_explain_timeseries(tmp_ts_csv):
    load_data(
        file_path=tmp_ts_csv,
        dataset_name="ts",
        dims=["region", "product"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_timeseries(
        dataset_name="ts",
        time_name="month",
    ))
    assert result["task"] == "time"
    assert len(result["segments"]) > 0


def test_explain_timeseries_with_image(tmp_ts_csv):
    load_data(
        file_path=tmp_ts_csv,
        dataset_name="ts",
        dims=["region", "product"],
        total_name="total",
        size_name="size",
    )
    result = json.loads(explain_timeseries(
        dataset_name="ts",
        time_name="month",
        image_size=[1200, 600],
    ))
    if result["image_path"] is not None:
        assert os.path.exists(result["image_path"])
        assert result["image_path"].endswith(".png")


def test_incompatible_datasets(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="ds1",
        dims=["region", "product"],
        total_name="total",
        size_name="size",
    )
    load_data(
        file_path=tmp_csv,
        dataset_name="ds2",
        dims=["region", "channel"],
        total_name="total",
        size_name="size",
    )
    with pytest.raises(ValueError, match="different dims"):
        explain_changes_in_totals(
            dataset_name_1="ds1",
            dataset_name_2="ds2",
        )


def test_missing_size_name(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="ds1",
        dims=["region"],
        total_name="total",
    )
    load_data(
        file_path=tmp_csv,
        dataset_name="ds2",
        dims=["region"],
        total_name="total",
    )
    with pytest.raises(ValueError, match="size_name"):
        explain_changes_in_totals(
            dataset_name_1="ds1",
            dataset_name_2="ds2",
        )


def test_missing_time_column(tmp_csv):
    load_data(
        file_path=tmp_csv,
        dataset_name="test",
        dims=["region"],
        total_name="total",
        size_name="size",
    )
    with pytest.raises(ValueError, match="Time column"):
        explain_timeseries(
            dataset_name="test",
            time_name="nonexistent_time",
        )
