from pathlib import Path
import json
from typing import Any

import numpy as np

from vna_system.core.logging.logger import get_component_logger
from vna_system.core.acquisition.sweep_buffer import SweepData
from vna_system.core.settings.preset_manager import ConfigPreset

logger = get_component_logger(__file__)


# -----------------------------------------------------------------------------
# Plot builders
# -----------------------------------------------------------------------------
def generate_magnitude_plot_from_sweep_data(
    sweep_data: SweepData,
    preset: ConfigPreset | None = None,
) -> dict[str, Any]:
    """
    Build a Plotly configuration for magnitude-vs-frequency from a `SweepData`.

    Parameters
    ----------
    sweep_data
        Sweep payload with `.points: list[tuple[float, float]]` of (real, imag).
    preset
        Optional preset carrying frequency range (Hz). If absent, defaults are used.

    Returns
    -------
    dict[str, Any]
        Plotly figure spec. If input is invalid, returns {"error": "..."}.
    """
    if not sweep_data or not sweep_data.points:
        logger.warning("Invalid sweep passed to magnitude plot")
        return {"error": "Invalid sweep data"}

    # Frequency range (Hz)
    start_freq = float(preset.start_freq) if (preset and preset.start_freq is not None) else 100e6
    stop_freq = float(preset.stop_freq) if (preset and preset.stop_freq is not None) else 8.8e9

    n = len(sweep_data.points)
    if n == 1:
        freqs = [start_freq]
    else:
        step = (stop_freq - start_freq) / max(1, n - 1)
        freqs = [start_freq + i * step for i in range(n)]

    # Magnitudes (dB). Clamp zero magnitude to -120 dB to avoid -inf.
    mags_db: list[float] = []
    for real, imag in sweep_data.points:
        mag = abs(complex(real, imag))
        mags_db.append(20.0 * np.log10(mag) if mag > 0.0 else -120.0)

    # Reasonable Y range with margin
    ymin = float(min(mags_db))
    ymax = float(max(mags_db))
    ymargin = (ymax - ymin) * 0.1 if ymax > ymin else 10.0
    y_min = max(ymin - ymargin, -120.0)
    y_max = min(ymax + ymargin, 20.0)

    trace = {
        "x": [f / 1e9 for f in freqs],  # Hz -> GHz
        "y": mags_db,
        "type": "scatter",
        "mode": "lines",
        "name": "Magnitude",
        "line": {"color": "#1f77b4", "width": 2},
    }

    fig = {
        "data": [trace],
        "layout": {
            "title": "Magnitude Response",
            "xaxis": {
                "title": "Frequency (GHz)",
                "showgrid": True,
                "gridcolor": "#e5e5e5",
                "gridwidth": 1,
            },
            "yaxis": {
                "title": "Magnitude (dB)",
                "range": [y_min, y_max],
                "showgrid": True,
                "gridcolor": "#e5e5e5",
                "gridwidth": 1,
            },
            "plot_bgcolor": "#fafafa",
            "paper_bgcolor": "#ffffff",
            "font": {"family": "Arial, sans-serif", "size": 12, "color": "#333333"},
            "hovermode": "x unified",
            "showlegend": True,
            "margin": {"l": 60, "r": 40, "t": 60, "b": 60},
        },
    }
    logger.debug("Magnitude plot generated", points=n)
    return fig


# -----------------------------------------------------------------------------
# IO helpers
# -----------------------------------------------------------------------------
def load_sweep_data_from_json(json_file: Path) -> SweepData:
    """
    Load a `SweepData` structure from a JSON file.

    The file is expected to contain:
      { "sweep_number": int, "timestamp": float, "points": [[r, i], ...], "total_points": int }

    If `total_points` is missing, it is derived from the length of `points`.
    """
    data = json.loads(Path(json_file).read_text(encoding="utf-8"))

    points = data.get("points", [])
    if not isinstance(points, list):
        raise ValueError(f"Invalid 'points' in file: {json_file}")

    # Normalize to list[tuple[float, float]]
    norm_points: list[tuple[float, float]] = []
    for pt in points:
        if not (isinstance(pt, (list, tuple)) and len(pt) == 2):
            raise ValueError(f"Invalid point format in {json_file}: {pt!r}")
        r, i = pt
        norm_points.append((float(r), float(i)))

    sweep = SweepData(
        sweep_number=int(data.get("sweep_number", 0)),
        timestamp=float(data.get("timestamp", 0.0)),
        points=norm_points,
        total_points=int(data.get("total_points", len(norm_points))),
    )
    return sweep


# -----------------------------------------------------------------------------
# Calibration plots
# -----------------------------------------------------------------------------
def generate_standards_magnitude_plots(
    calibration_path: Path,
    preset: ConfigPreset | None = None,
) -> dict[str, Any]:
    """
    Build individual magnitude plots for all calibration standards found under a folder.

    The function scans `calibration_path` for `*.json` (ignoring `*metadata.json`), loads each
    sweep, and produces a Plotly config per standard. Raw sweep data and (optional) frequency
    info are embedded into the output for convenience.

    Returns
    -------
    dict[str, Any]
        { "<standard>": <plotly fig | {'error': str}>, ... }
    """
    plots: dict[str, Any] = {}

    standard_colors = {
        "open": "#2ca02c",    # Green
        "short": "#d62728",   # Red
        "load": "#ff7f0e",    # Orange
        "through": "#1f77b4", # Blue
    }

    for standard_file in Path(calibration_path).glob("*.json"):
        standard_name = standard_file.stem
        if "metadata" in standard_name or "calibration_info" in standard_name:
            logger.debug(f"Skipping {standard_name=}")
            continue

        try:
            sweep = load_sweep_data_from_json(standard_file)
            fig = generate_magnitude_plot_from_sweep_data(sweep, preset)

            if "error" in fig:
                plots[standard_name] = fig
                continue

            # Customize per-standard appearance/title
            if fig.get("data"):
                fig["data"][0]["line"]["color"] = standard_colors.get(standard_name, "#1f77b4")
                fig["data"][0]["name"] = f"{standard_name.upper()} Standard"
                fig["layout"]["title"] = f"{standard_name.upper()} Standard Magnitude"

            # Attach raw sweep block for UI download/inspection
            fig["raw_sweep_data"] = {
                "sweep_number": sweep.sweep_number,
                "timestamp": sweep.timestamp,
                "total_points": sweep.total_points,
                "points": sweep.points,
                "file_path": str(standard_file),
            }

            # Optional frequency info
            if preset:
                fig["frequency_info"] = {
                    "start_freq": preset.start_freq,
                    "stop_freq": preset.stop_freq,
                    "points": preset.points,
                    "bandwidth": preset.bandwidth,
                }

            plots[standard_name] = fig

        except (json.JSONDecodeError, FileNotFoundError, KeyError, ValueError) as exc:
            logger.warning("Failed to load standard plot", file=str(standard_file), error=repr(exc))
            plots[standard_name] = {"error": f"Failed to load {standard_name}: {exc}"}

    return plots


def generate_combined_standards_plot(
    calibration_path: Path,
    preset: ConfigPreset | None = None,
) -> dict[str, Any]:
    """
    Build a combined Plotly figure that overlays all available calibration standards.

    Each standard is rendered as a separate trace with a canonical color.
    """
    traces: list[dict[str, Any]] = []
    standard_colors = {
        "open": "#2ca02c",    # Green
        "short": "#d62728",   # Red
        "load": "#ff7f0e",    # Orange
        "through": "#1f77b4", # Blue
    }

    # Initialize Y range trackers inversely so first update sets them correctly.
    y_min, y_max = 0.0, -120.0

    for standard_file in Path(calibration_path).glob("*.json"):
        name = standard_file.stem
        if "metadata" in name:
            continue

        try:
            sweep = load_sweep_data_from_json(standard_file)
            fig = generate_magnitude_plot_from_sweep_data(sweep, preset)
            if "error" in fig or not fig.get("data"):
                continue

            trace = dict(fig["data"][0])  # shallow copy
            trace["line"]["color"] = standard_colors.get(name, "#1f77b4")
            trace["name"] = f"{name.upper()} Standard"
            traces.append(trace)

            # Update Y range
            y_vals = trace.get("y") or []
            if y_vals:
                tmin = float(min(y_vals))
                tmax = float(max(y_vals))
                y_min = min(y_min, tmin)
                y_max = max(y_max, tmax)

        except (json.JSONDecodeError, FileNotFoundError, KeyError, ValueError) as exc:
            logger.warning("Failed to include standard in combined plot", file=str(standard_file), error=repr(exc))
            continue

    if not traces:
        return {"error": "No valid calibration standards found"}

    ymargin = (y_max - y_min) * 0.1 if y_max > y_min else 10.0
    y_min = max(y_min - ymargin, -120.0)
    y_max = min(y_max + ymargin, 20.0)

    fig = {
        "data": traces,
        "layout": {
            "title": "Calibration Standards Comparison",
            "xaxis": {
                "title": "Frequency (GHz)",
                "showgrid": True,
                "gridcolor": "#e5e5e5",
                "gridwidth": 1,
            },
            "yaxis": {
                "title": "Magnitude (dB)",
                "range": [y_min, y_max],
                "showgrid": True,
                "gridcolor": "#e5e5e5",
                "gridwidth": 1,
            },
            "plot_bgcolor": "#fafafa",
            "paper_bgcolor": "#ffffff",
            "font": {"family": "Arial, sans-serif", "size": 12, "color": "#333333"},
            "hovermode": "x unified",
            "showlegend": True,
            "margin": {"l": 60, "r": 40, "t": 60, "b": 60},
        },
    }
    logger.debug("Combined standards plot generated", traces=len(traces))
    return fig
