"""
Whisper Lite — Upload audio/video, get transcript.
No background jobs, no URLs, no tokens. Just transcription.
"""

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


ACCEPTED_EXTS = [
    ".mp3", ".mp4", ".wav", ".m4a", ".ogg", ".flac",
    ".webm", ".mkv", ".avi", ".mov", ".aac", ".opus",
]

# Pre-load base model at startup to warm the cache
import threading
def _warm_cache():
    try:
        import whisper as _w
        _w.load_model("base")
        print("[startup] Whisper base model loaded into cache")
    except Exception as e:
        print(f"[startup] Model pre-load skipped: {e}")
threading.Thread(target=_warm_cache, daemon=True).start()

def fmt_time(seconds: float) -> str:
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:05.2f}"

def detect_speaker(segments, idx, gap=1.5):
    if idx == 0:
        return "SPEAKER_01"
    g = segments[idx]["start"] - segments[idx - 1]["end"]
    prev = detect_speaker(segments, idx - 1) if idx > 1 else "SPEAKER_01"
    if g > gap:
        return "SPEAKER_02" if prev == "SPEAKER_01" else "SPEAKER_01"
    return prev

def transcribe(file_path, model_size, progress=gr.Progress(track_tqdm=False)):
    if file_path is None:
        raise gr.Error("Please upload a file first.")

    src = Path(file_path)
    if src.suffix.lower() not in ACCEPTED_EXTS:
        raise gr.Error(f"Unsupported format: {src.suffix}. Accepted: {', '.join(ACCEPTED_EXTS)}")

    workdir = Path(tempfile.mkdtemp())
    try:
        # Convert to 16kHz mono wav (optimal for Whisper)
        progress(0.05, desc="Converting audio…")
        wav_path = workdir / "audio.wav"
        ff = subprocess.run([
            "ffmpeg", "-y", "-i", str(src),
            "-ar", "16000", "-ac", "1", "-vn",
            str(wav_path)
        ], capture_output=True, text=True)
        if ff.returncode != 0:
            raise gr.Error(f"ffmpeg failed: {ff.stderr[-400:]}")

        # Run Whisper
        progress(0.15, desc=f"Loading {model_size} model…")
        cmd = [
            "whisper", str(wav_path),
            "--model", model_size,
            "--language", "auto",
            "--output_format", "json",
            "--output_dir", str(workdir),
            "--verbose", "False",
        ]

        progress(0.25, desc="Transcribing… (this may take a while)")
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        try:
            stdout_lines = []
            for line in proc.stdout:
                stdout_lines.append(line)
            proc.wait(timeout=3600)  # 1 hour max
        except subprocess.TimeoutExpired:
            proc.kill()
            raise gr.Error("Transcription timed out after 1 hour.")
        if proc.returncode != 0:
            raise gr.Error(f"Whisper failed: {''.join(stdout_lines[-10:])}")

        # Parse JSON → formatted transcript
        progress(0.90, desc="Formatting transcript…")
        import json
        json_files = list(workdir.glob("*.json"))
        if not json_files:
            raise gr.Error("Whisper produced no output — try a different model or file.")

        with open(json_files[0]) as f:
            data = json.load(f)

        segments = data.get("segments", [])
        if not segments:
            raise gr.Error("No speech detected in this file.")

        lines = []
        for i, seg in enumerate(segments):
            start   = fmt_time(seg["start"])
            end     = fmt_time(seg["end"])
            text    = seg["text"].strip()
            speaker = detect_speaker(segments, i)
            lines.append(f"[{start} → {end}] {speaker}: {text}")

        transcript = "\n".join(lines)

        # Save output txt
        out_txt = workdir / "transcript.txt"
        out_txt.write_text(transcript, encoding="utf-8")

        # Copy to a persistent temp location Gradio can serve
        final_txt = Path(tempfile.mktemp(suffix="_transcript.txt"))
        shutil.copy(out_txt, final_txt)

        progress(1.0, desc="Done!")
        return transcript, str(final_txt)

    finally:
        shutil.rmtree(workdir, ignore_errors=True)


# ── UI ────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@300;400;500;700&display=swap');

*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }

body, .gradio-container {
    background: #080808 !important;
    color: #d4d4d4 !important;
    font-family: 'JetBrains Mono', 'Courier New', monospace !important;
}

/* Hide Gradio chrome */
footer { display: none !important; }
.gr-prose { display: none !important; }

/* Header */
#header {
    padding: 2rem 0 1rem;
    border-bottom: 1px solid #1c1c1c;
    margin-bottom: 1.5rem;
}
#header .eyebrow {
    font-size: 10px;
    letter-spacing: .22em;
    color: #3a3a3a;
    text-transform: uppercase;
    margin-bottom: 6px;
}
#header h1 {
    font-size: 1.5rem;
    font-weight: 700;
    letter-spacing: .06em;
    color: #f0f0f0;
}
#header .sub {
    font-size: 11px;
    color: #3d3d3d;
    letter-spacing: .08em;
    margin-top: 4px;
}

/* Upload zone */
.gr-file-upload, .upload-container, [data-testid="file"] {
    background: #0d0d0d !important;
    border: 1px dashed #242424 !important;
    border-radius: 4px !important;
    min-height: 120px !important;
    transition: border-color .2s !important;
}
.gr-file-upload:hover { border-color: #f0a030 !important; }

/* Model selector */
.gr-dropdown select, select {
    background: #0d0d0d !important;
    color: #d4d4d4 !important;
    border: 1px solid #222 !important;
    border-radius: 3px !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 12px !important;
}

/* Labels */
label, .gr-label span, .block > label > span {
    font-size: 10px !important;
    letter-spacing: .14em !important;
    text-transform: uppercase !important;
    color: #3a3a3a !important;
    font-family: 'JetBrains Mono', monospace !important;
}

/* Button */
.run-btn button {
    background: #f0a030 !important;
    color: #080808 !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-weight: 700 !important;
    font-size: 12px !important;
    letter-spacing: .14em !important;
    text-transform: uppercase !important;
    border: none !important;
    border-radius: 2px !important;
    padding: 10px 28px !important;
    width: 100% !important;
    cursor: pointer !important;
    transition: opacity .15s !important;
}
.run-btn button:hover { opacity: .85 !important; }
.run-btn button:disabled { opacity: .35 !important; cursor: not-allowed !important; }

/* Transcript output */
.gr-textbox textarea {
    background: #0a0a0a !important;
    color: #b4b4b4 !important;
    border: 1px solid #1a1a1a !important;
    border-radius: 3px !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 12px !important;
    line-height: 1.7 !important;
}

/* Download area */
.gr-file, .file-preview {
    background: #0d0d0d !important;
    border: 1px solid #1a1a1a !important;
    border-radius: 3px !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 11px !important;
}

/* Progress */
.progress-bar { background: #f0a030 !important; }

/* Inputs row */
.inputs-row { gap: 12px !important; align-items: flex-end !important; }
"""

MODEL_INFO = {
    "tiny":   "~1 GB  · fastest · basic accuracy",
    "base":   "~1 GB  · fast    · good accuracy   ← recommended",
    "small":  "~2 GB  · medium  · better accuracy",
    "medium": "~5 GB  · slow    · great accuracy",
    "large":  "~10 GB · slowest · best accuracy",
}

with gr.Blocks(title="Whisper Lite") as demo:

    gr.HTML("""
    <div id="header">
        <div class="eyebrow">Lexical Space</div>
        <h1>WHISPER LITE</h1>
        <div class="sub">Upload audio or video · choose model · get transcript</div>
    </div>
    """)

    with gr.Row(elem_classes=["inputs-row"]):
        with gr.Column(scale=3):
            file_input = gr.File(
                label="Audio / Video file",
                file_types=ACCEPTED_EXTS,
                type="filepath",
            )
        with gr.Column(scale=1, min_width=200):
            model_dd = gr.Dropdown(
                choices=list(MODEL_INFO.keys()),
                value="base",
                label="Whisper model",
                info=None,
            )
            model_info_box = gr.Markdown(
                value=f"`base` — {MODEL_INFO['base']}",
                elem_id="model_info",
            )

    with gr.Row():
        with gr.Column(elem_classes=["run-btn"]):
            run_btn = gr.Button("▶  Transcribe", variant="primary")

    transcript_box = gr.Textbox(
        label="Transcript",
        lines=18,
        max_lines=40,
        interactive=False,
        placeholder="Transcript will appear here…",
    )

    dl_file = gr.File(label="Download .txt", visible=False)

    # ── Wire model info update ────────────────────────────────────────────────
    def update_model_info(m):
        return f"`{m}` — {MODEL_INFO.get(m, '')}"

    model_dd.change(update_model_info, inputs=[model_dd], outputs=[model_info_box])

    # ── Wire transcribe ───────────────────────────────────────────────────────
    def on_transcribe(file_path, model_size, progress=gr.Progress()):
        text, txt_path = transcribe(file_path, model_size, progress)
        return text, gr.update(value=txt_path, visible=True)

    run_btn.click(
        fn=on_transcribe,
        inputs=[file_input, model_dd],
        outputs=[transcript_box, dl_file],
        show_progress="full",
    )

    # Auto-start when file is uploaded
    file_input.upload(
        fn=on_transcribe,
        inputs=[file_input, model_dd],
        outputs=[transcript_box, dl_file],
        show_progress="full",
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, css=CSS)