#!/usr/bin/env python
"""Stack the pristine, current-pipeline and new-pipeline streams side by side.

`bench.py render` writes one composite per timestamp per variant.  This turns three
such streams into a single video with the panes labelled, so the repair can be judged
against both the truth and what the pipeline does today:

    bench.py render --case C --variant pristine --out /tmp/streams
    bench.py render --case C --variant today    --out /tmp/streams
    bench.py render --case C --variant new      --out /tmp/streams
    make_comparison_video.py --streams /tmp/streams --satellite 16 --out compare.mp4

`render_real.py` streams on real data work the same way (its ``--filler``
directories are variants), with ``@root`` per pane, ``--labels`` to name each
detection and filler, and ``--periods`` to show sample weeks back to back.

The panes must stay in step, which is the whole reason this does not simply hand
ffmpeg three directories: a variant that produced no composite for a timestamp -- what
the current pipeline does past its gap limit -- would otherwise shorten that pane and
slide it out of alignment with the others.  Every stream is emitted against the same
canonical timeline, with black where a composite is missing, so frame N is the same
instant in all three panes.  Those black runs are not padding; they are what the
current pipeline actually shows.
"""

import argparse
import datetime
import os
import re
import subprocess
import sys

import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from ffmpeg_video import find_ffmpeg
from suvi import paths

LABELS = {
    "pristine": "PRISTINE (ground truth)",
    "today": "TODAY (hold-last + black frames)",
    "new": "OPTICAL FLOW (disc+header detection)",
    "optical_flow": "OPTICAL FLOW",
    "learned": "LEARNED (two-satellite aligned fusion)",
    "damaged": "INPUT (gaps and damage, unrepaired)",
}
_STAMP = re.compile(r"Composite-(\d+)\.jpg$")
#: What a ``--labels`` entry may contain.  The label lands inside a drawtext
#: option and a shell command line, where quotes, colons, backslashes and the
#: like would end the option or the command.
_LABEL = re.compile(r"[A-Za-z0-9 ._+()|/-]+")


def parse_labels(spec, count):
    """'a,b,...' -> one on-screen label per pane, or None for the defaults."""
    if spec is None:
        return None
    labels = tuple(label.strip() for label in spec.split(","))
    if len(labels) != count:
        raise SystemExit(f"--labels names {len(labels)} panes; there are {count}")
    for label in labels:
        if not _LABEL.fullmatch(label):
            raise SystemExit(f"label {label!r}: use letters, digits, spaces and ._+()|/-")
    return labels


def parse_panes(spec, default_satellite, default_root=None):
    """'variant[:satellite][@root],...' -> ((variant, satellite, root), ...).

    ``@root`` reads that pane from a different ``bench.py render`` output
    directory.  A new model only needs its own ``learned`` stream rendered; the
    damaged inputs and the optical-flow repair of the same case already exist
    elsewhere, and this share supports neither symlinks nor cheap copies of a
    few thousand multi-megabyte composites.
    """
    panes = []
    for entry in spec.split(","):
        entry = entry.strip()
        if not entry:
            continue
        entry, _, root = entry.partition("@")
        variant, _, satellite = entry.partition(":")
        if variant not in LABELS:
            raise SystemExit(f"unknown variant {variant!r}; known: {sorted(LABELS)}")
        panes.append((variant, int(satellite) if satellite else default_satellite,
                      root or default_root))
    return tuple(panes)


def pane_label(variant, satellite, root, default_root):
    """The on-screen label; a pane read from its own root names that root, so two
    models' outputs can never be mistaken for one another."""
    text = f"{LABELS[variant]} | GOES-{satellite}"
    if root and root != default_root:
        text += f" | {os.path.basename(os.path.normpath(root))}"
    return text


def filter_graph(panes, default_root, font, labels=None):
    """The ffmpeg filter graph that labels each pane and lays them out as [out].

    One pane is a single product stream rather than a comparison, so it keeps its
    full resolution and is not stacked: ffmpeg's stack filters need two inputs.
    `labels`, from :func:`parse_labels`, replaces the default pane labels.

    Labels sit bottom centre, in the empty corona below the disc.  At the top
    they covered the composite's own title and timestamp, which is how a frame
    seen in the video is found again.
    """
    scale = "scale=960:-2," if len(panes) > 1 else ""
    labelled = []
    for index, (variant, satellite, root) in enumerate(panes):
        text = (labels[index] if labels
                else pane_label(variant, satellite, root, default_root))
        labelled.append(
            f"[{index}:v]{scale}"
            f"drawtext=fontfile='{font}':text='{text}':"
            f"x=(w-text_w)/2:y=h-text_h-12:fontsize=22:fontcolor=yellow:"
            f"box=1:boxcolor=black@0.5[v{index}]"
        )
    taps = "".join(f"[v{i}]" for i in range(len(panes)))
    if len(panes) == 1:
        stack = f"{taps}null[out]"
    elif len(panes) == 4:
        # Row-major 2x2: panes[0] top-left, [1] top-right, [2] bottom-left,
        # [3] bottom-right.
        stack = f"{taps}xstack=inputs=4:layout=0_0|w0_0|0_h0|w0_h0[out]"
    else:
        stack = f"{taps}hstack=inputs={len(panes)}[out]"
    return ";".join(labelled) + ";" + stack


def parse_periods(spec):
    """'2024-01-01..2024-01-08,...' -> ((t0, t1), ...) in UTC seconds, ends
    exclusive, in time order.  Overlaps are refused: they would show the same
    slots twice."""
    periods = []
    for entry in spec.split(","):
        start, separator, end = entry.strip().partition("..")
        try:
            t0, t1 = (int(datetime.datetime.strptime(text, "%Y-%m-%d").replace(
                tzinfo=datetime.timezone.utc).timestamp()) for text in (start, end))
        except ValueError:
            separator = ""
        if not separator:
            raise SystemExit(f"period {entry!r}: expected YYYY-MM-DD..YYYY-MM-DD")
        if t1 <= t0:
            raise SystemExit(f"period {entry!r} ends before it starts")
        periods.append((t0, t1))
    periods.sort()
    for (_, end), (start, _) in zip(periods, periods[1:]):
        if start < end:
            raise SystemExit(f"periods overlap: {spec}")
    return tuple(periods)


def build_timeline(streams, periods=None):
    """The slots every pane is emitted against.

    By default every cadence slot from the first composite in any pane to the
    last.  With `periods`, only the slots inside them, back to back: a pane read
    from a whole-year render would otherwise stretch three sample weeks into a
    year of mostly black video.
    """
    if periods:
        return [when for t0, t1 in periods for when in range(t0, t1, paths.CADENCE)]
    everything = set()
    for frames in streams.values():
        everything.update(frames)
    if not everything:
        return []
    return list(range(min(everything), max(everything) + paths.CADENCE, paths.CADENCE))


def stream_frames(directory):
    """{timestamp: path} for one rendered stream."""
    frames = {}
    try:
        entries = os.listdir(directory)
    except OSError:
        return frames
    for entry in entries:
        match = _STAMP.search(entry)
        if match:
            frames[int(match.group(1))] = os.path.join(directory, entry)
    return frames


def encode_stream(ffmpeg, frames, timeline, out_path, framerate, crf):
    """Encode one pane, emitting a black frame wherever the variant has nothing.

    Reads every composite in the pane -- three panes of a week-long window is over
    7,000 files -- which is enough to fill the guest's inode cache on its own.  Keeping
    that below the mount's ceiling is the reclaim watchdog's job, not this loop's; see
    reclaim.py.
    """
    command = (
        f'{ffmpeg} -y -f image2pipe -framerate {framerate} -i - '
        f'-c:v libx264 -crf {crf} -preset veryfast -pix_fmt yuv420p "{out_path}"'
    )
    process = subprocess.Popen(command, shell=True, stdin=subprocess.PIPE,
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    blank = None
    present = 0
    for when in timeline:
        path = frames.get(when)
        if path:
            with open(path, "rb") as handle:
                process.stdin.write(handle.read())
            present += 1
        else:
            if blank is None:
                size = Image.open(next(iter(frames.values()))).size if frames else (1920, 1080)
                blank = Image.fromarray(np.zeros((size[1], size[0], 3), dtype="uint8"))
            blank.save(process.stdin, "jpeg", quality=95)
    process.stdin.close()
    process.wait()
    return present


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--streams", required=True, help="root written by bench.py render")
    parser.add_argument("--variants", default="damaged:18,new,damaged,learned",
                        help="comma-separated pane order, each "
                             "'variant[:satellite][@root]' (default satellite "
                             "from --satellite, default root --streams).  Four "
                             "panes make a 2x2 grid (row-major); any other count "
                             "stacks horizontally")
    parser.add_argument("--labels", default=None,
                        help="comma-separated on-screen label per pane, replacing "
                             "the defaults (letters, digits, spaces and ._+()|/-)")
    parser.add_argument("--periods", default=None,
                        help="comma-separated UTC date ranges START..END (end "
                             "exclusive) to show, back to back (default: everything "
                             "from the first composite to the last)")
    parser.add_argument("--satellite", type=int, default=16)
    parser.add_argument("--out", required=True)
    parser.add_argument("--framerate", type=int, default=60)
    parser.add_argument("--crf", type=int, default=18)
    parser.add_argument("--keep-panes", action="store_true",
                        help="keep the intermediate per-pane videos")
    args = parser.parse_args(argv)

    ffmpeg = find_ffmpeg()
    panes_wanted = parse_panes(args.variants, args.satellite, args.streams)
    labels = parse_labels(args.labels, len(panes_wanted))
    periods = parse_periods(args.periods) if args.periods else None
    streams = {}
    for key in panes_wanted:
        variant, satellite, root = key
        directory = os.path.join(root, variant, f"goes{satellite}")
        streams[key] = stream_frames(directory)
        print(f"  {variant:>8} g{satellite}: "
              f"{len(streams[key])} composites in {directory}")
    # One canonical timeline across every variant, so the panes stay in step.
    timeline = build_timeline(streams, periods)
    if not any(when in frames for frames in streams.values() for when in timeline):
        print("No composites found; run bench.py render first.")
        return 1
    print(f"  timeline: {len(timeline)} slots "
          f"({len(timeline) / args.framerate:.0f}s per pane at {args.framerate}fps)")

    panes = []
    for index, key in enumerate(panes_wanted):
        variant, satellite, _ = key
        pane = f"{os.path.splitext(args.out)[0]}.{index}-{variant}-g{satellite}.mp4"
        present = encode_stream(
            ffmpeg, streams[key], timeline, pane, args.framerate, args.crf
        )
        missing = len(timeline) - present
        print(f"  encoded {variant} g{satellite}: {present} frames"
              + (f", {missing} black ({missing / len(timeline):.0%})" if missing else ""))
        panes.append(pane)

    inputs = " ".join(f'-i "{p}"' for p in panes)
    font = os.path.join(os.path.dirname(os.path.abspath(__file__)), "OpenSans-Regular.ttf")
    graph = filter_graph(panes_wanted, args.streams, font, labels)
    command = (
        f'{ffmpeg} -y {inputs} -filter_complex "{graph}" -map "[out]" '
        f'-c:v libx264 -crf {args.crf} -preset veryfast -pix_fmt yuv420p "{args.out}"'
    )
    print("  stacking panes...")
    result = subprocess.run(command, shell=True, stdout=subprocess.DEVNULL,
                            stderr=subprocess.PIPE)
    if result.returncode != 0:
        print(result.stderr.decode()[-1500:])
        return 1

    if not args.keep_panes:
        for pane in panes:
            try:
                os.remove(pane)
            except OSError:
                pass
    size = os.path.getsize(args.out) / 1024 / 1024
    print(f"Wrote {args.out} ({size:.0f} MB)")
    return 0


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