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")
    
    if not os.path.exists(video_path):
        raise RuntimeError(f"ffmpeg failed: Input file '{video_path}' does not exist.")

    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-vn",
        "-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):
        print(f"FFmpeg stderr: {proc.stderr}")
        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(media_file, language, task, progress=gr.Progress(track_tqdm=True)):
    if media_file is None:
        return "", ""

    # Ensure we have a valid string path
    # With type="filepath", media_file should be a string path directly
    if isinstance(media_file, str):
        src_path = media_file
    else:
        # Fallback for any unexpected types
        print(f"Unexpected media_file type: {type(media_file)}, value: {media_file}")
        return "Error: Invalid file input received. Please try uploading again.", ""

    # Validate the path
    if not os.path.exists(src_path):
        print(f"File not found: {src_path}")
        return f"Error: File not found at path: {src_path}", ""

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

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

        progress(0.3, desc="Transcribing (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="Finalizing...")
        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")

    return text, txt_path


# ---------------------------------------------------------------------------
# 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="Transcriber") as demo:
    gr.Markdown(
        """
# 🎙️ Audio/Video Transcriber
Upload any audio or video file and get a clean text transcription powered by **OpenAI Whisper**.
**How it works**
1. **Upload** an audio (.mp3, .wav, .m4a...) or video (.mp4, .mov, .mkv...) file.
2. Click **Transcribe**. The audio is extracted and transcribed automatically.
3. **Copy & paste** the transcript below or download it as a `.txt` file.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            # KEY FIX: Set type="filepath" to ensure we get a string path, not a tuple or object
            media_in = gr.Audio(label="Audio/Video", sources=["upload"], type="filepath")
            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 (copy & paste)", lines=25)
            txt_file = gr.File(label="⬇️ Download .txt")

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

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