#!/usr/bin/env python3
"""
ffmpeg-worker v3: a small, fixed-purpose video assembler for n8n.

It accepts ONLY these calls (no arbitrary commands or ffmpeg arguments):
  GET  /health
  POST /prepare    {"slug": "my-video"}   creates <slug>/shots, <slug>/audio
  POST /assemble   {"slug": "my-video", "deliver": false}
  POST /trim       {"input": "shorts/inbox/talk.mp4", "outDir": "shorts/out",
                    "clips": [{"start": 12.5, "end": 41.0, "name": "talk-short-01"}],
                    "layout": "original" | "vertical", "overwrite": false}
                   cuts clips out of a video. All paths are relative to WORK_DIR.

It reads  <WORK_DIR>/<slug>/manifest.json:
  {
    "width": 1920, "height": 1080, "fps": 30,

    # all optional, new in v2 (leave them out and the video renders like v1):
    "transition": {"type": "fade", "duration": 0.5},   # crossfade between different visuals
    "motion": true,                                      # slow zoom on still images
    "cardBackground": "assets/title-bg.jpg",             # under WORK_DIR, used by card segments

    "segments": [
      {"image": "shots/01.png", "audio": "audio/01.wav", "caption": "optional on-screen text"},
      {"card": {"kicker": "Step 2 of 5", "title": "Extract the text", "subtitle": "Apache Tika"},
       "audio": "audio/02.wav"},                        # or "duration": 3 for a silent card
      ...
    ]
  }
and writes <WORK_DIR>/<slug>/final.mp4.

Consecutive segments that show the same image are joined into one scene, so the
picture only changes (and only crossfades) when the visual actually changes.
Captions still appear at the start of the segment they belong to.

If deliver=true it also copies the video (and thumbnail.png/.jpg if present)
into QUEUE_DIR as <slug>.mp4 / <slug>.png, which your YouTube upload
workflow watches. It refuses to overwrite files already in the queue,
because the upload workflow re-processes any file that is added again.
"""
import json
import os
import re
import shutil
import subprocess
import tempfile
import textwrap
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

WORK = Path(os.environ.get("WORK_DIR", "/work")).resolve()
QUEUE = Path(os.environ.get("QUEUE_DIR", "/queue")).resolve()
FONT = os.environ.get("FONT_FILE", "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf")
FONT_LIGHT = os.environ.get("FONT_FILE_LIGHT", "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf")
PORT = int(os.environ.get("PORT", "8080"))
PAD_SECONDS = 0.4      # short pause after each narrated segment
ZOOM_AMOUNT = 0.05     # slow zoom: scale goes from 1.00 to 1.05 over a scene
ACCENT = "0x5fd3ff"    # kicker colour on title cards (matches the cyan in the background)
TRANSITIONS = {"fade", "fadeblack", "dissolve", "smoothleft", "smoothright", "slideleft", "slideright", "wipeleft"}
SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$")
CLIP_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,79}$")
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".webm", ".m4v"}
MAX_CLIPS = 10
MIN_CLIP_SEC, MAX_CLIP_SEC = 1.0, 180.0
LOCK = threading.Lock()


def log(msg):
    print(msg, flush=True)


def safe_path(base, rel):
    """Resolve rel inside base; refuse anything that escapes the folder."""
    p = (base / str(rel)).resolve()
    if p != base and base not in p.parents:
        raise ValueError("path escapes allowed folder: " + str(rel))
    return p


def run(cmd):
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        raise RuntimeError(cmd[0] + " failed: " + r.stderr[-1500:])
    return r.stdout


def probe_duration(path):
    out = run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
               "-of", "default=nw=1:nk=1", str(path)])
    return float(out.strip())


# ---------- audio ----------

def normalize_audio(src, dst):
    """Narration clip -> 48 kHz stereo PCM with the short pause added at the end."""
    run(["ffmpeg", "-y", "-i", str(src), "-af", "apad=pad_dur=%.2f" % PAD_SECONDS,
         "-ar", "48000", "-ac", "2", "-c:a", "pcm_s16le", str(dst)])
    return probe_duration(dst)


def make_silence(dst, seconds):
    run(["ffmpeg", "-y", "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo", "-t", "%.3f" % seconds,
         "-c:a", "pcm_s16le", str(dst)])
    return probe_duration(dst)


# ---------- title cards ----------

def write_text(tmp, name, text):
    f = tmp / name
    f.write_text(text, encoding="utf-8")
    return f


def render_card(card, out_png, bg, tmp, width, height, idx):
    """Background image with centred kicker / title / subtitle, as one PNG."""
    if bg is not None:
        inputs = ["-i", str(bg)]
        chain = "[0:v]scale=%d:%d:force_original_aspect_ratio=increase,crop=%d:%d,setsar=1" % (width, height, width, height)
    else:
        inputs = ["-f", "lavfi", "-i", "color=c=0x0d1117:s=%dx%d" % (width, height)]
        chain = "[0:v]setsar=1"

    kicker = str(card.get("kicker") or "").strip()[:60]
    title = str(card.get("title") or "").strip()[:120]
    subtitle = str(card.get("subtitle") or "").strip()[:120]
    if not title:
        raise ValueError("a card segment needs a title")

    title_size = max(40, height // 13)
    kicker_size = max(24, height // 30)
    sub_size = max(26, height // 26)
    lines = []  # (text, font, size, colour)
    if kicker:
        lines.append((kicker.upper(), FONT, kicker_size, ACCENT))
    for t in textwrap.wrap(title, width=max(12, int(width * 0.8 / (title_size * 0.6)))) or [title]:
        lines.append((t, FONT, title_size, "white"))
    if subtitle:
        for t in textwrap.wrap(subtitle, width=max(16, int(width * 0.8 / (sub_size * 0.55)))):
            lines.append((t, FONT_LIGHT, sub_size, "0xc9d1d9"))

    gap = height // 60
    block = sum(int(s * 1.25) for _, _, s, _ in lines) + gap * (len(lines) - 1)
    y = (height - block) // 2
    # soft dark panel behind the text so it stays readable on any background
    panel_h = block + height // 9
    chain += ",drawbox=x=0:y=%d:w=iw:h=%d:color=black@0.45:t=fill" % (max(0, y - height // 18), panel_h)
    for n, (text, font, size, colour) in enumerate(lines):
        f = write_text(tmp, "card%03d_%02d.txt" % (idx, n), text)
        chain += (",drawtext=fontfile=%s:textfile=%s:expansion=none:fontcolor=%s:fontsize=%d:"
                  "x=(w-text_w)/2:y=%d" % (font, f, colour, size, y))
        y += int(size * 1.25) + gap
    chain += "[v]"
    run(["ffmpeg", "-y"] + inputs + ["-filter_complex", chain, "-map", "[v]", "-frames:v", "1", str(out_png)])


# ---------- scenes ----------

def render_scene(idx, scene, tmp, width, height, fps, seconds, motion):
    """One still visual held for `seconds`, with timed captions and optional slow zoom. No audio."""
    out = tmp / ("scene%03d.mp4" % idx)
    # Compose the still frame ONCE (the blur is expensive), then animate that single image.
    # With motion on, it is composed at 2x so the slow zoom moves smoothly.
    k = 2 if motion else 1
    W, H = width * k, height * k
    frame = tmp / ("frame%03d.png" % idx)
    if scene["fill"]:
        comp = "[0:v]scale=%d:%d:force_original_aspect_ratio=increase,crop=%d:%d,setsar=1[v]" % (W, H, W, H)
    else:
        # Blurred, cropped copy of the image fills the frame; the sharp image sits on top, fitted.
        comp = (
            "[0:v]split[a][b];"
            "[a]scale=%d:%d:force_original_aspect_ratio=increase,crop=%d:%d,boxblur=%d:5[bg];"
            "[b]scale=%d:%d:force_original_aspect_ratio=decrease[fg];"
            "[bg][fg]overlay=(W-w)/2:(H-h)/2,setsar=1[v]"
            % (W, H, W, H, 30 * k, W - 120 * k, H - 120 * k)
        )
    run(["ffmpeg", "-y", "-i", str(scene["image"]), "-filter_complex", comp, "-map", "[v]",
         "-frames:v", "1", str(frame)])

    if motion:
        frames = max(1, int(round(seconds * fps)))
        chain = ("[0:v]zoompan=z='1+%.4f*on/%d':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)'"
                 ":d=%d:s=%dx%d:fps=%d" % (ZOOM_AMOUNT, frames, frames, width, height, fps))
        inputs = ["-i", str(frame)]
    else:
        chain = "[0:v]null"
        inputs = ["-loop", "1", "-framerate", str(fps), "-t", "%.3f" % seconds, "-i", str(frame)]
    for n, (text, start, end) in enumerate(scene["captions"]):
        f = write_text(tmp, "cap%03d_%02d.txt" % (idx, n), text)
        chain += (
            ",drawtext=fontfile=%s:textfile=%s:expansion=none:fontcolor=white:fontsize=%d:"
            "box=1:boxcolor=black@0.6:boxborderw=18:x=(w-text_w)/2:y=h-%d:enable='between(t,%.3f,%.3f)'"
            % (FONT, f, max(28, height // 24), height // 9, start, end)
        )
    chain += ",format=yuv420p[v]"
    run([
        "ffmpeg", "-y"] + inputs + [
        "-filter_complex", chain, "-map", "[v]",
        "-c:v", "libx264", "-preset", "veryfast", "-crf", "18", "-r", str(fps), "-an",
        "-t", "%.3f" % seconds, str(out),
    ])
    return out


def assemble(slug, deliver):
    if not SLUG_RE.match(slug or ""):
        raise ValueError("slug must match ^[a-z0-9][a-z0-9_-]{0,63}$")
    proj = safe_path(WORK, slug)
    manifest_path = proj / "manifest.json"
    if not manifest_path.is_file():
        raise FileNotFoundError("manifest.json not found for slug: " + slug)
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))

    width = int(manifest.get("width", 1920))
    height = int(manifest.get("height", 1080))
    fps = int(manifest.get("fps", 30))
    if not (640 <= width <= 3840 and 360 <= height <= 2160 and width % 2 == 0 and height % 2 == 0):
        raise ValueError("width/height must be even and between 640x360 and 3840x2160")
    if not (12 <= fps <= 60):
        raise ValueError("fps must be between 12 and 60")
    segments = manifest.get("segments") or []
    if not segments or len(segments) > 200:
        raise ValueError("manifest needs between 1 and 200 segments")

    tr = manifest.get("transition") or {}
    tr_type = str(tr.get("type", "fade"))
    tr_dur = float(tr.get("duration", 0) or 0)
    if tr_type not in TRANSITIONS:
        raise ValueError("transition type must be one of: " + ", ".join(sorted(TRANSITIONS)))
    if not (0 <= tr_dur <= 2):
        raise ValueError("transition duration must be between 0 and 2 seconds")
    motion = bool(manifest.get("motion", False))
    bg = None
    if manifest.get("cardBackground"):
        bg = safe_path(WORK, manifest["cardBackground"])
        if not bg.is_file():
            raise FileNotFoundError("card background not found: " + str(manifest["cardBackground"]))

    final = proj / "final.mp4"
    with tempfile.TemporaryDirectory(dir=str(proj)) as tmpname:
        tmp = Path(tmpname)

        # 1. audio for every segment, and its visual
        wavs, scenes = [], []
        t = 0.0
        for i, seg in enumerate(segments):
            log("[%s] segment %d/%d" % (slug, i + 1, len(segments)))
            wav = tmp / ("a%03d.wav" % i)
            if seg.get("audio"):
                aud = safe_path(proj, seg["audio"])
                if not aud.is_file():
                    raise FileNotFoundError("missing audio: " + str(seg["audio"]))
                dur = normalize_audio(aud, wav)
            else:
                secs = float(seg.get("duration", 0) or 0)
                if not (0.5 <= secs <= 30):
                    raise ValueError("segment %d has no audio, so it needs a duration of 0.5 to 30 seconds" % (i + 1))
                dur = make_silence(wav, secs)
            wavs.append(wav)

            if seg.get("card"):
                png = tmp / ("card%03d.png" % i)
                render_card(seg["card"], png, bg, tmp, width, height, i)
                key, image, fill = ("card", i), png, True
            else:
                image = safe_path(proj, seg.get("image", ""))
                if not image.is_file():
                    raise FileNotFoundError("missing image: " + str(seg.get("image")))
                key, fill = ("image", str(image)), False

            if scenes and scenes[-1]["key"] == key and not seg.get("newScene"):
                scene = scenes[-1]
            else:
                scene = {"key": key, "image": image, "fill": fill, "start": t, "audio": 0.0, "captions": []}
                scenes.append(scene)
            caption = str(seg.get("caption") or "").strip()
            if caption and not seg.get("card"):
                local = t - scene["start"]
                scene["captions"].append((caption, local, local + dur))
            scene["audio"] += dur
            t += dur
        total = t

        # 2. one video per scene. Every scene but the last is held tr_dur longer, so the
        #    crossfade overlap is paid for and the picture stays in sync with the narration.
        use_tr = tr_dur > 0 and len(scenes) > 1
        parts = []
        for k, scene in enumerate(scenes):
            log("[%s] scene %d/%d" % (slug, k + 1, len(scenes)))
            hold = scene["audio"] + (tr_dur if use_tr and k < len(scenes) - 1 else 0)
            parts.append(render_scene(k, scene, tmp, width, height, fps, hold, motion))

        # 3. join the narration into one track
        alist = write_text(tmp, "audio.txt", "".join("file '%s'\n" % w.name for w in wavs))
        full_wav = tmp / "narration.wav"
        run(["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(alist), "-c", "copy", str(full_wav)])

        # 4. join the scenes (crossfading between them) and add the narration
        cmd = ["ffmpeg", "-y"]
        for p in parts:
            cmd += ["-i", str(p)]
        cmd += ["-i", str(full_wav)]
        if use_tr:
            graph, prev, offset = [], "[0:v]", 0.0
            for k in range(1, len(parts)):
                offset += scenes[k - 1]["audio"]
                label = "[x%d]" % k
                graph.append("%s[%d:v]xfade=transition=%s:duration=%.3f:offset=%.3f%s"
                             % (prev, k, tr_type, tr_dur, offset, label))
                prev = label
            vout = prev
        else:
            graph = ["".join("[%d:v]" % k for k in range(len(parts))) + "concat=n=%d:v=1:a=0[cv]" % len(parts)]
            vout = "[cv]"
        part_out = tmp / "final.mp4"
        cmd += ["-filter_complex", ";".join(graph), "-map", vout, "-map", "%d:a" % len(parts),
                "-c:v", "libx264", "-preset", "veryfast", "-crf", "20", "-r", str(fps), "-pix_fmt", "yuv420p",
                "-c:a", "aac", "-b:a", "160k", "-t", "%.3f" % total, "-movflags", "+faststart", str(part_out)]
        run(cmd)
        os.replace(str(part_out), str(final))

    result = {"ok": True, "slug": slug, "output": str(final), "segments": len(segments),
              "scenes": len(scenes), "duration_sec": round(total, 2), "delivered": []}

    if deliver:
        dest_video = QUEUE / (slug + ".mp4")
        if dest_video.exists():
            raise FileExistsError("already in queue (would re-trigger upload): " + dest_video.name)
        # Thumbnail first, so it is in place before the upload workflow looks for it
        for ext in ("png", "jpg", "jpeg"):
            thumb = proj / ("thumbnail." + ext)
            if thumb.is_file():
                dest_thumb = QUEUE / (slug + "." + ext)
                if not dest_thumb.exists():
                    shutil.copyfile(str(thumb), str(dest_thumb))
                    result["delivered"].append(dest_thumb.name)
                break
        # Copy under a name the upload workflow ignores, then rename in one step
        part_file = QUEUE / ("." + slug + ".mp4.part")
        shutil.copyfile(str(final), str(part_file))
        os.replace(str(part_file), str(dest_video))
        result["delivered"].append(dest_video.name)
    return result


def prepare(slug):
    """Create the project folders n8n will write into (n8n's file node cannot mkdir)."""
    if not SLUG_RE.match(slug or ""):
        raise ValueError("slug must match ^[a-z0-9][a-z0-9_-]{0,63}$")
    proj = safe_path(WORK, slug)
    created = []
    for sub in ("shots", "audio"):
        d = proj / sub
        d.mkdir(parents=True, exist_ok=True)
        created.append(str(d))
    return {"ok": True, "slug": slug, "folders": created}


def has_audio(path):
    out = run(["ffprobe", "-v", "error", "-select_streams", "a", "-show_entries", "stream=index",
               "-of", "csv=p=0", str(path)])
    return bool(out.strip())


def trim(data):
    """Cut up to MAX_CLIPS clips out of one video. Re-encodes, so cuts are frame-accurate."""
    src = safe_path(WORK, str(data.get("input", "")))
    if src == WORK or not src.is_file():
        raise FileNotFoundError("input video not found: " + str(data.get("input")))
    if src.suffix.lower() not in VIDEO_EXTS:
        raise ValueError("input must be one of: " + ", ".join(sorted(VIDEO_EXTS)))
    out_dir = safe_path(WORK, str(data.get("outDir", "")))
    if out_dir == WORK:
        raise ValueError("outDir must be a folder inside the work folder, e.g. shorts/out")
    layout = str(data.get("layout") or "original")
    if layout not in ("original", "vertical"):
        raise ValueError("layout must be original or vertical")
    overwrite = bool(data.get("overwrite", False))
    clips = data.get("clips")
    if not isinstance(clips, list) or not 1 <= len(clips) <= MAX_CLIPS:
        raise ValueError("clips must be a list of 1 to %d items" % MAX_CLIPS)

    total = probe_duration(src)
    audio = has_audio(src)
    plan, names = [], set()
    for i, c in enumerate(clips):
        if not isinstance(c, dict):
            raise ValueError("clip %d must be an object with start and end" % (i + 1))
        try:
            start, end = float(c.get("start")), float(c.get("end"))
        except (TypeError, ValueError):
            raise ValueError("clip %d needs numeric start and end (seconds)" % (i + 1))
        if start < 0 or end <= start:
            raise ValueError("clip %d: end must be after start, and start at least 0" % (i + 1))
        if end > total + 0.5:
            raise ValueError("clip %d ends at %.2fs but the video is only %.2fs long" % (i + 1, end, total))
        end = min(end, total)
        dur = end - start
        if not MIN_CLIP_SEC <= dur <= MAX_CLIP_SEC:
            raise ValueError("clip %d is %.2fs; clips must be %d to %d seconds" % (i + 1, dur, MIN_CLIP_SEC, MAX_CLIP_SEC))
        name = str(c.get("name") or "%s-short-%02d" % (re.sub(r"[^A-Za-z0-9_-]+", "-", src.stem)[:60].strip("-") or "clip", i + 1))
        if not CLIP_NAME_RE.match(name):
            raise ValueError("clip %d name may only use letters, digits, - and _ (max 80)" % (i + 1))
        if name in names:
            raise ValueError("two clips are named " + name)
        names.add(name)
        dest = out_dir / (name + ".mp4")
        if dest.exists() and not overwrite:
            raise FileExistsError("clip already exists (set overwrite to replace it): " + str(dest.relative_to(WORK)))
        plan.append((start, dur, dest))

    if layout == "vertical":
        vf = ("split[a][b];[a]scale=1080:1920:force_original_aspect_ratio=increase,crop=1080:1920,boxblur=30:5[bg];"
              "[b]scale=1080:1920:force_original_aspect_ratio=decrease[fg];"
              "[bg][fg]overlay=(W-w)/2:(H-h)/2,setsar=1,format=yuv420p")
    else:
        vf = "scale=trunc(iw/2)*2:trunc(ih/2)*2,setsar=1,format=yuv420p"

    out_dir.mkdir(parents=True, exist_ok=True)
    made = []
    for n, (start, dur, dest) in enumerate(plan):
        log("[trim] %s clip %d/%d %.2f+%.2fs" % (src.name, n + 1, len(plan), start, dur))
        part = dest.with_name("." + dest.name + ".part.mp4")
        cmd = ["ffmpeg", "-y", "-ss", "%.3f" % start, "-i", str(src), "-t", "%.3f" % dur,
               "-filter_complex", "[0:v:0]" + vf + "[v]", "-map", "[v]"]
        if audio:
            cmd += ["-map", "0:a:0", "-c:a", "aac", "-b:a", "160k", "-ac", "2"]
        cmd += ["-c:v", "libx264", "-preset", "veryfast", "-crf", "20", "-movflags", "+faststart", str(part)]
        try:
            run(cmd)
            os.replace(str(part), str(dest))
        finally:
            if part.exists():
                part.unlink()
        made.append({"name": dest.stem, "path": str(dest.relative_to(WORK)), "start": round(start, 3),
                     "end": round(start + dur, 3), "duration_sec": round(probe_duration(dest), 2)})
    return {"ok": True, "input": str(src.relative_to(WORK)), "layout": layout,
            "source_duration_sec": round(total, 2), "clips": made}


class Handler(BaseHTTPRequestHandler):
    def _send(self, code, payload):
        body = json.dumps(payload).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path == "/health":
            self._send(200, {"ok": True, "version": 3})
        else:
            self._send(404, {"ok": False, "error": "not found"})

    def do_POST(self):
        if self.path not in ("/assemble", "/prepare", "/trim"):
            self._send(404, {"ok": False, "error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            data = json.loads(self.rfile.read(length) or b"{}")
        except Exception:
            self._send(400, {"ok": False, "error": "invalid JSON body"})
            return
        if not LOCK.acquire(blocking=False):
            self._send(409, {"ok": False, "error": "another render is running"})
            return
        try:
            if self.path == "/prepare":
                self._send(200, prepare(str(data.get("slug", ""))))
            elif self.path == "/trim":
                self._send(200, trim(data))
            else:
                self._send(200, assemble(str(data.get("slug", "")), bool(data.get("deliver", False))))
        except (ValueError, FileNotFoundError, FileExistsError, KeyError) as e:
            self._send(400, {"ok": False, "error": str(e)})
        except Exception as e:  # ffmpeg failures etc.
            log("ERROR: " + str(e))
            self._send(500, {"ok": False, "error": str(e)})
        finally:
            LOCK.release()

    def log_message(self, fmt, *args):
        log("%s - %s" % (self.address_string(), fmt % args))


if __name__ == "__main__":
    log("ffmpeg-worker v3 listening on :%d (work=%s queue=%s)" % (PORT, WORK, QUEUE))
    ThreadingHTTPServer(("0.0.0.0", PORT), Handler).serve_forever()