"""Chart rendering module for the KAPSARC MCP server.

Auto-detects chart types from data shape, renders matplotlib charts with
official KAPSARC branding, and returns PNG images via FastMCP's Image class.
"""

from __future__ import annotations

import io
from collections import Counter
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
from mcp.server.fastmcp.utilities.types import Image  # noqa: E402

# ---------------------------------------------------------------------------
# KAPSARC official brand palette
# ---------------------------------------------------------------------------
# Primary: KAPSARC Green + Dark Teal header from datasource.kapsarc.org
# Chart series: official data portal chart palette
KAPSARC_GREEN = "#019247"
KAPSARC_DARK = "#173744"
KAPSARC_ACCENT_BLUE = "#79bde8"
KAPSARC_YELLOW = "#ffd200"
KAPSARC_BG = "#F8F8F8"

# Data portal chart colors (from datasource.kapsarc.org theme)
COLORS = [
    "#019247",  # KAPSARC Green (primary)
    "#173744",  # KAPSARC Dark Teal
    "#79bde8",  # KAPSARC Accent Blue
    "#fc8d62",  # Warm orange
    "#e78ac3",  # Rose pink
    "#ffd200",  # KAPSARC Yellow
    "#a6d854",  # Lime green
    "#8da0cb",  # Soft blue
    "#e5c494",  # Tan
    "#66c2a5",  # Teal
]

# Logo path (bundled with the server)
_LOGO_PATH = Path(__file__).resolve().parent / "assets" / "kapsarc_logo.png"

MAX_CATEGORIES = 25
MAX_PIE_SLICES = 10
MAX_LABEL_LEN = 25
MAX_IMAGE_BYTES = 800_000
LOW_DPI = 80
DEFAULT_DPI = 100
SOURCE_TEXT = "Source: KAPSARC"

# ---------------------------------------------------------------------------
# Field classification
# ---------------------------------------------------------------------------

TEMPORAL_KEYWORDS = {"year", "date", "month", "period", "quarter", "time"}


def _classify_field(name: str, values: list) -> str:
    """Classify a field as temporal, numeric, categorical, or id."""
    non_null = [v for v in values if v is not None and v != ""]
    if not non_null:
        return "categorical"

    name_lower = name.lower()
    if any(kw in name_lower for kw in TEMPORAL_KEYWORDS):
        return "temporal"

    # Check if all values are numeric
    all_numeric = True
    all_int_like = True
    for v in non_null:
        if isinstance(v, (int, float)):
            if not isinstance(v, int) and not v.is_integer():
                all_int_like = False
        elif isinstance(v, str):
            try:
                f = float(v)
                if not f.is_integer():
                    all_int_like = False
            except (ValueError, TypeError):
                all_numeric = False
                break
        else:
            all_numeric = False
            break

    if all_numeric:
        # Integers in 1900-2100 range look temporal
        if all_int_like:
            int_vals = [int(float(v)) for v in non_null]
            if all(1900 <= iv <= 2100 for iv in int_vals):
                return "temporal"
        return "numeric"

    # Check if values are unique per record (id-like)
    if len(set(str(v) for v in non_null)) == len(non_null) and len(non_null) > 5:
        return "id"

    return "categorical"


def _classify_fields(records: list[dict]) -> dict[str, str]:
    """Classify all fields in the record set."""
    if not records:
        return {}
    fields = {}
    for key in records[0]:
        values = [r.get(key) for r in records]
        fields[key] = _classify_field(key, values)
    return fields


# ---------------------------------------------------------------------------
# Auto-detection: fields and chart type
# ---------------------------------------------------------------------------


def _auto_select_fields(
    records: list[dict],
    classifications: dict[str, str],
    x_field: str | None,
    y_field: str | None,
    series_field: str | None,
) -> tuple[str, str, str | None]:
    """Pick x, y, and optionally series fields from the data."""
    temporals = [k for k, v in classifications.items() if v == "temporal"]
    numerics = [k for k, v in classifications.items() if v == "numeric"]
    categoricals = [k for k, v in classifications.items() if v == "categorical"]

    if x_field is None:
        if temporals:
            x_field = temporals[0]
        elif categoricals:
            x_field = categoricals[0]
        elif numerics and len(numerics) >= 2:
            x_field = numerics[0]
        else:
            x_field = list(records[0].keys())[0]

    if y_field is None:
        # Pick the first numeric field that isn't x
        for n in numerics:
            if n != x_field:
                y_field = n
                break
        if y_field is None:
            # Fall back to counting occurrences
            y_field = x_field

    return x_field, y_field, series_field


def _detect_chart_type(
    records: list[dict],
    x_field: str,
    y_field: str,
    series_field: str | None,
    classifications: dict[str, str],
) -> str:
    """Auto-detect the best chart type for the data shape."""
    x_type = classifications.get(x_field, "categorical")
    y_type = classifications.get(y_field, "numeric")

    # Extract y values for analysis
    y_vals = []
    for r in records:
        v = r.get(y_field)
        if v is not None:
            try:
                y_vals.append(float(v))
            except (ValueError, TypeError):
                pass

    has_negatives = any(v < 0 for v in y_vals)
    x_values = [r.get(x_field) for r in records]
    unique_x = len(set(str(v) for v in x_values if v is not None))

    if series_field:
        if x_type == "temporal":
            return "multi_line"
        return "multi_bar"

    if x_type == "temporal":
        return "line"

    if x_type == "categorical":
        if (
            not has_negatives
            and unique_x <= 8
            and y_vals
            and 95 <= sum(y_vals) <= 105
        ):
            return "pie"
        if unique_x > 15:
            return "hbar"
        return "bar"

    if x_type == "numeric" and y_type == "numeric":
        return "scatter"

    return "bar"


# ---------------------------------------------------------------------------
# Data cleaning
# ---------------------------------------------------------------------------


def _clean_records(
    records: list[dict],
    x_field: str,
    y_field: str,
    classifications: dict[str, str],
) -> list[dict]:
    """Clean records: drop nulls, coerce types, sort, truncate."""
    cleaned = []
    for r in records:
        x_val = r.get(x_field)
        y_val = r.get(y_field)
        if x_val is None or x_val == "":
            continue
        if x_field != y_field:
            if y_val is None or y_val == "":
                continue
            try:
                y_val = float(y_val)
            except (ValueError, TypeError):
                continue
        new_r = dict(r)
        if x_field != y_field:
            new_r[y_field] = y_val
        cleaned.append(new_r)

    # Sort by x if temporal
    if classifications.get(x_field) == "temporal":
        try:
            cleaned.sort(key=lambda r: float(r[x_field]))
        except (ValueError, TypeError):
            cleaned.sort(key=lambda r: str(r[x_field]))

    # Truncate long category labels
    if classifications.get(x_field) == "categorical":
        for r in cleaned:
            label = str(r[x_field])
            if len(label) > MAX_LABEL_LEN:
                r[x_field] = label[:MAX_LABEL_LEN] + "…"

    return cleaned


def _group_excess_categories(
    x_vals: list,
    y_vals: list[float],
    max_cats: int = MAX_CATEGORIES,
) -> tuple[list, list[float]]:
    """Keep top categories, group the rest as 'Other'."""
    if len(x_vals) <= max_cats:
        return x_vals, y_vals

    paired = list(zip(x_vals, y_vals))
    paired.sort(key=lambda p: abs(p[1]), reverse=True)
    top = paired[:max_cats]
    rest = paired[max_cats:]
    other_sum = sum(v for _, v in rest)
    top_x = [p[0] for p in top]
    top_y = [p[1] for p in top]
    top_x.append("Other")
    top_y.append(other_sum)
    return top_x, top_y


# ---------------------------------------------------------------------------
# Styling helpers
# ---------------------------------------------------------------------------


def _load_logo():
    """Load the KAPSARC logo image, or return None if unavailable."""
    if _LOGO_PATH.exists():
        try:
            return plt.imread(str(_LOGO_PATH))
        except Exception:
            return None
    return None


def _apply_style(ax: plt.Axes, title: str | None) -> None:
    """Apply KAPSARC theme styling to an axes."""
    ax.set_facecolor(KAPSARC_BG)
    ax.grid(True, alpha=0.25, linestyle="--", color="#cccccc")
    if title:
        ax.set_title(title, fontsize=14, fontweight="bold", pad=12, color=KAPSARC_DARK)
    ax.tick_params(labelsize=10, colors="#333333")
    for spine in ax.spines.values():
        spine.set_color("#dddddd")


def _add_branding(fig: plt.Figure) -> None:
    """Add KAPSARC logo and source text to the figure."""
    # Source text in the dedicated footer band
    fig.text(
        0.97,
        0.025,
        SOURCE_TEXT,
        ha="right",
        va="center",
        fontsize=9,
        color="#666666",
        style="italic",
    )

    # Logo in the dedicated footer band
    logo_data = _load_logo()
    if logo_data is not None:
        logo_ax = fig.add_axes([0.01, 0.005, 0.12, 0.055], anchor="SW")
        logo_ax.imshow(logo_data)
        logo_ax.axis("off")


# Reserved space (fraction of figure height) for branding elements
_HEADER_HEIGHT = 0.025  # green bar thickness
_TOP_PAD = 0.09         # green bar + title clearance
_BOTTOM_PAD = 0.12      # logo + source text + x-label room


def _new_figure() -> tuple[plt.Figure, plt.Axes]:
    """Create a new styled figure with reserved space for branding."""
    fig = plt.figure(figsize=(10, 7), dpi=DEFAULT_DPI)
    fig.patch.set_facecolor("white")

    # Plot area sits safely between header and footer
    ax = fig.add_axes([0.10, _BOTTOM_PAD, 0.85, 1.0 - _TOP_PAD - _BOTTOM_PAD])

    # Green accent bar pinned to very top
    fig.patches.append(
        plt.Rectangle(
            (0, 1.0 - _HEADER_HEIGHT), 1, _HEADER_HEIGHT,
            transform=fig.transFigure,
            facecolor=KAPSARC_GREEN,
            edgecolor="none",
            zorder=10,
        )
    )
    return fig, ax


# ---------------------------------------------------------------------------
# Renderers
# ---------------------------------------------------------------------------


def _rotate_xlabels(ax: plt.Axes) -> None:
    """Rotate x-axis labels for readability without touching figure layout."""
    plt.setp(ax.get_xticklabels(), rotation=40, ha="right", rotation_mode="anchor")


def _render_bar(x_vals: list, y_vals: list[float], title: str | None) -> plt.Figure:
    fig, ax = _new_figure()
    colors = [COLORS[i % len(COLORS)] for i in range(len(x_vals))]
    ax.bar([str(v) for v in x_vals], y_vals, color=colors, edgecolor="white", linewidth=0.5)
    ax.set_xlabel("")
    _apply_style(ax, title)
    _rotate_xlabels(ax)
    return fig


def _render_hbar(x_vals: list, y_vals: list[float], title: str | None) -> plt.Figure:
    fig, ax = _new_figure()
    colors = [COLORS[i % len(COLORS)] for i in range(len(x_vals))]
    y_pos = range(len(x_vals))
    ax.barh(y_pos, y_vals, color=colors, edgecolor="white", linewidth=0.5)
    ax.set_yticks(y_pos)
    ax.set_yticklabels([str(v) for v in x_vals])
    _apply_style(ax, title)
    return fig


def _render_line(x_vals: list, y_vals: list[float], title: str | None) -> plt.Figure:
    fig, ax = _new_figure()
    ax.plot(
        [str(v) for v in x_vals],
        y_vals,
        color=COLORS[0],
        linewidth=2,
        marker="o",
        markersize=5,
    )
    ax.fill_between(
        range(len(x_vals)),
        y_vals,
        alpha=0.1,
        color=COLORS[0],
    )
    _apply_style(ax, title)
    _rotate_xlabels(ax)
    return fig


def _render_pie(x_vals: list, y_vals: list[float], title: str | None) -> plt.Figure:
    # Cap at MAX_PIE_SLICES
    if len(x_vals) > MAX_PIE_SLICES:
        paired = sorted(zip(x_vals, y_vals), key=lambda p: p[1], reverse=True)
        top = paired[:MAX_PIE_SLICES]
        rest_sum = sum(v for _, v in paired[MAX_PIE_SLICES:])
        x_vals = [p[0] for p in top] + ["Other"]
        y_vals = [p[1] for p in top] + [rest_sum]

    fig, ax = _new_figure()
    colors = [COLORS[i % len(COLORS)] for i in range(len(x_vals))]
    wedges, texts, autotexts = ax.pie(
        y_vals,
        labels=[str(v) for v in x_vals],
        colors=colors,
        autopct="%1.1f%%",
        startangle=90,
        pctdistance=0.8,
    )
    for t in autotexts:
        t.set_fontsize(9)
    _apply_style(ax, title)
    ax.grid(False)
    return fig


def _render_scatter(x_vals: list[float], y_vals: list[float], title: str | None, x_label: str, y_label: str) -> plt.Figure:
    fig, ax = _new_figure()
    ax.scatter(x_vals, y_vals, color=COLORS[0], alpha=0.7, edgecolors="white", linewidth=0.5, s=60)
    ax.set_xlabel(x_label, fontsize=11)
    ax.set_ylabel(y_label, fontsize=11)
    _apply_style(ax, title)
    return fig


def _render_multi_line(
    records: list[dict],
    x_field: str,
    y_field: str,
    series_field: str,
    title: str | None,
) -> plt.Figure:
    fig, ax = _new_figure()

    # Group by series
    series_data: dict[str, tuple[list, list[float]]] = {}
    for r in records:
        s = str(r.get(series_field, "?"))
        x = r.get(x_field)
        y = r.get(y_field)
        if x is None or y is None:
            continue
        try:
            y = float(y)
        except (ValueError, TypeError):
            continue
        if s not in series_data:
            series_data[s] = ([], [])
        series_data[s][0].append(x)
        series_data[s][1].append(y)

    for i, (label, (xs, ys)) in enumerate(series_data.items()):
        # Sort by x
        paired = sorted(zip(xs, ys), key=lambda p: str(p[0]))
        xs_sorted = [str(p[0]) for p in paired]
        ys_sorted = [p[1] for p in paired]
        ax.plot(
            xs_sorted,
            ys_sorted,
            color=COLORS[i % len(COLORS)],
            linewidth=2,
            marker="o",
            markersize=4,
            label=label,
        )

    ax.legend(fontsize=9, loc="best")
    _apply_style(ax, title)
    _rotate_xlabels(ax)
    return fig


def _render_multi_bar(
    records: list[dict],
    x_field: str,
    y_field: str,
    series_field: str,
    title: str | None,
) -> plt.Figure:
    # Group by series
    series_data: dict[str, dict[str, float]] = {}
    all_x = []
    for r in records:
        s = str(r.get(series_field, "?"))
        x = str(r.get(x_field, "?"))
        y = r.get(y_field)
        if y is None:
            continue
        try:
            y = float(y)
        except (ValueError, TypeError):
            continue
        if s not in series_data:
            series_data[s] = {}
        series_data[s][x] = y
        if x not in all_x:
            all_x.append(x)

    fig, ax = _new_figure()
    n_series = len(series_data)
    if n_series == 0:
        _apply_style(ax, title)
        return fig

    bar_width = 0.8 / n_series
    x_pos = np.arange(len(all_x))

    for i, (label, vals) in enumerate(series_data.items()):
        offsets = x_pos + i * bar_width - (n_series - 1) * bar_width / 2
        heights = [vals.get(x, 0) for x in all_x]
        ax.bar(
            offsets,
            heights,
            bar_width,
            color=COLORS[i % len(COLORS)],
            label=label,
            edgecolor="white",
            linewidth=0.5,
        )

    ax.set_xticks(x_pos)
    ax.set_xticklabels(all_x)
    ax.legend(fontsize=9, loc="best")
    _apply_style(ax, title)
    _rotate_xlabels(ax)
    return fig


# ---------------------------------------------------------------------------
# Figure to Image
# ---------------------------------------------------------------------------


def _fig_to_image(fig: plt.Figure) -> Image:
    """Add branding, convert figure to PNG Image, and close."""
    _add_branding(fig)
    buf = io.BytesIO()
    try:
        # Use the exact figure layout (no bbox_inches="tight" which would crop
        # the reserved branding margins)
        fig.savefig(buf, format="png", dpi=DEFAULT_DPI)
        data = buf.getvalue()
    finally:
        plt.close(fig)
        buf.close()
    return Image(data=data, format="png")


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def render_chart(
    records: list[dict],
    x_field: str | None = None,
    y_field: str | None = None,
    series_field: str | None = None,
    chart_type: str | None = None,
    title: str | None = None,
) -> Image:
    """Render a chart from records, auto-detecting fields and chart type.

    Args:
        records: List of dicts (rows from a dataset).
        x_field: Field for x-axis (auto-detected if None).
        y_field: Field for y-axis (auto-detected if None).
        series_field: Field for grouping into series (optional).
        chart_type: Explicit chart type override. One of:
            bar, hbar, line, pie, scatter, multi_line, multi_bar.
        title: Chart title (auto-generated if None).

    Returns:
        FastMCP Image with PNG data.

    Raises:
        ValueError: If records are empty or all values are null.
    """
    if not records:
        raise ValueError("No data to visualize")

    classifications = _classify_fields(records)
    x_field, y_field, series_field = _auto_select_fields(
        records, classifications, x_field, y_field, series_field
    )

    # Clean data
    cleaned = _clean_records(records, x_field, y_field, classifications)
    if not cleaned:
        raise ValueError("All values are null after cleaning")

    # Detect chart type
    if chart_type is None:
        if series_field:
            chart_type = _detect_chart_type(
                cleaned, x_field, y_field, series_field, classifications
            )
        else:
            chart_type = _detect_chart_type(
                cleaned, x_field, y_field, None, classifications
            )

    # Auto-generate title
    if title is None:
        y_label = y_field.replace("_", " ").title()
        x_label = x_field.replace("_", " ").title()
        if series_field:
            s_label = series_field.replace("_", " ").title()
            title = f"{y_label} by {x_label} ({s_label})"
        else:
            title = f"{y_label} by {x_label}"

    # Handle multi-series charts
    if chart_type in ("multi_line", "multi_bar") and series_field:
        fig = None
        try:
            if chart_type == "multi_line":
                fig = _render_multi_line(cleaned, x_field, y_field, series_field, title)
            else:
                fig = _render_multi_bar(cleaned, x_field, y_field, series_field, title)
            return _fig_to_image(fig)
        except Exception:
            if fig is not None:
                plt.close(fig)
            raise

    # Extract x and y values for single-series charts
    x_vals = [r[x_field] for r in cleaned]
    if x_field == y_field:
        # Counting mode — count occurrences
        counts = Counter(str(v) for v in x_vals)
        x_vals = list(counts.keys())
        y_vals = [float(c) for c in counts.values()]
    else:
        y_vals = [float(r[y_field]) for r in cleaned]

    # Safety: negative values force bar over pie
    has_negatives = any(v < 0 for v in y_vals)
    if chart_type == "pie" and has_negatives:
        chart_type = "bar"

    # Group excess categories
    if chart_type in ("bar", "hbar", "pie"):
        x_vals, y_vals = _group_excess_categories(x_vals, y_vals)

    # Render
    fig = None
    try:
        if chart_type == "bar":
            fig = _render_bar(x_vals, y_vals, title)
        elif chart_type == "hbar":
            fig = _render_hbar(x_vals, y_vals, title)
        elif chart_type == "line":
            fig = _render_line(x_vals, y_vals, title)
        elif chart_type == "pie":
            fig = _render_pie(x_vals, y_vals, title)
        elif chart_type == "scatter":
            x_floats = [float(v) for v in x_vals]
            x_label = x_field.replace("_", " ").title()
            y_label = y_field.replace("_", " ").title()
            fig = _render_scatter(x_floats, y_vals, title, x_label, y_label)
        else:
            fig = _render_bar(x_vals, y_vals, title)
        return _fig_to_image(fig)
    except Exception:
        if fig is not None:
            plt.close(fig)
        raise
