#!/usr/bin/env python3
"""rozbierz.py: split an ad video into shots for the reklamy skill.

Usage:
  python3 rozbierz.py <video file or URL> [--out folder] [--threshold 0.3] [--lang pl]

Needs ffmpeg + ffprobe. A URL needs yt-dlp. Speech-to-text runs only if
faster-whisper or openai-whisper is installed; otherwise the agent listens
itself or asks the user to paste the voice-over.

Writes to the out folder: shots.json (cut times, length, pacing), one frame per
shot (shot-01.jpg ...), first.jpg + last.jpg, audio.wav and, when possible,
transcript.json (speech with timestamps). Sends nothing anywhere.
"""
import argparse, json, re, shutil, subprocess, sys
from pathlib import Path


def run(cmd):
    return subprocess.run(cmd, capture_output=True, text=True)


def need(tool):
    if not shutil.which(tool):
        sys.exit(f"rozbierz: {tool} is not installed (https://ffmpeg.org/download.html, https://github.com/yt-dlp/yt-dlp)")


def fetch(url, out):
    need("yt-dlp")
    dst = out / "ad.mp4"
    r = run(["yt-dlp", "-q", "--no-playlist", "-f", "bv*[height<=1080]+ba/b", "--merge-output-format", "mp4", "-o", str(dst), url])
    if r.returncode or not dst.exists():
        sys.exit("rozbierz: download failed. Save the ad as a file yourself (screen recording works) and pass the file.\n" + r.stderr[-400:])
    return dst


def duration(video):
    return float(run(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(video)]).stdout.strip())


def cuts(video, threshold):
    # scene score per frame; showinfo prints pts_time for every frame the select keeps
    r = run(["ffmpeg", "-nostdin", "-i", str(video), "-filter:v", f"select='gt(scene,{threshold})',showinfo", "-f", "null", "-"])
    return [float(t) for t in re.findall(r"pts_time:([0-9.]+)", r.stderr)]


def frame(video, t, path):
    run(["ffmpeg", "-nostdin", "-v", "error", "-y", "-ss", f"{t:.3f}", "-i", str(video), "-frames:v", "1", "-vf", "scale=-2:540", str(path)])


def transcribe(audio, lang):
    try:
        from faster_whisper import WhisperModel
        segs, _ = WhisperModel("small", compute_type="int8").transcribe(str(audio), language=lang or None)
        return [{"start": round(s.start, 2), "end": round(s.end, 2), "text": s.text.strip()} for s in segs]
    except ImportError:
        pass
    try:
        import whisper
        res = whisper.load_model("small").transcribe(str(audio), language=lang or None)
        return [{"start": round(s["start"], 2), "end": round(s["end"], 2), "text": s["text"].strip()} for s in res["segments"]]
    except ImportError:
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("src")
    ap.add_argument("--out", default="ad-breakdown")
    ap.add_argument("--threshold", type=float, default=0.3, help="lower = more cuts found (0.2-0.4)")
    ap.add_argument("--lang", default=None, help="speech language, e.g. pl or en (default: auto)")
    a = ap.parse_args()
    need("ffmpeg"); need("ffprobe")
    out = Path(a.out).expanduser(); out.mkdir(parents=True, exist_ok=True)
    video = fetch(a.src, out) if re.match(r"https?://", a.src) else Path(a.src).expanduser()
    if not video.exists():
        sys.exit(f"rozbierz: no such file: {video}")

    total = duration(video)
    bounds = [0.0] + [t for t in cuts(video, a.threshold) if 0.4 < t < total - 0.2] + [total]
    edges = [bounds[0]]
    for t in bounds[1:]:
        if t - edges[-1] >= 0.4:  # a flash or a glitch is not a new shot
            edges.append(t)
    edges[-1] = total
    shots = []
    for i, (s, e) in enumerate(zip(edges, edges[1:]), 1):
        f = out / f"shot-{i:02d}.jpg"
        frame(video, (s + e) / 2, f)
        shots.append({"n": i, "start": round(s, 2), "end": round(e, 2), "length": round(e - s, 2), "frame": f.name})
    frame(video, 0.1, out / "first.jpg")
    frame(video, max(total - 0.2, 0), out / "last.jpg")

    run(["ffmpeg", "-nostdin", "-v", "error", "-y", "-i", str(video), "-vn", "-ac", "1", "-ar", "16000", str(out / "audio.wav")])
    speech = transcribe(out / "audio.wav", a.lang) if (out / "audio.wav").exists() else None
    if speech is not None:
        (out / "transcript.json").write_text(json.dumps(speech, ensure_ascii=False, indent=1))

    report = {
        "video": str(video), "duration": round(total, 2), "shots": shots,
        "pacing": {"shot_count": len(shots), "avg_shot_s": round(total / len(shots), 2),
                   "cuts_per_10s": round((len(shots) - 1) / total * 10, 1)},
        "transcript": "transcript.json" if speech is not None else None,
        "note": None if speech is not None else "no speech-to-text installed (pip install faster-whisper); listen to audio.wav or ask the user for the voice-over",
    }
    (out / "shots.json").write_text(json.dumps(report, ensure_ascii=False, indent=1))
    print(json.dumps(report["pacing"]), f"-> {out}/shots.json, {len(shots)} frames")


if __name__ == "__main__":
    main()