#!/usr/bin/env python3
"""Stitch reconnect fragments of one stream back into one recording (run by normalize-sweep.sh).

Every time a publisher drops and comes back, MediaMTX opens a new file under streams/<user>/<Title>/.live/,
so a night with a flaky uplink leaves one stream as a row of fragments: 20 min, 2 s, 14 min, 1 s, ... The
fragments are the SAME stream (same user, same title) and the file names carry the server arrival time to
the microsecond, so we know exactly where each piece sits on the clock. This joins them losslessly:

  * only when the stream is over: nothing in .live/ open (fuser) or modified in the last IDLE_MIN minutes
  * only fewer than MAX_PARTS fragments (default 500 — a lossy uplink splits a half-hour stream into ~180
    files, see normalize-sweep.sh): beyond that is a broken night, left for vgstuck.sh / the streaming page
  * a chain breaks where the gap between fragments is longer than MAX_GAP_S (default 10 min: a reconnect
    is seconds, a second session the same evening is not), or where the video/audio parameters change
    (resolution, codec, rate): those cannot be copy-concatenated, so each side is stitched on its own
  * ffmpeg concat demuxer with -c copy — no re-encode, the sweep's own pass (timecode, band removal,
    Opus->AAC) runs on the result as on any other recording
  * the gaps between fragments are COLLAPSED (the picture jumps), and recorded in <stem>.stitch.json which
    the sweep merges into the published sidecar: part files, each part's arrival time, its offset in the
    stitched file, and the wall-clock time missing at every join
  * the stitched file takes the FIRST fragment's name (so the arrival-time fallback for the timecode still
    points at the real start), the fragments move to trash/<night>/ for the usual 3-day undo
  * a fragment ffprobe cannot read is binned (reason=unreadable) rather than blocking its stream forever

Prints logfmt lines (component=sweep) for Kibana. Exit 0 unless the arguments are wrong.
    vgstitch.py <media root> [--max-parts N] [--max-gap-s S] [--idle-min M] [--dry-run]
"""
import argparse, json, os, re, subprocess, sys, time
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo

TZ = ZoneInfo("America/New_York")
NAME_RE = re.compile(r"(\d{4})-(\d{2})-(\d{2})_(\d{2})-(\d{2})-(\d{2})(?:-(\d{1,6}))?")

def vglog(level, **kv):
    parts = [f"ts={datetime.now().astimezone().isoformat(timespec='seconds')}", f"level={level}", "component=sweep"]
    for k, v in kv.items():
        v = str(v)
        parts.append(f'{k}="{v}"' if (" " in v or "/" in v or v == "") else f"{k}={v}")
    print(" ".join(parts), flush=True)

def arrival_ms(stem):
    m = NAME_RE.search(stem)
    if not m: return None
    y, mo, d, h, mi, s, us = m.groups()
    us = int((us or "0").ljust(6, "0"))
    return int(datetime(int(y), int(mo), int(d), int(h), int(mi), int(s), us, tzinfo=TZ).timestamp() * 1000)

def ffprobe(path):
    """(params tuple, duration_s) or None when unreadable."""
    try:
        out = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                              "stream=codec_type,codec_name,width,height,pix_fmt,sample_rate,channels:format=duration",
                              "-of", "json", str(path)], capture_output=True, text=True, timeout=120, check=True).stdout
        j = json.loads(out)
        v = [s for s in j["streams"] if s["codec_type"] == "video"]
        a = [s for s in j["streams"] if s["codec_type"] == "audio"]
        if not v: return None
        dur = float(j["format"].get("duration") or 0)
        if dur <= 0: return None
        params = (v[0].get("codec_name"), v[0].get("width"), v[0].get("height"), v[0].get("pix_fmt"),
                  tuple((s.get("codec_name"), s.get("sample_rate"), s.get("channels")) for s in a))
        return params, dur
    except Exception:
        return None

def busy(path):
    return subprocess.run(["fuser", "-s", str(path)], capture_output=True).returncode == 0

def to_trash(root, path, night):
    rel = path.relative_to(root)
    dst = root / "trash" / night / rel
    dst.parent.mkdir(parents=True, exist_ok=True)
    os.replace(path, dst)
    return dst

def stitch_dir(root, live, args):
    titledir = live.parent
    stream = f"{titledir.parent.name}/{titledir.name}"
    files = sorted(p for p in live.iterdir() if p.is_file() and p.suffix.lower() == ".mp4" and not p.name.endswith(".part"))
    if len(files) < 2: return
    now = time.time()
    for p in files:                                   # still streaming? leave the whole set alone
        if busy(p) or now - p.stat().st_mtime < args.idle_min * 60:
            vglog("info", event="stitch_wait", stream=stream, parts=len(files), reason="stream still going")
            return
    if len(files) >= args.max_parts:
        vglog("warn", event="stitch_skipped", stream=stream, parts=len(files), reason="too_many", max=args.max_parts)
        return
    night = datetime.now().strftime("%Y-%m-%d")
    parts = []
    for p in files:
        t = arrival_ms(p.stem)
        info = ffprobe(p)
        if t is None:
            vglog("warn", event="stitch_skipped", stream=stream, file=p, reason="name has no timestamp"); return
        if info is None:
            if not args.dry_run: dst = to_trash(root, p, night)
            else: dst = "(dry run)"
            vglog("warn", event="discarded", reason="unreadable", file=p, stream=stream, trash=dst)
            continue
        parts.append({"file": p, "arrival_ms": t, "params": info[0], "duration_ms": int(info[1] * 1000)})
    parts.sort(key=lambda x: x["arrival_ms"])
    # chains: consecutive, same parameters, gap under the limit
    chains, cur = [], []
    for part in parts:
        if cur:
            prev = cur[-1]
            gap = part["arrival_ms"] - (prev["arrival_ms"] + prev["duration_ms"])
            if part["params"] != prev["params"]:
                vglog("info", event="stitch_break", stream=stream, reason="params_changed", file=part["file"],
                      before=f"{prev['params'][1]}x{prev['params'][2]}", after=f"{part['params'][1]}x{part['params'][2]}")
                chains.append(cur); cur = []
            elif gap > args.max_gap_s * 1000:
                vglog("info", event="stitch_break", stream=stream, reason="gap", gap_s=gap // 1000, file=part["file"])
                chains.append(cur); cur = []
        cur.append(part)
    if cur: chains.append(cur)
    for chain in chains:
        if len(chain) < 2: continue
        first = chain[0]["file"]
        offset, joins, plist = 0, [], []
        for i, part in enumerate(chain):
            if i:
                prev = chain[i - 1]
                missing = part["arrival_ms"] - (prev["arrival_ms"] + prev["duration_ms"])
                joins.append({"at_s": round(offset / 1000, 3), "missing_ms": max(0, missing)})
            plist.append({"file": part["file"].name, "arrival_ms": part["arrival_ms"],
                          "duration_ms": part["duration_ms"], "offset_ms": offset})
            offset += part["duration_ms"]
        if args.dry_run:
            vglog("info", event="stitched", stream=stream, parts=len(chain), joins=len(joins), dry_run=1,
                  missing_ms=sum(j["missing_ms"] for j in joins), out=first); continue
        listing = live / f"{first.stem}.concat.txt"
        tmp = live / f"{first.stem}.stitch.part"
        listing.write_text("".join("file '%s'\n" % str(p["file"]).replace("'", r"'\''") for p in chain))
        run = subprocess.run(["nice", "-n", "15", "ionice", "-c", "3", "ffmpeg", "-nostdin", "-hide_banner", "-loglevel", "error", "-y",
                              "-f", "concat", "-safe", "0", "-i", str(listing), "-map", "0", "-c", "copy",
                              "-movflags", "+faststart", "-f", "mp4", str(tmp)], capture_output=True, text=True)
        listing.unlink(missing_ok=True)
        got = ffprobe(tmp) if run.returncode == 0 else None
        want = offset / 1000
        if got is None or abs(got[1] - want) > 2 + len(chain):
            tmp.unlink(missing_ok=True)
            why = (run.stderr.strip().splitlines() or ["no stderr"])[-1][:200]
            vglog("error", event="stitch_failed", stream=stream, parts=len(chain), rc=run.returncode,
                  want_s=int(want), got_s=int(got[1]) if got else "?", reason=why)
            continue                                    # fragments untouched: the sweep publishes them one by one
        last_mtime = max(p["file"].stat().st_mtime for p in chain)
        for p in chain: to_trash(root, p["file"], night)
        final = live / f"{first.stem}.mp4"
        os.replace(tmp, final)
        os.utime(final, (last_mtime, last_mtime))        # keep it "idle" so this very sweep publishes it
        meta = {"parts": plist, "joins": joins, "missing_ms": sum(j["missing_ms"] for j in joins),
                "trash": str(root / "trash" / night / live.relative_to(root)), "stitched_at": datetime.now(TZ).isoformat(timespec="seconds")}
        (live / f"{first.stem}.stitch.json").write_text(json.dumps(meta))
        vglog("info", event="stitched", stream=stream, parts=len(chain), joins=len(joins),
              missing_ms=meta["missing_ms"], seconds=int(want), out=final, trash=meta["trash"])

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("root")
    ap.add_argument("--max-parts", type=int, default=500)
    ap.add_argument("--max-gap-s", type=int, default=600)
    ap.add_argument("--idle-min", type=int, default=3)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()
    root = Path(args.root)
    streams = root / "streams"
    if not streams.is_dir(): return 0
    for live in sorted(streams.glob("*/*/.live")):
        if live.is_dir():
            try: stitch_dir(root, live, args)
            except Exception as error:
                vglog("error", event="stitch_failed", dir=live, reason=str(error)[:200])
    return 0

if __name__ == "__main__":
    sys.exit(main())