#!/usr/bin/env python3
"""
shipowiz.py — audio-reactive Shipow backdrop renderer.

Maps bass kick → subtle zoom pulse, highs → brightness flash.
Optional: cross-fade through several Shipow loops (sinusoidal blend).

Usage:
    tools/shipowiz.py -a MASTER.mp4 -s ~/Videos/Shipow/nb3.mp4 -o OUT.mp4
    tools/shipowiz.py -a take_five.flac -s nb1.mp4 nb3.mp4 nb7.mp4 -o OUT.mp4
    tools/shipowiz.py -a M.mp4 -s nb*.mp4 --opacity 0.35 --zoom 0.05 --flash 0.8

If -a has video, the rendered backdrop is screen-blended onto it.
If -a is audio-only, the backdrop becomes the video and audio is muxed.
"""
from __future__ import annotations
import argparse
import subprocess
import sys
import tempfile
import wave
from pathlib import Path

import numpy as np
from PIL import Image

FPS = 30
SR = 48000
W, H = 1920, 1080            # final composite resolution
BG_W, BG_H = 960, 540        # backdrop render resolution (upscaled at composite)


def run(cmd, **kw):
    return subprocess.run(cmd, check=True, **kw)


# ─── audio ─────────────────────────────────────────────────────

def extract_audio(src: Path, out_wav: Path):
    run(["ffmpeg", "-y", "-i", str(src), "-vn",
         "-ac", "1", "-ar", str(SR), "-f", "wav", str(out_wav)],
        capture_output=True)


def load_wav(path: Path) -> np.ndarray:
    with wave.open(str(path), "rb") as w:
        raw = w.readframes(w.getnframes())
    return np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0


def band_envelope(audio: np.ndarray, lo_hz: float, hi_hz: float) -> np.ndarray:
    """Per-video-frame band energy via overlapping STFT, normalized to 98th pct."""
    hop = SR // FPS                       # 1600 samples = one video frame
    win = 2048
    n = len(audio) // hop
    window = np.hanning(win).astype(np.float32)
    freqs = np.fft.rfftfreq(win, 1 / SR)
    mask = (freqs >= lo_hz) & (freqs <= hi_hz)
    env = np.zeros(n, dtype=np.float32)
    pad = np.zeros(win, dtype=np.float32)
    for i in range(n):
        s = i * hop
        chunk = audio[s:s + win]
        if len(chunk) < win:
            pad[:] = 0
            pad[:len(chunk)] = chunk
            chunk = pad
        spec = np.abs(np.fft.rfft(chunk * window))
        env[i] = spec[mask].mean()
    p = np.percentile(env, 98)
    if p > 1e-9:
        env = np.clip(env / p, 0, 1)
    return env


def follower(env: np.ndarray, attack: float, release: float) -> np.ndarray:
    """Asymmetric one-pole envelope follower. attack/release in 0..1."""
    out = np.zeros_like(env)
    prev = 0.0
    for i, x in enumerate(env):
        coef = attack if x > prev else release
        prev = prev + coef * (x - prev)
        out[i] = prev
    return out


# ─── video ─────────────────────────────────────────────────────

class FrameSource:
    """Infinite-looping rgb24 frame reader via an ffmpeg subprocess."""

    def __init__(self, path: Path, w: int, h: int, fps: int):
        self.proc = subprocess.Popen([
            "ffmpeg", "-stream_loop", "-1", "-i", str(path),
            "-vf", f"scale={w}:{h}:force_original_aspect_ratio=increase,"
                   f"crop={w}:{h}",
            "-f", "rawvideo", "-pix_fmt", "rgb24", "-r", str(fps),
            "-loglevel", "error", "-"
        ], stdout=subprocess.PIPE)
        self.size = w * h * 3
        self.w, self.h = w, h

    def read(self):
        buf = self.proc.stdout.read(self.size)
        if len(buf) < self.size:
            return None
        return np.frombuffer(buf, dtype=np.uint8).reshape(self.h, self.w, 3)

    def close(self):
        try:
            self.proc.stdout.close()
        except Exception:
            pass
        self.proc.terminate()
        try:
            self.proc.wait(timeout=2)
        except Exception:
            self.proc.kill()


class Looper:
    """Reads from N FrameSources, cross-fades sequentially with smoothstep.

    Reads only the currently-active source(s) — idle sources don't advance,
    they catch up when they become active (acceptable for ambient loops)."""

    def __init__(self, paths, w, h, period_seconds=12.0):
        self.sources = [FrameSource(p, w, h, FPS) for p in paths]
        self.period = max(1, int(period_seconds * FPS))
        self.i = 0

    def frame(self):
        n = len(self.sources)
        if n == 1:
            self.i += 1
            return self.sources[0].read()
        phase = (self.i / self.period) % n
        a_idx = int(phase) % n
        b_idx = (a_idx + 1) % n
        t = phase - int(phase)
        t = t * t * (3 - 2 * t)              # smoothstep
        self.i += 1
        a = self.sources[a_idx].read()
        if t < 1e-3:
            return a
        b = self.sources[b_idx].read()
        if a is None or b is None:
            return None
        # in-place blend without float64 promotion
        return (a.astype(np.float32) * (1 - t) + b.astype(np.float32) * t).astype(np.uint8)

    def close(self):
        for s in self.sources:
            s.close()


def zoom_frame(arr: np.ndarray, factor: float) -> np.ndarray:
    """Crop center, upscale back — zooms in by `factor` (>=1).

    BILINEAR is intentional: at 1.0-1.06× on monochrome content the difference
    vs LANCZOS is invisible and BILINEAR is ~3× faster."""
    if factor <= 1.002:
        return arr
    h, w = arr.shape[:2]
    cw = int(w / factor)
    ch = int(h / factor)
    x = (w - cw) // 2
    y = (h - ch) // 2
    crop = arr[y:y + ch, x:x + cw]
    return np.asarray(Image.fromarray(crop).resize((w, h), Image.BILINEAR))


def brighten(arr: np.ndarray, amount: float,
             gain: float = 0.18, lift: float = 42.0) -> np.ndarray:
    if amount <= 0:
        return arr
    out = arr.astype(np.float32) * (1.0 + gain * amount) + lift * amount
    return np.clip(out, 0, 255).astype(np.uint8)


def has_video_stream(path: Path) -> bool:
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=codec_type", "-of", "csv=p=0", str(path)],
        capture_output=True, text=True)
    return "video" in r.stdout


# ─── main ──────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("-a", "--audio", required=True,
                    help="audio source — accepts any media (audio or video)")
    ap.add_argument("-s", "--src", nargs="+", required=True,
                    help="one or more Shipow MP4s; multiple → cross-fade")
    ap.add_argument("-o", "--out", required=True)
    ap.add_argument("--bass-lo", type=float, default=40)
    ap.add_argument("--bass-hi", type=float, default=160)
    ap.add_argument("--high-lo", type=float, default=4000)
    ap.add_argument("--high-hi", type=float, default=12000)
    ap.add_argument("--zoom", type=float, default=0.06,
                    help="max zoom factor on bass peak (0.06 = +6%%)")
    ap.add_argument("--flash", type=float, default=1.0,
                    help="brightness flash multiplier on highs")
    ap.add_argument("--opacity", type=float, default=0.4,
                    help="bg blend strength when overlaying on code video")
    ap.add_argument("--morph-period", type=float, default=12.0,
                    help="seconds per cross-fade cycle when multi-src")
    ap.add_argument("--bg-only", action="store_true",
                    help="output the reactive backdrop alone (no composite)")
    args = ap.parse_args()

    src_audio = Path(args.audio).expanduser().resolve()
    paths = [Path(p).expanduser().resolve() for p in args.src]
    out = Path(args.out).expanduser().resolve()

    for p in [src_audio, *paths]:
        if not p.exists():
            sys.exit(f"missing: {p}")

    with tempfile.TemporaryDirectory() as td:
        td = Path(td)
        wav = td / "a.wav"
        print(f"[1/4] extracting audio → {wav.name}")
        extract_audio(src_audio, wav)
        audio = load_wav(wav)
        n_frames = len(audio) * FPS // SR
        print(f"      {n_frames} frames @ {FPS}fps  ({n_frames / FPS:.1f}s)")

        print("[2/4] computing FFT band envelopes")
        bass = follower(band_envelope(audio, args.bass_lo, args.bass_hi),
                        attack=0.65, release=0.10)
        high = follower(band_envelope(audio, args.high_lo, args.high_hi),
                        attack=0.85, release=0.28)
        print(f"      bass mean={bass.mean():.2f} peak={bass.max():.2f}  "
              f"high mean={high.mean():.2f} peak={high.max():.2f}")

        print(f"[3/4] rendering visual layer ({len(paths)} src"
              f"{' cross-fading' if len(paths) > 1 else ''}) at {BG_W}x{BG_H}")
        looper = Looper(paths, BG_W, BG_H, period_seconds=args.morph_period)
        bg_path = td / "bg.mp4"
        enc = subprocess.Popen([
            "ffmpeg", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24",
            "-s", f"{BG_W}x{BG_H}", "-r", str(FPS), "-i", "-",
            "-c:v", "libx264", "-preset", "ultrafast", "-crf", "18",
            "-pix_fmt", "yuv420p", "-fps_mode", "cfr", "-r", str(FPS),
            "-loglevel", "error", str(bg_path)
        ], stdin=subprocess.PIPE)

        try:
            n = min(n_frames, len(bass), len(high))
            for i in range(n):
                f = looper.frame()
                if f is None:
                    break
                f = zoom_frame(f, 1.0 + args.zoom * float(bass[i]))
                f = brighten(f, args.flash * float(high[i]))
                enc.stdin.write(f.tobytes())
                if i % (FPS * 10) == 0 and i > 0:
                    print(f"      {i}/{n} ({100 * i // n}%)")
        finally:
            try:
                enc.stdin.close()
            except Exception:
                pass
            enc.wait()
            looper.close()

        if args.bg_only:
            print(f"[4/4] muxing backdrop + audio")
            run(["ffmpeg", "-y", "-i", str(bg_path), "-i", str(src_audio),
                 "-map", "0:v", "-map", "1:a?",
                 "-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
                 "-shortest", "-movflags", "+faststart", str(out)])
        elif has_video_stream(src_audio):
            print(f"[4/4] screen-blending backdrop onto code (opacity {args.opacity})")
            # Convert both streams to planar RGB before blend — screen-mode on YUV
            # chroma pushes neutral grayscale off-center and tints output pink.
            run(["ffmpeg", "-y", "-i", str(src_audio), "-i", str(bg_path),
                 "-filter_complex",
                 f"[0:v]scale={W}:{H},format=gbrp[code];"
                 f"[1:v]scale={W}:{H}:flags=bilinear,format=gbrp[bg];"
                 f"[code][bg]blend=all_mode=screen:all_opacity={args.opacity},"
                 f"format=yuv420p[v]",
                 "-map", "[v]", "-map", "0:a?",
                 "-c:v", "libx264", "-preset", "medium", "-crf", "20",
                 "-pix_fmt", "yuv420p", "-fps_mode", "cfr", "-r", str(FPS),
                 "-c:a", "aac", "-b:a", "192k",
                 "-movflags", "+faststart", str(out)])
        else:
            print(f"[4/4] audio-only source — muxing backdrop + audio")
            run(["ffmpeg", "-y", "-i", str(bg_path), "-i", str(src_audio),
                 "-map", "0:v", "-map", "1:a",
                 "-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
                 "-shortest", "-movflags", "+faststart", str(out)])

    print(f"✓ {out}")


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