import os
import re
import shutil
import subprocess
import tempfile
from pathlib import Path

import gradio as gr
import torch
from transformers import pipeline

# ---------------------------------------------------------------------------
# Model setup
# ---------------------------------------------------------------------------
MODEL_ID = "openai/whisper-base"  # CPU-friendly; swap for whisper-large-v3 on GPU

device = 0 if torch.cuda.is_available() else -1
asr = pipeline(
    "automatic-speech-recognition",
    model=MODEL_ID,
    chunk_length_s=30,
    device=device,
    return_timestamps=True,
)

SUPPORTED = [".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v", ".mp3", ".wav", ".m4a", ".flac", ".ogg"]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def extract_audio(video_path: str, workdir: str) -> str:
    """Extract mono 16 kHz WAV audio from a media file using ffmpeg."""
    audio_path = os.path.join(workdir, "audio.wav")
    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-vn",               # drop video stream
        "-acodec", "pcm_s16le",
        "-ar", "16000",
        "-ac", "1",
        audio_path,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0 or not os.path.exists(audio_path):
        raise RuntimeError(f"ffmpeg failed to extract audio:\n{proc.stderr[-1500:]}")
    return audio_path


def format_timestamp(seconds: float) -> str:
    """Format seconds as SRT timestamp: HH:MM:SS,mmm"""
    if seconds is None:
        seconds = 0.0
    ms = int(round(seconds * 1000))
    h, ms = divmod(ms, 3600_000)
    m, ms = divmod(ms, 60_000)
    s, ms = divmod(ms, 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def build_srt(chunks) -> str:
    blocks = []
    for i, chunk in enumerate(chunks, start=1):
        start, end = chunk.get("timestamp", (0.0, 0.0))
        if end is None:
            end = (start or 0.0) + 2.0
        text = chunk["text"].strip()
        blocks.append(
            f"{i}\n{format_timestamp(start)} --> {format_timestamp(end)}\n{text}"
        )
    return "\n\n".join(blocks) + "\n"


def safe_stem(filename: str) -> str:
    stem = Path(filename).stem
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", stem) or "transcript"


# ---------------------------------------------------------------------------
# Main transcription handler
# ---------------------------------------------------------------------------
def transcribe(video_file, language, task, progress=gr.Progress(track_tqdm=True)):
    if video_file is None:
        return "Please upload a video or audio file first.", None, None

    src_path = video_file if isinstance(video_file, str) else video_file.name
    ext = Path(src_path).suffix.lower()
    if ext and ext not in SUPPORTED:
        return f"Unsupported file type: `{ext}`", None, None

    with tempfile.TemporaryDirectory() as workdir:
        progress(0.1, desc="Extracting audio with ffmpeg...")
        audio_path = extract_audio(src_path, workdir)

        progress(0.3, desc="Transcribing with Whisper (this can take a while)...")
        kwargs = {"generate_kwargs": {"task": task}}
        if language != "auto":
            kwargs["generate_kwargs"]["language"] = language
        result = asr(audio_path, **kwargs)

        text = result["text"].strip()
        chunks = result.get("chunks") or []

        progress(0.9, desc="Preparing download files...")
        stem = safe_stem(src_path)

        txt_path = os.path.join(workdir, f"{stem}_transcript.txt")
        with open(txt_path, "w", encoding="utf-8") as f:
            f.write(text + "\n")
        final_txt = os.path.join(tempfile.gettempdir(), f"{stem}_transcript.txt")
        shutil.copy(txt_path, final_txt)

        final_srt = None
        if chunks:
            srt_path = os.path.join(workdir, f"{stem}_subtitles.srt")
            with open(srt_path, "w", encoding="utf-8") as f:
                f.write(build_srt(chunks))
            final_srt = os.path.join(tempfile.gettempdir(), f"{stem}_subtitles.srt")
            shutil.copy(srt_path, final_srt)

    return text, final_txt, final_srt


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
LANGUAGES = [
    "auto", "english", "spanish", "french", "german", "italian", "portuguese",
    "dutch", "russian", "chinese", "japanese", "korean", "arabic", "hindi",
    "turkish", "polish", "swedish", "ukrainian", "vietnamese", "indonesian",
]

with gr.Blocks(title="Video Transcriber") as demo:
    gr.Markdown(
        """
# 🎬 → 📝 Video Transcriber (Speech-to-Text)

Upload a video (or audio) file and get a full text transcription powered by
**OpenAI Whisper** running on 🤗 Transformers.

**How it works**
1. **Upload** a video file (`.mp4`, `.mov`, `.mkv`, `.webm`, ...) — audio files work too.
2. Click **Transcribe**. The audio track is extracted with `ffmpeg` and transcribed.
3. **Read** the transcript on the right and **download** it as `.txt` or timed subtitles `.srt`.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            video_in = gr.Video(label="Upload video (or drop an audio file)", sources=["upload"])
            language = gr.Dropdown(
                choices=LANGUAGES, value="auto",
                label="Language (optional — auto-detect by default)",
            )
            task = gr.Radio(
                choices=["transcribe", "translate"], value="transcribe",
                label="Task ('translate' translates speech into English)",
            )
            run_btn = gr.Button("🚀 Transcribe", variant="primary")
        with gr.Column(scale=1):
            transcript_out = gr.Textbox(label="Transcript", lines=22)
            with gr.Row():
                txt_file = gr.File(label="⬇️ Download transcript (.txt)")
                srt_file = gr.File(label="⬇️ Download subtitles (.srt)")

    run_btn.click(
        fn=transcribe,
        inputs=[video_in, language, task],
        outputs=[transcript_out, txt_file, srt_file],
    )

if __name__ == "__main__":
    demo.queue(max_size=8).launch(server_name="0.0.0.0", server_port=7860)