"""S-parameter plotting utilities for MCP server.

This module provides utilities for rendering S-parameter simulation results
to PNG spectral plots, enabling visual feedback when running simulations.
"""

from __future__ import annotations

import base64
import logging
import math
from io import BytesIO
from typing import Any

from mcp.types import ImageContent

__all__ = ["plot_s_parameters", "render_simulation_plot", "HAS_MATPLOTLIB"]

logger = logging.getLogger(__name__)

try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False
    plt = None  # type: ignore[assignment]

DB_FLOOR = -100.0


def _magnitude_squared_db(real: list[float], imag: list[float]) -> list[float]:
    """Compute |S|^2 in dB from real and imaginary parts.

    Args:
        real: Real parts of S-parameter values
        imag: Imaginary parts of S-parameter values

    Returns:
        List of |S|^2 values in dB, floored at DB_FLOOR
    """
    result = []
    for r, i in zip(real, imag):
        mag_sq = r * r + i * i
        if mag_sq <= 0:
            result.append(DB_FLOOR)
        else:
            db_val = 10.0 * math.log10(mag_sq)
            result.append(max(db_val, DB_FLOOR))
    return result


def plot_s_parameters(
    wavelengths: list[float],
    sdict: dict[str, dict[str, dict[str, list[float]]]],
    title: str = "S-Parameters",
    width: int = 800,
    height: int = 500,
) -> bytes | None:
    """Plot S-parameter spectral data as |S|^2 in dB vs wavelength.

    Args:
        wavelengths: List of wavelength values (in micrometers)
        sdict: Nested dict of S-parameters: {port_out: {port_in: {real: [...], imag: [...]}}}
        title: Plot title
        width: Image width in pixels
        height: Image height in pixels

    Returns:
        PNG image as bytes, or None if matplotlib is unavailable or plotting fails
    """
    if not HAS_MATPLOTLIB:
        logger.debug("matplotlib not available, skipping simulation plot")
        return None

    if not wavelengths or not sdict:
        logger.debug("No data to plot")
        return None

    try:
        dpi = 100
        fig, ax = plt.subplots(figsize=(width / dpi, height / dpi), dpi=dpi)

        for port_out, inputs in sorted(sdict.items()):
            for port_in, values in sorted(inputs.items()):
                real = values.get("real", [])
                imag = values.get("imag", [])

                if not real or len(real) != len(wavelengths):
                    continue

                if not imag:
                    imag = [0.0] * len(real)

                db_values = _magnitude_squared_db(real, imag)
                label = f"|S({port_out},{port_in})|²"
                ax.plot(wavelengths, db_values, label=label)

        ax.set_xlabel("Wavelength (μm)")
        ax.set_ylabel("|S|² (dB)")
        ax.set_title(title)
        ax.legend(fontsize="small", loc="best")
        ax.grid(True, alpha=0.3)
        fig.tight_layout()

        buf = BytesIO()
        fig.savefig(buf, format="png")
        plt.close(fig)
        buf.seek(0)
        png_bytes = buf.read()

        logger.debug("Rendered simulation plot: %s (%dx%d)", title, width, height)
        return png_bytes

    except Exception as e:
        logger.warning("Failed to render simulation plot: %s", e)
        return None


def render_simulation_plot(
    response: dict[str, Any],
    component_name: str,
) -> ImageContent | None:
    """Render S-parameter response as a PNG spectral plot.

    Args:
        response: Raw simulation response dict with 'wavelengths' and 'sdict' keys
        component_name: Name of the simulated component (used in plot title)

    Returns:
        ImageContent with base64-encoded PNG, or None if rendering fails
    """
    wavelengths = response.get("wavelengths")
    sdict = response.get("sdict")

    if not wavelengths or not sdict:
        logger.debug("Simulation response missing wavelengths or sdict")
        return None

    png_bytes = plot_s_parameters(
        wavelengths=wavelengths,
        sdict=sdict,
        title=f"{component_name} — S-Parameters",
    )

    if png_bytes is None:
        return None

    return ImageContent(
        type="image",
        data=base64.b64encode(png_bytes).decode("utf-8"),
        mimeType="image/png",
    )
