app.py

import gradio as gr
import subprocess
import json
import os
import tempfile
from pydub import AudioSegment


def hms_to_seconds(hms: str) -> float:
    """Convert hh:mm:ss or mm:ss or ss to seconds."""
    parts = hms.strip().split(":")
    if len(parts) == 1:
        return float(parts[0])
    elif len(parts) == 2:
        m, s = parts
        return float(m) * 60 + float(s)
    elif len(parts) == 3:
        h, m, s = parts
        return float(h) * 3600 + float(m) * 60 + float(s)
    return 0.0


def seconds_to_hms(seconds: float) -> str:
    """Convert seconds to hh:mm:ss string."""
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}"


def get_video_duration(video_path: str) -> float:
    """Use ffprobe to get video duration in seconds."""
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "json",
        video_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    data = json.loads(result.stdout)
    return float(data["format"]["duration"])


def extract_audio(video_path: str, output_wav: str, start_s: float = None, end_s: float = None) -> str:
    """Extract audio from video to WAV using ffmpeg, optionally clipping to interval."""
    cmd = ["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "pcm_s16le", "-ar", "44100", "-ac", "2"]
    if start_s is not None:
        cmd += ["-ss", str(start_s)]
    if end_s is not None:
        cmd += ["-to", str(end_s)]
    cmd += [output_wav]
    subprocess.run(cmd, capture_output=True, check=True)
    return output_wav


FORMAT_EXT = {"MP3": "mp3", "M4A": "m4a", "WAV": "wav", "OGG": "ogg", "WEBM": "webm"}
MAX_SEGMENTS = 10


def segment_audio(
    audio_path: str,
    segment_sec: float,
    overlap_sec: float,
    n_segments: int = None,
    output_dir: str = None,
    out_format: str = "mp3",
) -> list:
    """Segment audio by fixed time or fixed number of segments, with overlap.

    If segment_sec >= duration, return one segment of the entire audio.
    Shorter last segment is included if it has non-zero length.
    Segments are capped at MAX_SEGMENTS (10).
    """
    audio = AudioSegment.from_file(audio_path)
    duration_ms = len(audio)

    # If asking for 1 segment, return entire interval
    if n_segments == 1:
        ext = FORMAT_EXT.get(out_format.upper(), "mp3")
        seg_path = os.path.join(output_dir, f"seg_000.{ext}")
        audio.export(seg_path, format=ext)
        return [seg_path]

    # If segment length >= duration, return one segment of entire audio
    segment_ms = int(segment_sec * 1000)
    overlap_ms = int(overlap_sec * 1000)
    if segment_ms >= duration_ms:
        ext = FORMAT_EXT.get(out_format.upper(), "mp3")
        seg_path = os.path.join(output_dir, f"seg_000.{ext}")
        audio.export(seg_path, format=ext)
        return [seg_path]

    step_ms = segment_ms - overlap_ms

    if n_segments is not None and n_segments > 0:
        # Fixed number of segments — cap at MAX_SEGMENTS
        n_segments = min(n_segments, MAX_SEGMENTS)
        if n_segments == 1:
            ext = FORMAT_EXT.get(out_format.upper(), "mp3")
            seg_path = os.path.join(output_dir, f"seg_000.{ext}")
            audio.export(seg_path, format=ext)
            return [seg_path]
        step_ms = (duration_ms - segment_ms) // max(n_segments - 1, 1)

    ext = FORMAT_EXT.get(out_format.upper(), "mp3")
    segments = []
    start = 0
    idx = 0

    while start < duration_ms and idx < MAX_SEGMENTS:
        end = min(start + segment_ms, duration_ms)
        if end <= start:
            break
        seg = audio[start:end]
        seg_path = os.path.join(output_dir, f"seg_{idx:03d}.{ext}")
        seg.export(seg_path, format=ext)
        segments.append(seg_path)
        start += step_ms
        idx += 1
        if n_segments is not None and idx >= n_segments:
            break

    return segments


def process_video(
    video_file,
    start_hms,
    end_hms,
    mode,
    segment_len,
    n_segments,
    overlap,
    out_format,
):
    """Main processing function."""
    if video_file is None:
        return [], gr.update(visible=False), gr.update(visible=True, value="Please upload a video.")

    start_s = hms_to_seconds(start_hms)
    end_s = hms_to_seconds(end_hms)
    start_s = max(0, float(start_s))
    end_s = max(start_s, float(end_s))

    tmpdir = tempfile.mkdtemp()

    wav_path = os.path.join(tmpdir, "audio.wav")
    extract_audio(video_file, wav_path, start_s, end_s)

    seg_len = float(segment_len)
    ov = float(overlap)
    n_seg = int(n_segments) if mode == "Fixed number of segments" else None

    if seg_len <= 0 and n_seg is None:
        return [], gr.update(visible=False), gr.update(visible=True, value="Segment length must be > 0.")

    paths = segment_audio(wav_path, seg_len, ov, n_segments=n_seg, output_dir=tmpdir, out_format=out_format)

    if not paths:
        return [], gr.update(visible=False), gr.update(visible=True, value="No segments produced. Check your parameters.")

    return paths, gr.update(visible=True), gr.update(visible=False)


# ---------------------------------------------------------------------------
# Build UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Video Audio Extractor & Segmenter") as demo:
    gr.Markdown("""
    # 🎬 Video → 🔊 Audio Extractor & Segmenter
    Upload a video file (MP4, WEBM, AVI, MOV, MKV, etc.), pick a time interval (hh:mm:ss), choose segmentation mode, and extract audio segments.
    """)

    # Use gr.File with explicit file types so WEBM is not dimmed in the picker.
    # gr.Video has no file_types parameter; the frontend hardcodes a restrictive accept filter.
    video_input = gr.File(
        file_types=["video", ".webm", ".mp4", ".m4v", ".avi", ".mov", ".mkv", ".ogv"],
        label="Upload Video",
        file_count="single",
    )

    duration_state = gr.State(value=60.0)

    with gr.Row():
        discover_btn = gr.Button("🔄 Discover Duration", scale=1)
        duration_text = gr.Textbox(label="Video Duration", value="Not discovered yet", interactive=False, scale=2)

    with gr.Row():
        start_input = gr.Textbox(
            value="00:00:00",
            label="Start Time (hh:mm:ss)",
            placeholder="00:00:00",
        )
        end_input = gr.Textbox(
            value="00:00:10",
            label="End Time (hh:mm:ss)",
            placeholder="00:00:10",
        )

    with gr.Row():
        mode = gr.Radio(
            ["Fixed segment length", "Fixed number of segments"],
            value="Fixed segment length",
            label="Segmentation Mode",
        )
        seg_len = gr.Number(value=5, minimum=0.5, step=0.5, label="Segment Length (s)")
        n_seg = gr.Number(value=3, minimum=1, maximum=MAX_SEGMENTS, step=1, label=f"Number of Segments (max {MAX_SEGMENTS})", visible=False)
        overlap = gr.Slider(0, 20, value=0, step=0.5, label="Overlap (s)")
        out_format = gr.Dropdown(
            ["MP3", "M4A", "WAV", "OGG", "WEBM"],
            value="MP3",
            label="Output Format",
        )

    process_btn = gr.Button("▶️ Extract & Segment Audio", variant="primary")

    error_msg = gr.Markdown(visible=False)
    segments_container = gr.Markdown(visible=False)

    segment_state = gr.State([])

    @gr.render(inputs=segment_state)
    def show_segments(paths):
        if not paths:
            gr.Markdown("Segments will appear here after processing.")
            return
        gr.Markdown(f"**Generated {len(paths)} segment(s):**")
        for i, p in enumerate(paths):
            with gr.Row():
                gr.Audio(value=p, label=f"Segment {i + 1}", interactive=False)

    def discover_duration(video_file):
        if video_file is None:
            return 60.0, "No video uploaded", "00:00:00", "00:01:00"
        dur = get_video_duration(video_file)
        return dur, f"{seconds_to_hms(dur)} (total {dur:.2f} s)", "00:00:00", seconds_to_hms(min(10.0, dur))

    discover_btn.click(
        fn=discover_duration,
        inputs=video_input,
        outputs=[duration_state, duration_text, start_input, end_input],
    )

    def toggle_mode(m):
        if m == "Fixed segment length":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    mode.change(
        fn=toggle_mode,
        inputs=mode,
        outputs=[seg_len, n_seg],
    )

    def on_process(video_file, start_hms, end_hms, mode_val, seg_len_val, n_seg_val, overlap_val, out_format_val):
        paths, container_update, err_update = process_video(
            video_file, start_hms, end_hms, mode_val, seg_len_val, n_seg_val, overlap_val, out_format_val
        )
        return paths, container_update, err_update

    process_btn.click(
        fn=on_process,
        inputs=[video_input, start_input, end_input, mode, seg_len, n_seg, overlap, out_format],
        outputs=[segment_state, segments_container, error_msg],
    )

if __name__ == "__main__":
    demo.launch()