import gradio as gr
import os, json, shutil, tempfile, subprocess
from pathlib import Path

# ---------- helpers ----------
def run(cmd):
    r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if r.returncode != 0:
        raise RuntimeError(r.stderr.strip())
    return r.stdout

def ffprobe_json(path):
    return json.loads(run([
        "ffprobe", "-v", "error", "-print_format", "json",
        "-show_format", "-show_streams", path
    ]))

def get_duration_seconds(meta):
    if "format" in meta and "duration" in meta["format"]:
        return float(meta["format"]["duration"])
    dur = 0.0
    for s in meta.get("streams", []):
        if "duration" in s:
            try:
                dur = max(dur, float(s["duration"]))
            except:
                pass
    if dur <= 0:
        raise ValueError("Could not determine video duration.")
    return dur

def has_audio_stream(meta):
    return any(s.get("codec_type") == "audio" for s in meta.get("streams", []))

# ---------- core logic ----------
def split_and_zip(video_path, n_parts):
    if video_path is None:
        raise gr.Error("Please upload a video.")
    try:
        n_parts = int(n_parts)
    except:
        raise gr.Error("Number of parts must be an integer.")
    if n_parts < 1:
        raise gr.Error("Number of parts must be ≥ 1.")

    # probe
    meta = ffprobe_json(video_path)
    duration = get_duration_seconds(meta)
    audio_present = has_audio_stream(meta)

    stem = Path(video_path).stem

    # temp working dir
    workdir = tempfile.mkdtemp(prefix="split_")
    outdir = Path(workdir) / f"{stem}_parts"
    outdir.mkdir(parents=True, exist_ok=True)

    seg_len = duration / n_parts
    logs = [f"Input: {Path(video_path).name} ({duration:.3f}s) | Parts: {n_parts} | Segment ~ {seg_len:.3f}s"]

    # split with ffmpeg (re-encode for clean frame-accurate cuts)
    for i in range(n_parts):
        start = max(0.0, min(i * seg_len, duration))
        length = duration - start if i == n_parts - 1 else seg_len
        length = max(0.0, min(length, duration - start))

        # Filename: Filename_Part_XX.mp4 (with leading zero up to 09)
        out_file = outdir / f"{stem}_Part_{i+1:02d}.mp4"

        cmd = [
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
            "-ss", f"{start:.6f}",
            "-i", video_path,
            "-t", f"{length:.6f}",
            "-c:v", "libx264", "-preset", "veryfast", "-crf", "18",
            "-movflags", "+faststart"
        ]
        if audio_present:
            cmd += ["-c:a", "aac", "-b:a", "128k"]
        else:
            cmd += ["-an"]
        cmd += [str(out_file)]
        try:
            run(cmd)
            logs.append(f"✔️ Saved {out_file.name} ({length:.3f}s)")
        except Exception as e:
            logs.append(f"❌ Failed part {i+1}: {e}")

    # ZIP filename: Filename_Split.zip
    zip_path = Path(workdir) / f"{stem}_Split.zip"
    shutil.make_archive(str(zip_path.with_suffix("")), "zip", root_dir=outdir)

    logs.append(f"ZIP ready: {zip_path.name}")
    return str(zip_path), "\n".join(logs)

# ---------- UI ----------
with gr.Blocks(title="Video Splitter") as demo:
    gr.Markdown("## Split a video into N equal parts and download a ZIP\n"
                "Upload a video, choose how many parts, and get a zipped folder of clips.")

    with gr.Row():
        video = gr.File(label="Upload video", file_count="single", type="filepath")
        parts = gr.Number(label="Number of parts", value=10, precision=0)

    with gr.Row():
        run_btn = gr.Button("Split Video", variant="primary")

    out_zip = gr.File(label="Download ZIP")
    log_box = gr.Textbox(label="Process Log", lines=10)

    run_btn.click(split_and_zip, inputs=[video, parts], outputs=[out_zip, log_box])

demo.launch()