"""Shared utilities for the quiver demo generators.

Mirrors scripts/charts/_shared.py, with two differences forced by the fact
that the quiver trace is new in plotly.js 4.0:

  * Figures are plain dicts, not go.Figure — plotly.py doesn't know the
    `quiver` trace type yet, so there is nothing to validate against.
  * No Kaleido. The static PNGs are rendered by scripts/quiver/render_pngs.mjs,
    which loads each JSON into headless Chromium with the actual
    plotly-4.0.0 bundle (so the PNG is exactly what the CDN build draws).

`save()` therefore only writes the figure JSON; run render_pngs.mjs afterwards
for the PNGs + thumbnails.
"""
import io
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd
import requests

# src/assets/static/quiver is copied verbatim into build/ by the templater.
OUT = Path(__file__).parent.parent.parent / "src/assets/static/quiver"
OUT.mkdir(parents=True, exist_ok=True)

# ── Palette (same canonical colormap as scripts/charts/_shared.py) ──────────
VIOLET   = "#845EEE"
TEAL     = "#52B3D0"
GREEN    = "#55B685"
PINK     = "#DA5597"
ORANGE   = "#E9A23B"
PRIMARY   = VIOLET
SECONDARY = TEAL
COLORWAY = [VIOLET, TEAL, GREEN, PINK, ORANGE]
BG    = "#ffffff"
TEXT  = "#1c2024"
GRID  = "#d9d9e0"
MUTED = "#60646c"
FONT  = "Inter, -apple-system, BlinkMacSystemFont, sans-serif"

# Violet-forward sequential colorscale for magnitude coloring.
VIOLET_SCALE = [
    [0.0, "rgba(132, 94, 238, 0.25)"],
    [1.0, "rgba(132, 94, 238, 1.0)"],
]
# Light-grey → violet, for fields where the low end should recede entirely.
GREY_VIOLET_SCALE = [
    [0.0, "#d9d9e0"],
    [1.0, VIOLET],
]
# Teal → violet → pink, a punchier three-stop ramp for speed maps.
SPEED_SCALE = [
    [0.0, "#b7e0ec"],
    [0.45, TEAL],
    [0.8, VIOLET],
    [1.0, PINK],
]


def themed_layout(**overrides) -> dict:
    """Layout dict pre-filled with the gallery's light theme."""
    layout = {
        "paper_bgcolor": BG,
        "plot_bgcolor": BG,
        "colorway": COLORWAY,
        "font": {"family": FONT, "color": TEXT, "size": 12},
        "hoverlabel": {
            "bgcolor": "#f0f0f3",
            "font": {"color": TEXT, "family": FONT},
            "bordercolor": GRID,
        },
        "margin": {"t": 60, "b": 50, "l": 50, "r": 40},
        "height": 500,
        "xaxis": {"gridcolor": GRID, "linecolor": GRID, "zerolinecolor": GRID,
                  "automargin": True},
        "yaxis": {"gridcolor": GRID, "linecolor": GRID, "zerolinecolor": GRID,
                  "automargin": True},
    }
    for key, val in overrides.items():
        if isinstance(val, dict) and isinstance(layout.get(key), dict):
            layout[key] = {**layout[key], **val}
        else:
            layout[key] = val
    return layout


# ── Network helpers ──────────────────────────────────────────────────────────
def fetch_csv(url: str, **kwargs) -> pd.DataFrame:
    r = requests.get(url, timeout=120)
    r.raise_for_status()
    return pd.read_csv(io.StringIO(r.text), **kwargs)


def fetch_json(url: str):
    r = requests.get(url, timeout=120)
    r.raise_for_status()
    return r.json()


# ── Basemap outlines (Natural Earth GeoJSON → one None-gapped line trace) ───
NE_BASE = "https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/geojson"


def _geojson_lines(url: str, bbox=None):
    """Flatten a GeoJSON of lines/polygons into None-gapped lon/lat lists,
    keeping only the segments inside `bbox` (lon0, lat0, lon1, lat1)."""
    gj = fetch_json(url)
    lons, lats = [], []

    def add_line(coords):
        run = []
        for lon, lat in coords:
            inside = bbox is None or (
                bbox[0] <= lon <= bbox[2] and bbox[1] <= lat <= bbox[3]
            )
            if inside:
                run.append((lon, lat))
            elif run:
                _flush(run)
                run = []
        if run:
            _flush(run)

    def _flush(run):
        if len(run) < 2:
            return
        for lon, lat in run:
            lons.append(lon)
            lats.append(lat)
        lons.append(None)
        lats.append(None)

    for feat in gj["features"]:
        geom = feat["geometry"]
        if geom is None:
            continue
        gtype, coords = geom["type"], geom["coordinates"]
        if gtype == "LineString":
            add_line(coords)
        elif gtype in ("MultiLineString", "Polygon"):
            for part in coords:
                add_line(part)
        elif gtype == "MultiPolygon":
            for poly in coords:
                for part in poly:
                    add_line(part)
    return lons, lats


def outline_trace(kind: str, bbox=None, res: str = "50m", color=GRID, width=1) -> dict:
    """A muted basemap line trace. kind: 'coastline' or 'states'."""
    fname = {
        "coastline": f"ne_{res}_coastline.geojson",
        "states": f"ne_{res}_admin_1_states_provinces_lines.geojson",
    }[kind]
    lons, lats = _geojson_lines(f"{NE_BASE}/{fname}", bbox)
    return {
        "type": "scatter",
        "x": lons,
        "y": lats,
        "mode": "lines",
        "line": {"color": color, "width": width},
        "hoverinfo": "skip",
        "showlegend": False,
    }


# ── Save ─────────────────────────────────────────────────────────────────────
def _round_floats(obj, ndigits: int):
    """Round every float in a nested structure — keeps the JSON small."""
    if isinstance(obj, float):
        if math.isnan(obj):
            return None
        return round(obj, ndigits)
    if isinstance(obj, (np.floating,)):
        return _round_floats(float(obj), ndigits)
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, np.ndarray):
        return _round_floats(obj.tolist(), ndigits)
    if isinstance(obj, dict):
        return {k: _round_floats(v, ndigits) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_round_floats(v, ndigits) for v in obj]
    return obj


def save(n: int, fig: dict, ndigits: int = 4) -> None:
    """Write the figure JSON for demo `n` (PNGs come from render_pngs.mjs)."""
    fig = _round_floats(fig, ndigits)
    path_json = OUT / f"quiver-{n:02d}.json"
    path_json.write_text(json.dumps(fig, separators=(",", ":")))
    print(f"  saved quiver-{n:02d}.json ({path_json.stat().st_size // 1024} KB)")