src/charting.py

"""Chart rendering utilities for the MCP charts server."""

from __future__ import annotations

import base64
import io
from typing import Iterable

import matplotlib
import matplotlib.pyplot as plt
import numpy as np

matplotlib.use("Agg")

ALLOWED_CHART_TYPES = {"line", "bar", "scatter", "pie", "grouped_bar", "stacked_bar"}


def _validate_lengths(x_values: list[float], y_values: list[float]) -> None:
    if not x_values or not y_values:
        raise ValueError("x_values and y_values must not be empty.")
    if len(x_values) != len(y_values):
        raise ValueError("x_values and y_values must have the same length.")


def _to_float_list(values: Iterable[float | int | str], field_name: str) -> list[float]:
    try:
        return [float(v) for v in values]
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} must contain only numeric values.") from exc


def _to_float_matrix(
    values: Iterable[Iterable[float | int | str]],
    field_name: str,
) -> list[list[float]]:
    rows: list[list[float]] = []
    try:
        for idx, row in enumerate(values):
            rows.append(_to_float_list(row, f"{field_name}[{idx}]"))
    except TypeError as exc:
        raise ValueError(f"{field_name} must be a list of numeric lists.") from exc
    if not rows:
        raise ValueError(f"{field_name} must not be empty.")
    return rows


def _to_str_list(values: Iterable[float | int | str], field_name: str) -> list[str]:
    try:
        return [str(v) for v in values]
    except TypeError as exc:
        raise ValueError(f"{field_name} must be a list.") from exc


def render_chart_png_base64(
    chart_type: str,
    x_values: list[float | int | str],
    y_values: list[float | int | str],
    *,
    y_series: list[list[float | int | str]] | None = None,
    series_labels: list[str] | None = None,
    pie_labels: list[str] | None = None,
    title: str = "Generated Chart",
    x_label: str = "X",
    y_label: str = "Y",
    color: str = "#1f77b4",
) -> str:
    """Render a chart and return PNG bytes as base64."""

    chart_type = chart_type.lower().strip()
    if chart_type not in ALLOWED_CHART_TYPES:
        raise ValueError(
            f"Unsupported chart_type '{chart_type}'. Choose one of {sorted(ALLOWED_CHART_TYPES)}."
        )

    fig, ax = plt.subplots(figsize=(9, 5), dpi=120)
    try:
        if chart_type == "line":
            x = _to_float_list(x_values, "x_values")
            y = _to_float_list(y_values, "y_values")
            _validate_lengths(x, y)
            ax.plot(x, y, color=color, linewidth=2.2, marker="o")
        elif chart_type == "bar":
            x = _to_str_list(x_values, "x_values")
            y = _to_float_list(y_values, "y_values")
            _validate_lengths(x, y)
            ax.bar(x, y, color=color, alpha=0.9)
        elif chart_type == "scatter":
            x = _to_float_list(x_values, "x_values")
            y = _to_float_list(y_values, "y_values")
            _validate_lengths(x, y)
            ax.scatter(x, y, color=color, s=50)
        elif chart_type == "pie":
            y = _to_float_list(y_values, "y_values")
            if not y:
                raise ValueError("y_values must not be empty for pie charts.")
            if pie_labels is not None:
                labels = [str(label) for label in pie_labels]
            elif x_values:
                labels = _to_str_list(x_values, "x_values")
            else:
                labels = [f"Slice {i + 1}" for i in range(len(y))]
            if len(labels) != len(y):
                raise ValueError("pie_labels/x_values length must match y_values length.")
            ax.pie(y, labels=labels, autopct="%1.1f%%", startangle=90)
            ax.axis("equal")
        elif chart_type in {"grouped_bar", "stacked_bar"}:
            categories = _to_str_list(x_values, "x_values")
            matrix = _to_float_matrix(y_series or [], "y_series")
            num_categories = len(categories)
            if num_categories == 0:
                raise ValueError("x_values must not be empty for grouped/stacked bar charts.")
            for idx, row in enumerate(matrix):
                if len(row) != num_categories:
                    raise ValueError(
                        f"y_series[{idx}] must have {num_categories} values to match x_values."
                    )

            labels = (
                [str(label) for label in series_labels]
                if series_labels is not None
                else [f"Series {i + 1}" for i in range(len(matrix))]
            )
            if len(labels) != len(matrix):
                raise ValueError("series_labels length must match y_series length.")

            x_pos = np.arange(num_categories)
            if chart_type == "grouped_bar":
                width = 0.8 / len(matrix)
                for i, row in enumerate(matrix):
                    offset = (i - (len(matrix) - 1) / 2) * width
                    ax.bar(x_pos + offset, row, width=width, label=labels[i])
            else:
                bottoms = np.zeros(num_categories)
                for i, row in enumerate(matrix):
                    ax.bar(x_pos, row, bottom=bottoms, label=labels[i], alpha=0.9)
                    bottoms = bottoms + np.array(row)

            ax.set_xticks(x_pos)
            ax.set_xticklabels(categories)
            ax.legend()

        ax.set_title(title)
        if chart_type != "pie":
            ax.set_xlabel(x_label)
            ax.set_ylabel(y_label)
            ax.grid(True, linestyle="--", alpha=0.35)
        fig.tight_layout()

        buffer = io.BytesIO()
        fig.savefig(buffer, format="png")
        buffer.seek(0)
        image_base64 = base64.b64encode(buffer.read()).decode("ascii")
        return image_base64
    finally:
        plt.close(fig)