rustwx_stormscope/ops_render.py

#!/usr/bin/env python
"""Build the StormScope operational web view from cycle .npz outputs.

Rendering is delegated to rust/stormscope_render, which uses rustwx-render for
the basemap, projection, palettes, and PNG writing. This script only watches
cycle manifests, invokes the renderer, assembles GIFs, and writes the HTML shell.
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
import traceback
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path

from PIL import Image

from .paths import default_exports_dir, default_outdir, default_site_dir, default_spatial_root, repo_root
from .wxstore_export import export_cycle as export_wxstore_cycle


@dataclass(frozen=True)
class Product:
    slug: str
    renderer_key: str
    field: str
    title: str
    duration_ms: int


@dataclass(frozen=True)
class Region:
    slug: str
    title: str
    west: float
    east: float
    south: float
    north: float
    pad_frac: float = 0.0


PRODUCTS = [
    Product("ir", "ir", "ch_ir_1035", "IR Brightness Temperature", 650),
    Product("refc", "refc", "refc", "Composite Reflectivity", 650),
]

REGIONS = {
    "conus": Region("conus", "CONUS", -126.0, -66.0, 23.0, 50.0, 0.0),
    "central_plains": Region("central_plains", "Central Plains", -106.5, -90.0, 31.5, 47.2, 0.0),
}


def parse_args() -> argparse.Namespace:
    root = repo_root()
    renderer_name = "stormscope_render.exe" if os.name == "nt" else "stormscope_render"
    wxstore_name = "stormscope_wxstore.exe" if os.name == "nt" else "stormscope_wxstore"
    ap = argparse.ArgumentParser(description="Render StormScope operational site")
    ap.add_argument("--outdir", default=str(default_outdir()))
    ap.add_argument("--site-dir", default=str(default_site_dir()))
    ap.add_argument("--renderer", default=str(root / "rust" / "stormscope_render" / "target" / "release" / renderer_name))
    ap.add_argument("--wxstore-exporter", default=str(root / "rust" / "stormscope_wxstore" / "target" / "release" / wxstore_name))
    ap.add_argument("--wxstore-spatial-root", default=str(default_spatial_root()))
    ap.add_argument("--wxstore-fields", default="",
                    help="comma-separated WXA field allowlist; default exports all StormScope 2D fields")
    ap.add_argument("--no-wxstore-export", action="store_true",
                    help="skip native rustwx WXA export")
    ap.add_argument("--visual-dir", default=str(default_exports_dir()))
    ap.add_argument("--export-dir", default="",
                    help="alias for --visual-dir; copies latest exported GIFs/PNGs there")
    ap.add_argument("--regions", default="conus,central_plains",
                    help="comma-separated region slugs; available: conus,central_plains")
    ap.add_argument("--loop", action="store_true")
    ap.add_argument("--interval", type=int, default=60, help="seconds between render scans")
    ap.add_argument("--width", type=int, default=1200)
    ap.add_argument("--height", type=int, default=820)
    ap.add_argument("--max-cycles-per-mode", type=int, default=2,
                    help="render only this many latest cycles per mode; 0 renders all")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()
    if args.export_dir and not args.visual_dir:
        args.visual_dir = args.export_dir
    args.region_list = parse_regions(args.regions)
    return args


def parse_regions(value: str) -> list[Region]:
    regions = []
    for raw in value.split(","):
        slug = raw.strip()
        if not slug:
            continue
        if slug not in REGIONS:
            raise ValueError(f"unknown region {slug!r}; available: {', '.join(REGIONS)}")
        regions.append(REGIONS[slug])
    return regions or [REGIONS["conus"]]


def read_json(path: Path) -> dict:
    return json.loads(path.read_text())


def display_mode(mode: str) -> str:
    return "nearcast" if mode == "forecast" else mode


def mode_title(mode: str) -> str:
    return display_mode(mode).capitalize()


def parse_dt(value: str) -> datetime:
    return datetime.fromisoformat(value.replace("Z", "+00:00"))


def dt_label(value: str) -> str:
    return parse_dt(value).strftime("%Y-%m-%d %H:%MZ")


def cycle_sort_key(meta: dict) -> datetime:
    return parse_dt(meta["init_time"])


def relpath(path: Path, root: Path) -> str:
    return path.relative_to(root).as_posix()


def should_update(target: Path, sources: list[Path], force: bool) -> bool:
    if force or not target.exists():
        return True
    try:
        if target.stat().st_size == 0:
            return True
    except OSError:
        return True
    target_mtime = target.stat().st_mtime
    return any(src.exists() and src.stat().st_mtime > target_mtime for src in sources)


def frame_number(filename: str, index: int) -> str:
    stem = Path(filename).stem
    tail = stem.rsplit("_", 1)[-1]
    if tail.startswith("f") and tail[1:].isdigit():
        return tail.upper()
    return f"F{index + 1:02d}"


def visible_frame_indices(meta: dict) -> list[int]:
    files = meta.get("files", [])
    if display_mode(meta["mode"]) != "nearcast":
        return list(range(len(files)))

    valid_times = meta.get("valid_times", [])
    cutoff = datetime.now(timezone.utc) - timedelta(minutes=15)
    selected = [
        i for i in range(len(files))
        if i < len(valid_times) and parse_dt(valid_times[i]) >= cutoff
    ]
    return selected or list(range(len(files)))


def render_png(renderer: Path, npz: Path, output: Path, product: Product, region: Region, meta: dict,
               index: int, width: int, height: int, force: bool) -> None:
    if not should_update(output, [npz, renderer], force):
        return
    valid = meta["valid_times"][index] if index < len(meta.get("valid_times", [])) else ""
    step_minutes = int(meta.get("step_minutes", 0))
    lead_minutes = step_minutes * (index + 1)
    if lead_minutes and lead_minutes % 60 == 0:
        lead = f"+{lead_minutes // 60} h"
    elif lead_minutes:
        lead = f"+{lead_minutes} min"
    else:
        lead = frame_number(npz.name, index)
    title = f"StormScope {mode_title(meta['mode'])} {region.title} {product.title}"
    subtitle_left = f"Init {dt_label(meta['init_time'])} {frame_number(npz.name, index)} ({lead})"
    if valid:
        subtitle_left += f" valid {dt_label(valid)}"
    subtitle_right = "rustwx-render / rustwx-stormscope"

    cmd = [
        str(renderer),
        "--input", str(npz),
        "--output", str(output),
        "--product", product.renderer_key,
        "--title", title,
        "--subtitle-left", subtitle_left,
        "--subtitle-right", subtitle_right,
        "--width", str(width),
        "--height", str(height),
        f"--west={region.west}",
        f"--east={region.east}",
        f"--south={region.south}",
        f"--north={region.north}",
        f"--bounds-pad-frac={region.pad_frac}",
    ]
    env = os.environ.copy()
    vendor_assets = repo_root() / "rust" / "vendor" / "rustwx" / "assets"
    if vendor_assets.exists():
        env.setdefault("RUSTWX_ASSETS_DIR", str(vendor_assets))
    subprocess.run(cmd, check=True, env=env)


def make_gif(frames: list[Path], gif_path: Path, duration_ms: int, force: bool) -> None:
    if not frames:
        return
    if not should_update(gif_path, frames, force):
        return
    gif_path.parent.mkdir(parents=True, exist_ok=True)
    images = []
    for frame in frames:
        try:
            if not frame.exists() or frame.stat().st_size == 0:
                frame.unlink(missing_ok=True)
                continue
            image = Image.open(frame)
            image.load()
            images.append(image.convert("RGB"))
            image.close()
        except Exception:
            frame.unlink(missing_ok=True)
    if not images:
        return
    try:
        images[0].save(
            gif_path,
            save_all=True,
            append_images=images[1:],
            duration=duration_ms,
            loop=0,
            optimize=False,
        )
    finally:
        for image in images:
            image.close()


def latest_meta_paths(meta_paths: list[Path], max_per_mode: int) -> list[Path]:
    if max_per_mode <= 0:
        return meta_paths
    buckets: dict[str, list[tuple[datetime, Path]]] = {}
    for path in meta_paths:
        try:
            meta = read_json(path)
            key = display_mode(meta["mode"])
            init = parse_dt(meta["init_time"])
        except Exception:
            continue
        buckets.setdefault(key, []).append((init, path))
    selected = []
    for entries in buckets.values():
        selected.extend(path for _, path in sorted(entries, reverse=True)[:max_per_mode])
    return sorted(set(selected))


def render_cycle(meta_path: Path, outdir: Path, site_dir: Path, renderer: Path,
                 regions: list[Region],
                 width: int, height: int, force: bool) -> dict:
    meta = read_json(meta_path)
    mode = display_mode(meta["mode"])
    cycle = meta["cycle"]
    products_out = []
    fields = set(meta.get("fields", []))

    for region in regions:
        cycle_dir = site_dir / "assets" / region.slug / mode / cycle
        for product in PRODUCTS:
            if product.field not in fields:
                continue
            frame_paths = []
            png_dir = cycle_dir / product.slug / "frames"
            files = meta.get("files", [])
            frame_indices = visible_frame_indices(meta)
            for index in frame_indices:
                filename = files[index]
                npz = outdir / filename
                if not npz.exists():
                    continue
                png = png_dir / f"{Path(filename).stem}_{region.slug}_{product.slug}.png"
                render_png(renderer, npz, png, product, region, meta, index, width, height, force)
                frame_paths.append(png)
            gif_path = cycle_dir / f"{product.slug}.gif"
            make_gif(frame_paths, gif_path, product.duration_ms, force)
            products_out.append({
                "slug": product.slug,
                "title": product.title,
                "region_slug": region.slug,
                "region_title": region.title,
                "gif": gif_path,
                "frames": frame_paths,
                "n_frames": len(frame_paths),
            })

    return {
        "mode": mode,
        "source_mode": meta["mode"],
        "cycle": cycle,
        "init_time": meta["init_time"],
        "valid_times": [
            meta.get("valid_times", [])[i]
            for i in visible_frame_indices(meta)
            if i < len(meta.get("valid_times", []))
        ],
        "step_minutes": meta.get("step_minutes"),
        "resolution": meta.get("resolution"),
        "products": products_out,
        "meta_path": meta_path,
    }


def latest_by_mode(cycles: list[dict]) -> list[dict]:
    latest = {}
    for cycle in cycles:
        key = cycle["mode"]
        if key not in latest or parse_dt(cycle["init_time"]) > parse_dt(latest[key]["init_time"]):
            latest[key] = cycle
    return [latest[key] for key in ("nowcast", "nearcast") if key in latest]


def write_site(site_dir: Path, outdir: Path, cycles: list[dict]) -> None:
    site_dir.mkdir(parents=True, exist_ok=True)
    now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%SZ")
    service_status = {}
    status_path = outdir / "status.json"
    if status_path.exists():
        try:
            service_status = read_json(status_path)
        except Exception:
            service_status = {"state": "status unreadable"}

    manifest_cycles = []
    for cycle in cycles:
        manifest_products = []
        for product in cycle["products"]:
            manifest_products.append({
                "slug": product["slug"],
                "title": product["title"],
                "region_slug": product["region_slug"],
                "region_title": product["region_title"],
                "gif": relpath(product["gif"], site_dir) if product["gif"].exists() else None,
                "frames": [relpath(frame, site_dir) for frame in product["frames"] if frame.exists()],
                "n_frames": product["n_frames"],
            })
        manifest_cycles.append({**{k: cycle[k] for k in (
            "mode", "source_mode", "cycle", "init_time", "valid_times", "step_minutes", "resolution"
        )}, "products": manifest_products})

    (site_dir / "manifest.json").write_text(json.dumps({
        "updated": now,
        "service_status": service_status,
        "cycles": manifest_cycles,
    }, indent=2))

    cards = []
    for cycle in latest_by_mode(cycles):
        product_cards = []
        for product in cycle["products"]:
            if not product["gif"].exists():
                continue
            gif_rel = relpath(product["gif"], site_dir)
            download_name = f"stormscope_{cycle['mode']}_{cycle['cycle']}_{product['region_slug']}_{product['slug']}.gif"
            product_cards.append(f"""
          <article class="product">
            <div class="product-head">
              <div>
                <h3>{html_escape(product["title"])}</h3>
                <span>{html_escape(product["region_title"])} | {product["n_frames"]} frames</span>
              </div>
              <a class="export" href="{gif_rel}" download="{html_escape(download_name)}">Export GIF</a>
            </div>
            <img src="{gif_rel}" alt="{html_escape(cycle["mode"])} {html_escape(product["region_title"])} {html_escape(product["title"])} animation">
          </article>""")
        valid_times = cycle.get("valid_times") or []
        valid_span = ""
        if valid_times:
            valid_span = f"{dt_label(valid_times[0])} to {dt_label(valid_times[-1])}"
        cards.append(f"""
      <section class="cycle {cycle["mode"]}">
        <div class="cycle-head">
          <div>
            <p class="kicker">{html_escape(cycle["mode"])}</p>
            <h2>{html_escape(cycle["resolution"])} cycle {html_escape(cycle["cycle"])}</h2>
          </div>
          <dl>
            <div><dt>Init</dt><dd>{dt_label(cycle["init_time"])}</dd></div>
            <div><dt>Valid</dt><dd>{html_escape(valid_span)}</dd></div>
            <div><dt>Step</dt><dd>{cycle["step_minutes"]} min</dd></div>
          </dl>
        </div>
        <div class="products">
          {"".join(product_cards) if product_cards else "<p>No rendered products yet.</p>"}
        </div>
      </section>""")

    state = service_status.get("state", "unknown")
    last_success = service_status.get("last_success", "")
    next_run = service_status.get("next_run", "")
    active_modes = {cycle["mode"] for cycle in latest_by_mode(cycles)}
    if active_modes == {"nowcast"}:
        page_subtitle = "Nowcast, refreshed from the latest completed cycle."
    elif active_modes == {"nearcast"}:
        page_subtitle = "GFS-conditioned nearcast, refreshed from the latest completed cycle."
    else:
        page_subtitle = "Nowcast and GFS-conditioned nearcast, refreshed from the latest completed cycle."
    body = "".join(cards) if cards else """
      <section class="empty">
        <h2>Waiting for StormScope cycles</h2>
        <p>The renderer is ready; the first products will appear after the service writes a cycle manifest.</p>
      </section>"""

    (site_dir / "index.html").write_text(f"""<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <meta http-equiv="refresh" content="60">
  <title>StormScope Operational</title>
  <style>
    :root {{
      color-scheme: light;
      --bg: #f6f7f4;
      --ink: #1d211f;
      --muted: #657069;
      --line: #d8ddd7;
      --surface: #ffffff;
      --green: #0f766e;
      --orange: #b45309;
      --red: #b91c1c;
    }}
    * {{ box-sizing: border-box; }}
    body {{
      margin: 0;
      background: var(--bg);
      color: var(--ink);
      font: 15px/1.45 system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
    }}
    header {{
      padding: 18px 22px 12px;
      border-bottom: 1px solid var(--line);
      background: var(--surface);
      display: flex;
      align-items: end;
      justify-content: space-between;
      gap: 18px;
      flex-wrap: wrap;
    }}
    h1, h2, h3, p {{ margin: 0; }}
    h1 {{ font-size: clamp(24px, 3vw, 34px); letter-spacing: 0; }}
    .status {{
      display: grid;
      grid-template-columns: repeat(3, max-content);
      gap: 8px 14px;
      color: var(--muted);
      font-size: 13px;
    }}
    .status strong {{ color: var(--ink); }}
    main {{ padding: 18px 22px 28px; display: grid; gap: 18px; }}
    .cycle {{
      background: var(--surface);
      border: 1px solid var(--line);
      border-top: 4px solid var(--green);
      border-radius: 8px;
      overflow: hidden;
    }}
    .cycle.nearcast {{ border-top-color: var(--orange); }}
    .cycle-head {{
      padding: 14px 16px;
      display: flex;
      align-items: start;
      justify-content: space-between;
      gap: 16px;
      border-bottom: 1px solid var(--line);
    }}
    .kicker {{
      text-transform: uppercase;
      color: var(--muted);
      font-size: 12px;
      font-weight: 700;
    }}
    h2 {{ font-size: 20px; letter-spacing: 0; }}
    dl {{
      margin: 0;
      display: grid;
      grid-template-columns: repeat(3, max-content);
      gap: 10px 16px;
      font-size: 13px;
      color: var(--muted);
    }}
    dt {{ font-weight: 700; color: var(--ink); }}
    dd {{ margin: 2px 0 0; }}
    .products {{
      padding: 14px;
      display: grid;
      grid-template-columns: repeat(auto-fit, minmax(min(100%, 520px), 1fr));
      gap: 14px;
    }}
    .product {{
      border: 1px solid var(--line);
      border-radius: 8px;
      background: #fbfcfa;
      overflow: hidden;
    }}
    .product-head {{
      min-height: 46px;
      padding: 10px 12px;
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 10px;
      border-bottom: 1px solid var(--line);
    }}
    h3 {{ font-size: 15px; letter-spacing: 0; }}
    .product-head span {{ color: var(--muted); font-size: 12px; white-space: nowrap; }}
    .export {{
      display: inline-flex;
      align-items: center;
      justify-content: center;
      min-height: 32px;
      padding: 0 11px;
      border: 1px solid var(--line);
      border-radius: 6px;
      color: var(--ink);
      background: #fff;
      font-size: 13px;
      font-weight: 700;
      text-decoration: none;
      white-space: nowrap;
    }}
    .export:hover {{ border-color: var(--green); color: var(--green); }}
    img {{ width: 100%; height: auto; display: block; background: #fff; }}
    .empty {{
      background: var(--surface);
      border: 1px solid var(--line);
      border-radius: 8px;
      padding: 18px;
    }}
    @media (max-width: 800px) {{
      header, main {{ padding-left: 12px; padding-right: 12px; }}
      .status, dl {{ grid-template-columns: 1fr; }}
      .cycle-head {{ display: grid; }}
    }}
  </style>
</head>
<body>
  <header>
    <div>
      <h1>StormScope Operational</h1>
      <p>{html_escape(page_subtitle)}</p>
    </div>
    <div class="status">
      <span>Service <strong>{html_escape(str(state))}</strong></span>
      <span>Updated <strong>{html_escape(now)}</strong></span>
      <span>Next <strong>{html_escape(str(next_run or last_success or "pending"))}</strong></span>
    </div>
  </header>
  <main>
{body}
  </main>
</body>
</html>
""")


def html_escape(value: str) -> str:
    return (
        value.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


def mirror_latest_visuals(site_dir: Path, cycles: list[dict], visual_dir: Path | None) -> None:
    if not visual_dir:
        return
    visual_dir.mkdir(parents=True, exist_ok=True)
    for cycle in latest_by_mode(cycles):
        for product in cycle["products"]:
            if product["gif"].exists():
                shutil.copy2(
                    product["gif"],
                    visual_dir / f"latest_{cycle['mode']}_{product['region_slug']}_{product['slug']}.gif",
                )
            frames = [frame for frame in product["frames"] if frame.exists()]
            if frames:
                shutil.copy2(
                    frames[-1],
                    visual_dir / f"latest_{cycle['mode']}_{product['region_slug']}_{product['slug']}_last.png",
                )
    index = site_dir / "index.html"
    if index.exists():
        shutil.copy2(index, visual_dir / "index.html")


def export_wxstore_cycles(cycles: list[dict], outdir: Path, args: argparse.Namespace) -> None:
    if args.no_wxstore_export:
        return
    spatial_root = Path(args.wxstore_spatial_root)
    fields = args.wxstore_fields or None
    for cycle in cycles:
        meta_path = Path(cycle["meta_path"])
        try:
            result = export_wxstore_cycle(
                meta_path,
                outdir,
                spatial_root,
                binary=args.wxstore_exporter,
                fields=fields,
                force=args.force,
            )
            if result.get("skipped_current"):
                continue
            print(
                f"[ops_render] WXA exported {result.get('model')} {result.get('cycle')} "
                f"products={result.get('products', 'unknown')}",
                flush=True,
            )
        except FileNotFoundError as exc:
            print(f"[ops_render] WXA export unavailable: {exc}", flush=True)
            return
        except Exception as exc:
            print(f"[ops_render] WXA export failed for {meta_path.name}: {exc}", flush=True)


def render_once(args: argparse.Namespace) -> list[dict]:
    outdir = Path(args.outdir)
    site_dir = Path(args.site_dir)
    renderer = Path(args.renderer)
    if not renderer.exists():
        raise FileNotFoundError(f"renderer binary not found: {renderer}")
    all_meta_paths = sorted(outdir.glob("stormscope_*_*_*_meta.json"))
    complete_meta_paths = []
    for path in all_meta_paths:
        try:
            if read_json(path).get("complete", True):
                complete_meta_paths.append(path)
        except Exception:
            continue
    meta_paths = latest_meta_paths(
        complete_meta_paths,
        args.max_cycles_per_mode,
    )
    cycles = [
        render_cycle(path, outdir, site_dir, renderer, args.region_list, args.width, args.height, args.force)
        for path in meta_paths
    ]
    export_wxstore_cycles(cycles, outdir, args)
    write_site(site_dir, outdir, cycles)
    visual_dir = Path(args.visual_dir) if args.visual_dir else None
    mirror_latest_visuals(site_dir, cycles, visual_dir)
    return cycles


def main() -> int:
    args = parse_args()
    while True:
        try:
            cycles = render_once(args)
            print(f"[ops_render] rendered/checked {len(cycles)} cycles at {datetime.now(timezone.utc).isoformat()}", flush=True)
        except Exception:
            traceback.print_exc()
            if not args.loop:
                return 1
        if not args.loop:
            return 0
        time.sleep(max(5, args.interval))


if __name__ == "__main__":
    raise SystemExit(main())