"""
Custom captioning tool: transcribes audio/video with Whisper (large-v3) and chunks
captions by a user-specified number of words per caption, outputting an SRT file
ready to import into CapCut or any other editor.

Designed for Hugging Face Spaces ZeroGPU (dynamic H200 access).
Uses openai-whisper (PyTorch-based) rather than faster-whisper/CTranslate2,
since ZeroGPU's GPU-attach mechanism is built around torch.cuda and is most
reliable with PyTorch-native models.
"""

import os
import tempfile

import gradio as gr
import spaces
import torch
import whisper

# ---------------------------------------------------------------------------
# Model loads once at startup on CPU. It's moved to GPU only inside the
# @spaces.GPU-decorated function below, since real CUDA access on ZeroGPU
# only exists during that call.
# ---------------------------------------------------------------------------
MODEL_SIZE = "large-v3"
model = whisper.load_model(MODEL_SIZE, device="cpu")


def format_timestamp(seconds: float) -> str:
    """Convert seconds (float) to SRT timestamp format: HH:MM:SS,mmm"""
    ms_total = int(round(seconds * 1000))
    hours, rem = divmod(ms_total, 3600_000)
    minutes, rem = divmod(rem, 60_000)
    secs, ms = divmod(rem, 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"


def chunk_words(words, words_per_caption, max_chars=None):
    """
    Group a flat list of {"word","start","end"} dicts into caption chunks
    of `words_per_caption` words each. Optionally caps chunk length by
    max_chars so long words don't overflow a caption line.
    """
    captions = []
    current = []

    def flush():
        if current:
            captions.append({
                "text": " ".join(w["word"] for w in current).strip(),
                "start": current[0]["start"],
                "end": current[-1]["end"],
            })

    for w in words:
        current.append(w)
        text_len = len(" ".join(x["word"] for x in current))
        hit_word_limit = len(current) >= words_per_caption
        hit_char_limit = max_chars is not None and text_len >= max_chars
        if hit_word_limit or hit_char_limit:
            flush()
            current = []
    flush()

    return captions


def write_srt(captions, path):
    with open(path, "w", encoding="utf-8") as f:
        for i, cap in enumerate(captions, start=1):
            f.write(f"{i}\n")
            f.write(f"{format_timestamp(cap['start'])} --> {format_timestamp(cap['end'])}\n")
            f.write(f"{cap['text']}\n\n")


@spaces.GPU(duration=120)  # per-call GPU time budget; counts against your daily quota
def run_transcription(media_path):
    """The only part that touches CUDA -- kept as small as possible to conserve quota."""
    global model
    model = model.to("cuda")
    result = model.transcribe(media_path, word_timestamps=True, fp16=True)
    model = model.to("cpu")  # release VRAM before the GPU is handed back
    torch.cuda.empty_cache()
    return result


def transcribe_and_caption(media_file, words_per_caption, max_chars, progress=gr.Progress()):
    if media_file is None:
        return None, "Upload an audio or video file first."

    words_per_caption = max(1, int(words_per_caption))
    max_chars_val = int(max_chars) if max_chars and max_chars > 0 else None

    progress(0.2, desc="Requesting GPU and transcribing (uses your ZeroGPU quota)...")
    result = run_transcription(media_file)

    words = []
    for segment in result.get("segments", []):
        for w in segment.get("words", []):
            words.append({
                "word": w["word"].strip(),
                "start": w["start"],
                "end": w["end"],
            })

    if not words:
        return None, "No speech detected in the file."

    progress(0.8, desc="Building captions...")
    captions = chunk_words(words, words_per_caption, max_chars_val)

    out_path = os.path.join(tempfile.gettempdir(), "captions.srt")
    write_srt(captions, out_path)

    preview_lines = [f"[{format_timestamp(c['start'])}] {c['text']}" for c in captions[:15]]
    preview = "\n".join(preview_lines)
    if len(captions) > 15:
        preview += f"\n... ({len(captions) - 15} more captions)"

    detected_lang = result.get("language", "unknown")
    summary = (
        f"Detected language: {detected_lang}\n"
        f"Total captions: {len(captions)}\n\n"
        f"Preview:\n{preview}"
    )

    progress(1.0, desc="Done")
    return out_path, summary


with gr.Blocks(title="Word-Count Caption Generator") as demo:
    gr.Markdown(
        "# Word-Count Caption Generator\n"
        "Upload a video or audio file, choose how many words should appear per caption, "
        "and get back an `.srt` file to import into CapCut (or any editor).\n\n"
        f"Running **{MODEL_SIZE}** on ZeroGPU (H200). Each run uses part of your daily "
        "GPU quota, so batch your clips rather than testing repeatedly."
    )

    with gr.Row():
        with gr.Column():
            media_input = gr.File(
                label="Audio or video file",
                file_types=["audio", "video"],
            )
            words_slider = gr.Slider(
                minimum=1, maximum=10, value=3, step=1,
                label="Words per caption",
            )
            chars_slider = gr.Slider(
                minimum=0, maximum=60, value=0, step=1,
                label="Max characters per caption (0 = no limit)",
            )
            run_btn = gr.Button("Generate captions", variant="primary")

        with gr.Column():
            srt_output = gr.File(label="Download .srt")
            summary_output = gr.Textbox(label="Summary / preview", lines=18)

    run_btn.click(
        fn=transcribe_and_caption,
        inputs=[media_input, words_slider, chars_slider],
        outputs=[srt_output, summary_output],
    )

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