import tempfile
from typing import List, Optional

import pandas as pd
import plotly.io as pio
import wise_pizza
from mcp.server.fastmcp import FastMCP

from wise_pizza_mcp.extract import extract_result, result_to_json
from wise_pizza_mcp.models import LoadDataResult
from wise_pizza_mcp.store import (
    clear_all,
    get_dataset,
    list_datasets,
    remove_dataset,
    store_dataset,
)

mcp = FastMCP("wise-pizza")


def _generate_image(result: object, image_size: List[int]) -> Optional[str]:
    """Generate a PNG from a SliceFinder/SlicerPair plot and return the file path."""
    try:
        fig = result.plot(return_fig=True)
        # plot() may return a list [main_fig, cluster_table] if clusters exist
        if isinstance(fig, list):
            fig = fig[0]
        fig.update_layout(width=image_size[0], height=image_size[1])
        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        pio.write_image(fig, tmp.name, format="png", scale=2)
        return tmp.name
    except Exception:
        return None


@mcp.tool()
def load_data(
    file_path: str,
    dataset_name: str,
    dims: List[str],
    total_name: str,
    size_name: Optional[str] = None,
) -> str:
    """Load a CSV file into the in-memory dataset store.

    Args:
        file_path: Path to a CSV file on disk
        dataset_name: A unique name to refer to this dataset later
        dims: List of column names to use as dimensions (categorical columns to segment by)
        total_name: Column name containing the totals (the value to analyze)
        size_name: Optional column name containing segment sizes/weights
    """
    df = pd.read_csv(file_path)

    missing = [c for c in dims + [total_name] if c not in df.columns]
    if size_name is not None and size_name not in df.columns:
        missing.append(size_name)
    if missing:
        raise ValueError(f"Columns not found in CSV: {missing}. Available: {list(df.columns)}")

    store_dataset(dataset_name, df, dims, total_name, size_name)

    result = LoadDataResult(
        dataset_name=dataset_name,
        num_rows=len(df),
        num_columns=len(df.columns),
        columns=list(df.columns),
        dims=dims,
        total_name=total_name,
        size_name=size_name,
    )
    return result.model_dump_json()


@mcp.tool()
def list_loaded_datasets() -> str:
    """List all datasets currently loaded in the store."""
    return str(list_datasets())


@mcp.tool()
def remove_loaded_dataset(dataset_name: str) -> str:
    """Remove a dataset from the store.

    Args:
        dataset_name: Name of the dataset to remove
    """
    remove_dataset(dataset_name)
    return f"Dataset '{dataset_name}' removed."


@mcp.tool()
def explain_levels(
    dataset_name: str,
    max_segments: Optional[int] = None,
    min_depth: int = 1,
    max_depth: int = 2,
    constrain_signs: bool = True,
    cluster_values: bool = False,
    image_size: Optional[List[int]] = None,
) -> str:
    """Find segments whose average is most different from the global one.

    This identifies the most unusual slices in a single dataset — segments where the
    average value deviates most from the overall average.

    Args:
        dataset_name: Name of a previously loaded dataset
        max_segments: Maximum number of segments to find (default: auto)
        min_depth: Minimum number of dimensions to constrain per segment (default: 1)
        max_depth: Maximum number of dimensions to constrain per segment (default: 2)
        constrain_signs: Force segment weights to match their direction (default: True)
        cluster_values: Also consider clusters of similar dimension values (default: False)
        image_size: Optional [width, height] in pixels to generate a PNG chart
    """
    ds = get_dataset(dataset_name)
    df, dims, total_name, size_name = ds["df"], ds["dims"], ds["total_name"], ds["size_name"]

    result = wise_pizza.explain_levels(
        df=df,
        dims=dims,
        total_name=total_name,
        size_name=size_name,
        max_segments=max_segments,
        min_depth=min_depth,
        max_depth=max_depth,
        solver="lasso",
        constrain_signs=constrain_signs,
        cluster_values=cluster_values,
    )

    extracted = extract_result(result)

    if image_size is not None:
        extracted.image_path = _generate_image(result, image_size)

    return result_to_json(extracted)


@mcp.tool()
def explain_changes_in_totals(
    dataset_name_1: str,
    dataset_name_2: str,
    max_segments: Optional[int] = None,
    min_depth: int = 1,
    max_depth: int = 2,
    how: str = "totals",
    constrain_signs: bool = True,
    cluster_values: bool = False,
    image_size: Optional[List[int]] = None,
) -> str:
    """Find segments most useful in explaining differences between totals of two datasets.

    Compares two time periods or cohorts and identifies which segments drove the change
    in total values.

    Args:
        dataset_name_1: Name of the first (baseline) dataset
        dataset_name_2: Name of the second (comparison) dataset
        max_segments: Maximum number of segments to find (default: auto)
        min_depth: Minimum number of dimensions to constrain per segment (default: 1)
        max_depth: Maximum number of dimensions to constrain per segment (default: 2)
        how: Decomposition method. "totals" for combined analysis, "split_fits" to
             separately analyze size vs average changes, "extra_dim" or "force_dim"
             for alternative decompositions (default: "totals")
        constrain_signs: Force segment weights to match their direction (default: True)
        cluster_values: Also consider clusters of similar dimension values (default: False)
        image_size: Optional [width, height] in pixels to generate a PNG chart
    """
    ds1 = get_dataset(dataset_name_1)
    ds2 = get_dataset(dataset_name_2)

    if ds1["dims"] != ds2["dims"]:
        raise ValueError(
            f"Datasets have different dims: {ds1['dims']} vs {ds2['dims']}"
        )
    if ds1["total_name"] != ds2["total_name"]:
        raise ValueError(
            f"Datasets have different total_name: {ds1['total_name']} vs {ds2['total_name']}"
        )
    if ds1["size_name"] is None or ds2["size_name"] is None:
        raise ValueError("Both datasets must have size_name set for changes analysis")

    result = wise_pizza.explain_changes_in_totals(
        df1=ds1["df"],
        df2=ds2["df"],
        dims=ds1["dims"],
        total_name=ds1["total_name"],
        size_name=ds1["size_name"],
        max_segments=max_segments,
        min_depth=min_depth,
        max_depth=max_depth,
        solver="lasso",
        how=how,
        constrain_signs=constrain_signs,
        cluster_values=cluster_values,
    )

    extracted = extract_result(result)

    if image_size is not None:
        image_path = _generate_image(result, image_size)
        extracted.image_path = image_path

    return result_to_json(extracted)


@mcp.tool()
def explain_changes_in_average(
    dataset_name_1: str,
    dataset_name_2: str,
    max_segments: Optional[int] = None,
    min_depth: int = 1,
    max_depth: int = 2,
    how: str = "totals",
    constrain_signs: bool = True,
    cluster_values: bool = False,
    image_size: Optional[List[int]] = None,
) -> str:
    """Find segments most useful in explaining differences between averages of two datasets.

    Similar to explain_changes_in_totals but focuses on average value changes rather
    than total volume changes.

    Args:
        dataset_name_1: Name of the first (baseline) dataset
        dataset_name_2: Name of the second (comparison) dataset
        max_segments: Maximum number of segments to find (default: auto)
        min_depth: Minimum number of dimensions to constrain per segment (default: 1)
        max_depth: Maximum number of dimensions to constrain per segment (default: 2)
        how: Decomposition method. "totals" for combined analysis, "split_fits" to
             separately analyze size vs average changes (default: "totals")
        constrain_signs: Force segment weights to match their direction (default: True)
        cluster_values: Also consider clusters of similar dimension values (default: False)
        image_size: Optional [width, height] in pixels to generate a PNG chart
    """
    ds1 = get_dataset(dataset_name_1)
    ds2 = get_dataset(dataset_name_2)

    if ds1["dims"] != ds2["dims"]:
        raise ValueError(
            f"Datasets have different dims: {ds1['dims']} vs {ds2['dims']}"
        )
    if ds1["total_name"] != ds2["total_name"]:
        raise ValueError(
            f"Datasets have different total_name: {ds1['total_name']} vs {ds2['total_name']}"
        )
    if ds1["size_name"] is None or ds2["size_name"] is None:
        raise ValueError("Both datasets must have size_name set for changes analysis")

    result = wise_pizza.explain_changes_in_average(
        df1=ds1["df"],
        df2=ds2["df"],
        dims=ds1["dims"],
        total_name=ds1["total_name"],
        size_name=ds1["size_name"],
        max_segments=max_segments,
        min_depth=min_depth,
        max_depth=max_depth,
        solver="lasso",
        how=how,
        constrain_signs=constrain_signs,
        cluster_values=cluster_values,
    )

    extracted = extract_result(result)

    if image_size is not None:
        image_path = _generate_image(result, image_size)
        extracted.image_path = image_path

    return result_to_json(extracted)


@mcp.tool()
def explain_timeseries(
    dataset_name: str,
    time_name: str,
    num_segments: Optional[int] = None,
    max_depth: int = 2,
    fit_sizes: Optional[bool] = None,
    num_breaks: int = 3,
    n_jobs: int = 10,
    ignore_averages: bool = True,
    image_size: Optional[List[int]] = None,
) -> str:
    """Split a time series panel dataset into segments with different temporal patterns.

    Finds groups of dimension values that behave differently over time — e.g., one
    segment is growing while another is declining.

    Args:
        dataset_name: Name of a previously loaded dataset
        time_name: Column name containing time values (e.g. date or month)
        num_segments: Number of segments to find (default: auto)
        max_depth: Maximum number of dimensions to constrain per segment (default: 2)
        fit_sizes: Whether to also fit segment sizes over time (default: auto)
        num_breaks: Number of breakpoints in stylized time series for comparison (default: 3)
        n_jobs: Number of parallel jobs for tree building (default: 10)
        ignore_averages: Ignore overall level when comparing segments (default: True)
        image_size: Optional [width, height] in pixels to generate a PNG chart
    """
    ds = get_dataset(dataset_name)
    df, dims, total_name, size_name = ds["df"], ds["dims"], ds["total_name"], ds["size_name"]

    if time_name not in df.columns:
        raise ValueError(
            f"Time column '{time_name}' not found. Available: {list(df.columns)}"
        )

    result = wise_pizza.explain_timeseries(
        df=df,
        dims=dims,
        total_name=total_name,
        time_name=time_name,
        size_name=size_name,
        num_segments=num_segments,
        max_depth=max_depth,
        solver="tree",
        fit_sizes=fit_sizes,
        num_breaks=num_breaks,
        n_jobs=n_jobs,
        ignore_averages=ignore_averages,
    )

    extracted = extract_result(result)

    if image_size is not None:
        extracted.image_path = _generate_image(result, image_size)

    return result_to_json(extracted)


def main():
    mcp.run()


if __name__ == "__main__":
    main()
