#!/usr/bin/env python3
"""
mpc_to_orbitals.py — fetch comet orbital elements from the Minor Planet Center
and write them in the binary format used by the N.I.N.A. "Orbitals" plugin.

Background: the plugin downloads CometEls.txt from the MPC, parses it in memory
and stores the result as a GZip-compressed protobuf-net stream. The text file
itself is never saved and never read back from disk. So if the built-in download
does not work for you, you have to produce that binary file yourself.

Target:  %LOCALAPPDATA%\\NINA\\OrbitalElements\\MPC_CometElements.bin.gz
         (only read when the plugin's comet source is set to MPC)

Usage:
    python mpc_to_orbitals.py                 # download from MPC and write
    python mpc_to_orbitals.py --input C.txt   # use a local CometEls.txt
    python mpc_to_orbitals.py --verify        # read the written file back
    python mpc_to_orbitals.py --out X.bin.gz  # different target path

Runs on Windows, macOS and Linux. Standard library only, no dependencies.
"""

import argparse, gzip, math, os, struct, sys, urllib.request

MPC_URL = "https://www.minorplanetcenter.net/iau/MPCORB/CometEls.txt"
GM_SUN  = 1.32712440018e20          # m^3/s^2, same as GravitationalParameter.Sun
RAD     = math.pi / 180.0

# ---------------------------------------------------------------- protobuf ---

def _varint(n):
    out = bytearray()
    while True:
        b = n & 0x7F
        n >>= 7
        out.append(b | (0x80 if n else 0))
        if not n:
            return bytes(out)

def _f64(field, value):
    return bytes([(field << 3) | 1]) + struct.pack("<d", value)

def _bytes(field, payload):
    return bytes([(field << 3) | 2]) + _varint(len(payload)) + payload

def _gm(field, value):
    inner = _f64(1, value) if value != 0.0 else b""   # protobuf-net omits defaults
    return _bytes(field, inner)

def encode_element(name, epoch_jd, q_au, ecc, i_rad, w_rad, node_rad, tp_jd):
    """One OrbitalElements record as declared in Kepler.cs."""
    m  = _bytes(1, name.encode("utf-8"))
    m += _gm(2, GM_SUN)      # PrimaryGravitationalParameter
    m += _gm(3, 0.0)         # SecondaryGravitationalParameter
    m += _f64(4, epoch_jd)
    m += _f64(5, q_au)
    m += _f64(6, ecc)
    m += _f64(7, i_rad)
    m += _f64(8, w_rad)
    m += _f64(9, node_rad)
    m += _f64(10, tp_jd)
    # Fields 11 (M_MeanAnomalyAtEpoch) and 12 (a_SemiMajorAxis_au) are left out
    # (null), exactly like MPCCometElements.ToOrbitalElements() does.
    return _bytes(1, m)      # SerializeWithLengthPrefix(Base128, field number 1)

# -------------------------------------------------------------------- date ---

def julian_date(year, month, day, hour=0.0):
    """NOVAS julian_date(): integer day number plus fraction of day.

    Careful: C integer division truncates toward zero, Python's // floors.
    For (month-14)/12 that differs for months 1..13 and puts the result
    two days off.
    """
    a = math.trunc((month - 14) / 12)
    jd = (day - 32075
          + 1461 * (year + 4800 + a) // 4
          + 367 * (month - 2 - a * 12) // 12
          - 3 * ((year + 4900 + a) // 100) // 4)
    return jd - 0.5 + hour / 24.0

# ------------------------------------------------------------------ parser ---

COLS = dict(number=(0,4), otype=(4,5), prov=(5,14), tp_y=(14,19), tp_m=(19,22),
            tp_d=(22,30), q=(30,41), e=(41,51), w=(51,61), node=(61,71),
            i=(71,81), epoch=(81,89), name=(100,159))

def _num(s):
    s = s.strip()
    if not s or "*" in s:
        # The MPC occasionally ships asterisks instead of a value. The plugin
        # aborts the whole update on those (issue #12); we skip the record.
        raise ValueError(f"unusable value {s!r}")
    return float(s)

def parse_line(line):
    g = lambda k: line[COLS[k][0]:COLS[k][1]]
    name = g("name").strip()
    if not name:
        raise ValueError("no name")
    tp_y, tp_m = int(g("tp_y")), int(g("tp_m"))
    tp_d = _num(g("tp_d"))
    day  = int(tp_d)
    tp_jd = julian_date(tp_y, tp_m, day, (tp_d - day) * 24.0)

    ep = g("epoch").strip()
    epoch_jd = (julian_date(int(ep[0:4]), int(ep[4:6]), int(ep[6:8]))
                if len(ep) == 8 else float("nan"))

    # Angles go in as radians; a is derived from q and e by the plugin at runtime.
    return dict(name=name, epoch_jd=epoch_jd, q_au=_num(g("q")), ecc=_num(g("e")),
                i_rad=_num(g("i")) * RAD, w_rad=_num(g("w")) * RAD,
                node_rad=_num(g("node")) * RAD, tp_jd=tp_jd)

# ------------------------------------------------------------------ reader ---

def read_back(path, limit=5):
    """Minimal decoder, used by --verify to sanity-check what was written."""
    data = gzip.open(path, "rb").read()
    pos, n, shown = 0, 0, []
    while pos < len(data):
        assert data[pos] == 0x0A, f"unexpected tag at {pos}"
        pos += 1
        ln, shift = 0, 0
        while True:
            b = data[pos]; pos += 1
            ln |= (b & 0x7F) << shift; shift += 7
            if not b & 0x80: break
        msg, pos = data[pos:pos+ln], pos + ln
        n += 1
        if len(shown) < limit:
            p, rec = 0, {}
            while p < len(msg):
                tag = msg[p]; p += 1
                fld, wt = tag >> 3, tag & 7
                if wt == 1:
                    rec[fld] = struct.unpack("<d", msg[p:p+8])[0]; p += 8
                elif wt == 2:
                    l2, sh = 0, 0
                    while True:
                        b = msg[p]; p += 1
                        l2 |= (b & 0x7F) << sh; sh += 7
                        if not b & 0x80: break
                    val, p = msg[p:p+l2], p + l2
                    rec[fld] = val.decode("utf-8") if fld == 1 else val
            shown.append(rec)
    return n, shown

# -------------------------------------------------------------------- main ---

def default_out():
    """Default target on Windows only; elsewhere --out is required."""
    base = os.environ.get("LOCALAPPDATA")
    if not base:
        sys.exit("No %LOCALAPPDATA% found (not Windows?).\n"
                 "Please give a target, e.g.:\n"
                 "  --out ~/Desktop/MPC_CometElements.bin.gz\n"
                 "then copy the finished file to\n"
                 "  C:\\Users\\<user>\\AppData\\Local\\NINA\\OrbitalElements\\\n"
                 "without unpacking it.")
    return os.path.join(base, "NINA", "OrbitalElements", "MPC_CometElements.bin.gz")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", help="local CometEls.txt instead of downloading")
    ap.add_argument("--out", default=None, help="target path for the .bin.gz")
    ap.add_argument("--timeout", type=int, default=600,
                    help="download timeout in seconds (default 600; "
                         "the plugin itself is stuck at 100)")
    ap.add_argument("--verify", action="store_true",
                    help="read the written file back and show the first records")
    a = ap.parse_args()
    out = a.out or default_out()

    if a.verify:
        n, recs = read_back(out)
        print(f"{out}\n{n} records\n")
        for r in recs:
            print(f"  {r[1]:<34} q={r[5]:.6f} e={r[6]:.6f} "
                  f"i={math.degrees(r[7]):8.4f}deg tp={r[10]:.4f}")
        return

    if a.input:
        text = open(a.input, encoding="utf-8", errors="replace").read()
    else:
        print(f"Downloading {MPC_URL} (timeout {a.timeout}s) ...")
        req = urllib.request.Request(MPC_URL, headers={
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
            "Accept": "text/plain,*/*"})
        with urllib.request.urlopen(req, timeout=a.timeout) as r:
            text = r.read().decode("utf-8", errors="replace")
        print(f"  {len(text):,} characters received")

    ok, bad = [], 0
    for line in text.splitlines():
        if len(line) < 100:
            continue
        try:
            ok.append(parse_line(line))
        except Exception:
            bad += 1

    if not ok:
        sys.exit("No valid records found - is this really CometEls.txt?")

    # Write to a temp file first, then swap it in, so a crash cannot leave
    # the plugin with a half-written file.
    d = os.path.dirname(out)
    if d:
        os.makedirs(d, exist_ok=True)
    tmp = out + ".tmp"
    with gzip.open(tmp, "wb", compresslevel=9) as f:
        for e in ok:
            f.write(encode_element(e["name"], e["epoch_jd"], e["q_au"], e["ecc"],
                                   e["i_rad"], e["w_rad"], e["node_rad"], e["tp_jd"]))
    os.replace(tmp, out)

    print(f"\n{len(ok)} comets written, {bad} bad lines skipped")
    print(f"-> {out}")
    print("\nRestart N.I.N.A. (or toggle the source MPC->JPL->MPC), then check")
    print("that 'Comets (...)' shows the expected count.")

if __name__ == "__main__":
    main()