#!/usr/bin/env python3

"""
Holt ein Sonnen-Vollscheibenbild von Helioviewer zu einem Zeitpunkt.

Beispiel:
    ./helioviewer_grab.py 2026-08-12T18:18:00Z hmi_continuum
    ./helioviewer_grab.py --list
    ./helioviewer_grab.py --sources
"""

import json
import sys
import urllib.parse
import urllib.request

API = "https://api.helioviewer.org/v2"

# Kuratierte Kurzerklaerung der gebraeuchlichsten Quellen (sourceId -> Text).
# Nur eine Auswahl, keine vollstaendige Doku -- fuer alle IDs siehe --list.
SOURCE_INFO = {
    8: "AIA 94 - hot flare plasma (~6 MK)",
    9: "AIA 131 - flaring regions (~10 MK)",
    10: "AIA 171 - quiet corona, coronal loops (~0.8 MK)",
    11: "AIA 193 - corona, coronal holes (~1.2 MK)",
    12: "AIA 211 - active region corona (~2 MK)",
    13: "AIA 304 - chromosphere/transition region, prominences (~0.05 MK)",
    14: "AIA 335 - active region corona (~2.5 MK)",
    15: "AIA 1600 - upper photosphere/transition region, flare ribbons",
    16: "AIA 1700 - photosphere/temperature minimum",
    17: "AIA 4500 - near-visible continuum, sunspots",
    18: "HMI continuum - white-light photosphere: sunspots, granulation "
        "(highest full-disk resolution, 0.5\"/px)",
    19: "HMI magnetogram - photospheric magnetic field map",
    94: "GONG H-alpha - chromosphere: filaments, plage, flare ribbons "
        "(ground-based network)",
}


def call(endpoint, **params):
    """GET auf die Helioviewer-API. Gibt (content_type, bytes) zurueck."""
    url = f"{API}/{endpoint}/?" + urllib.parse.urlencode(params)
    with urllib.request.urlopen(url, timeout=120) as r:
        return r.headers.get_content_type(), r.read()


def sources():
    """Flache Liste aller Datenquellen: [(sourceId, layer_path, nickname), ...]

    Der Pfad durch den verschachtelten JSON-Baum IST die Layer-Kette,
    die takeScreenshot erwartet (z.B. SDO/HMI/HMI/continuum).
    """
    _, raw = call("getDataSources")
    out = []

    def walk(node, path):
        if not isinstance(node, dict):
            return
        if "sourceId" in node:
            out.append((node["sourceId"], path, node.get("nickname", "")))
            return
        for key, val in node.items():
            walk(val, path + [key])

    walk(json.loads(raw), [])
    return sorted(out)


def pick(needle):
    """Sucht eine Datenquelle anhand eines Substrings im Nickname/Pfad."""
    hits = [
        s
        for s in sources()
        if needle.lower().replace("_", " ")
        in (" ".join(s[1]) + " " + s[2]).lower().replace("_", " ")
    ]
    if not hits:
        sys.exit(f"no data source matches {needle!r} (see --list)")
    return hits[0]


def main():
    if "--list" in sys.argv:
        for sid, path, nick in sources():
            print(f"{sid:4d}  {'/'.join(path):40s}  {nick}")
        return

    if "--sources" in sys.argv:
        for sid, desc in sorted(SOURCE_INFO.items()):
            print(f"{sid:4d}  {desc}")
        return

    date = sys.argv[1] if len(sys.argv) > 1 else "2026-08-12T18:18:00Z"
    what = sys.argv[2] if len(sys.argv) > 2 else "continuum"

    sid, path, nick = pick(what)
    print(f"Source: [{sid}] {'/'.join(path)}  ({nick})")

    # Welches Bild liegt dem Wunschzeitpunkt am naechsten? Wichtig:
    # Helioviewer liefert IMMER etwas, auch wenn es Stunden daneben liegt.
    _, raw = call("getClosestImage", date=date, sourceId=sid)
    meta = json.loads(raw)
    print(f"requested: {date}\ndelivered: {meta['date']}")

    # getClosestImage liefert "YYYY-MM-DD HH:MM:SS" (Leerzeichen, kein
    # Zeitzonen-Suffix). takeScreenshot verlangt dagegen ISO 8601 mit "T"
    # und "Z", sonst antwortet die API mit HTTP 400.
    shot_date = meta["date"].replace(" ", "T") + "Z"

    layers = "[" + ",".join(path) + ",1,100]"
    ctype, png = call(
        "takeScreenshot",
        date=shot_date,
        imageScale=0.6, # Bogensekunden/Pixel -> ~3400 px Vollscheibe
        layers=layers,
        x1=-1100,
        y1=-1100,
        x2=1100,
        y2=1100, # Bogensekunden ab Scheibenmitte
        display="true",
        watermark="false",
    )
    if not ctype.startswith("image"):
        sys.exit(f"unexpected response ({ctype}): {png[:200]!r}")

    stamp = meta["date"].replace(":", "").replace("-", "").replace(" ", "_")
    name = f"{stamp}_{nick.replace(' ', '')}.png"
    with open(name, "wb") as f:
        f.write(png)
    print(f"-> {name}  ({len(png) / 1e6:.1f} MB)")


if __name__ == "__main__":
    main()