import gradio as gr
import time
import os
import wave
import threading
from typing import Union, Optional
from faster_whisper import WhisperModel

print("Loading whisper-small...")
model = WhisperModel("small", device="cpu", compute_type="int8")
print("✅ Ready")

# ══════════════════════════════════════════════════════════════════════════
# Concurrency / capacity tracking
# ══════════════════════════════════════════════════════════════════════════
# MAX_CONCURRENT_JOBS controls how many transcriptions Gradio's queue will
# actually run AT THE SAME TIME (passed as concurrency_limit on the
# gr.api() registration below). faster_whisper / ctranslate2 is safe to
# call from multiple threads on one shared `model` instance — it handles
# its own internal locking — so we don't need one model copy per slot.
#
# This number is a tradeoff: more parallel jobs = more users serviced at
# once, but each individual job runs SLOWER because they're sharing the
# same CPU cores. On HF Spaces' free CPU tier (2 vCPUs), 3 is a reasonable
# starting point — raise it if you upgrade hardware, lower it if requests
# are timing out under load.
MAX_CONCURRENT_JOBS = 5

_active_jobs_lock = threading.Lock()
_active_jobs = 0          # currently transcribing right now
_total_jobs_served = 0    # lifetime counter, for the status endpoint


def _job_started():
    global _active_jobs, _total_jobs_served
    with _active_jobs_lock:
        _active_jobs += 1
        _total_jobs_served += 1


def _job_finished():
    global _active_jobs
    with _active_jobs_lock:
        _active_jobs = max(0, _active_jobs - 1)


def _get_capacity_snapshot() -> dict:
    with _active_jobs_lock:
        active = _active_jobs
        total_served = _total_jobs_served
    return {
        "max_concurrent_jobs": MAX_CONCURRENT_JOBS,
        "active_jobs": active,
        "available_slots": max(0, MAX_CONCURRENT_JOBS - active),
        "total_jobs_served_since_startup": total_served,
    }


LANGUAGES = {
    "Auto Detect": None,
    "Arabic (ar)": "ar",   "Hindi (hi)": "hi",    "English (en)": "en",
    "French (fr)": "fr",   "German (de)": "de",   "Spanish (es)": "es",
    "Chinese (zh)": "zh",  "Japanese (ja)": "ja", "Russian (ru)": "ru",
    "Urdu (ur)": "ur",     "Turkish (tr)": "tr",  "Korean (ko)": "ko",
    "Italian (it)": "it",  "Portuguese (pt)": "pt",
}
VALID_CODES = {v for v in LANGUAGES.values() if v is not None}

# ── Big-file safety limits ──────────────────────────────────────────────────
# WAV is uncompressed, so file size scales directly with duration+quality
# (e.g. 16-bit/44.1kHz mono ≈ 5.3 MB/min). Adjust these to your needs.
MAX_AUDIO_SECONDS = 3 * 60 * 60        # 3 hours
MAX_FILE_BYTES    = 900 * 1024 * 1024  # 900 MB hard cap

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


def resolve_language(language_code: str):
    """Accept '', 'auto', None, a raw ISO code ('en'), or a UI label ('English (en)')."""
    if not language_code or language_code.lower() == "auto":
        return None
    if language_code in VALID_CODES:
        return language_code
    if language_code in LANGUAGES:
        return LANGUAGES[language_code]
    raise ValueError(f"Unknown language code: {language_code!r}")


def get_wav_duration_seconds(path: str) -> Optional[float]:
    """
    Cheap, near-instant duration check using Python's stdlib `wave` module —
    reads only the WAV header (frame count + sample rate), never decodes
    audio samples. Used purely for the MAX_AUDIO_SECONDS safety check
    before committing to a full transcription pass.

    Returns None if the file isn't a standard PCM WAV the `wave` module can
    parse (e.g. some compressed codec inside a .wav container) — in that
    case we simply skip the pre-check and let faster_whisper's own
    ffmpeg/PyAV-backed decoder (which is far more format-tolerant) handle
    it directly.
    """
    try:
        with wave.open(path, "rb") as wf:
            frames = wf.getnframes()
            rate = wf.getframerate()
            if rate > 0:
                return frames / float(rate)
    except Exception:
        pass
    return None


def transcribe_audio_stream(audio_path: str, language):
    """
    Generator version of the transcription core. Yields a series of
    progress dicts like:
        {"status": "progress", "fraction": 0.42, "desc": "Transcribed 00:08.10 / 00:19.30"}
    and finally yields ONE result dict:
        {"status": "done", "language": ..., "words": [...], ...}

    Used by BOTH gradio_decode (consumes the generator directly instead of
    a progress_cb) and api_transcribe (yields straight through — gr.api()
    auto-streams generator yields as SSE events, which is how a client's
    poll loop can show a real percentage instead of a guessed one).

    ── DIRECT WAV DECODING ──────────────────────────────────────────────
    `audio_path` is passed straight to faster_whisper's `transcribe()`,
    which decodes and resamples the file itself via its ffmpeg/PyAV
    backend. No custom PCM parsing, no manual resampling, no memmap
    bookkeeping — that entire layer is gone. This is both simpler and
    faster: ffmpeg's native decode path is implemented in C and is
    considerably quicker than any hand-rolled Python/NumPy pipeline.

    ── SPEED NOTES ──────────────────────────────────────────────────────
    - BatchedInferencePipeline (see `batched_model` at module load) is
      tried first — it's the primary speed lever for CPU throughput.
    - beam_size=2 (down from 5) trades a small amount of decode robustness
      for speed. This is safe now: the earlier word-duplication /
      hallucination issue was traced to a bug in the Android app's custom
      PCM resampler corrupting samples before they ever reached Whisper —
      not to beam_size. Since audio now arrives as an unmodified WAV file,
      that failure mode no longer applies. compression_ratio_threshold and
      no_repeat_ngram_size are kept as cheap safety nets regardless.
    - Includes automatic fallback: if the batched/VAD pass finds no
      speech at all (common on quiet or corrupted audio), automatically
      retries once with the plain (non-batched) pipeline and VAD disabled.
    """
    t0 = time.time()

    def _run(vad_filter: bool, vad_params):
        kwargs = dict(
            audio=audio_path,
            language=language,
            beam_size=2,
            word_timestamps=True,
            vad_filter=vad_filter,
            vad_parameters=vad_params,
            condition_on_previous_text=False,
            compression_ratio_threshold=2.4,
            no_repeat_ngram_size=3,
        )

        segments_iter, info = model.transcribe(**kwargs)
        audio_duration = info.duration

        if audio_duration > 0:
            yield {
                "status":   "progress",
                "fraction": 0.08,
                "desc":     f"Loaded {fmt_time(audio_duration)} of audio",
            }

        segments, words = [], []
        seg_count = 0
        for seg in segments_iter:
            seg_count += 1
            segments.append({
                "start": round(float(seg.start), 3),
                "end":   round(float(seg.end), 3),
                "text":  seg.text.strip(),
            })
            if seg.words:
                for w in seg.words:
                    words.append({
                        "start": round(float(w.start), 3),
                        "end":   round(float(w.end), 3),
                        "word":  w.word.strip(),
                    })

            if audio_duration > 0:
                frac = min(0.95, 0.1 + 0.85 * (seg.end / audio_duration))
                elapsed_so_far = time.time() - t0

                yield {
                    "status":   "progress",
                    "fraction": round(frac, 4),
                    "desc":     f"{fmt_time(seg.end)} / {fmt_time(audio_duration)}  (elapsed: {elapsed_so_far:.1f}s)",
                }

        return_value["segments"] = segments
        return_value["words"] = words
        return_value["info"] = info
        return_value["seg_count"] = seg_count
        return_value["audio_duration"] = audio_duration

    yield {"status": "progress", "fraction": 0.05, "desc": "Loading audio..."}

    return_value = {}
    yield from _run(vad_filter=True, vad_params=dict(min_silence_duration_ms=500))
    segments, words = return_value["segments"], return_value["words"]
    info, seg_count = return_value["info"], return_value["seg_count"]
    audio_duration = return_value["audio_duration"]

    vad_fallback_used = False
    if not segments:
        vad_fallback_used = True
        yield {"status": "progress", "fraction": 0.5,
               "desc": "No speech detected — retrying without VAD..."}
        return_value = {}
        yield from _run(vad_filter=False, vad_params=None)
        segments, words = return_value["segments"], return_value["words"]
        info, seg_count = return_value["info"], return_value["seg_count"]
        audio_duration = return_value["audio_duration"]

    elapsed_s = time.time() - t0
    rtf = elapsed_s / max(audio_duration, 0.01)

    yield {"status": "progress", "fraction": 1.0, "desc": "Done"}

    note = None
    if not words:
        note = "No speech detected in this audio."
    elif vad_fallback_used and words:
        note = "VAD pass found nothing; retried with VAD disabled."

    yield {
        "status":           "done",
        "language":         info.language,
        "language_prob":    round(info.language_probability, 3),
        "duration_seconds": round(audio_duration, 2),
        "decode_seconds":   round(elapsed_s, 2),
        "real_time_factor": round(rtf, 3),
        "segment_count":    seg_count,
        "segments":         segments,
        "words":            words,
        "diagnostics": {
            "vad_fallback_used": vad_fallback_used,
            "note": note,
        },
    }


# ══════════════════════════════════════════════════════════════════════════
# Gradio UI
# ══════════════════════════════════════════════════════════════════════════

def gradio_decode(wav_file, language_label, progress=gr.Progress()):
    if wav_file is None:
        return "❌ Please upload a WAV file"

    _job_started()
    try:
        file_size = os.path.getsize(wav_file)
        if file_size > MAX_FILE_BYTES:
            return (f"❌ File too large: {file_size/1e6:.0f} MB "
                    f"(limit {MAX_FILE_BYTES/1e6:.0f} MB)")

        progress(0, desc="Checking duration...")
        duration_check = get_wav_duration_seconds(wav_file)
        if duration_check is not None and duration_check > MAX_AUDIO_SECONDS:
            return (f"❌ Audio too long: {duration_check/60:.1f} min "
                    f"(limit {MAX_AUDIO_SECONDS/60:.0f} min)")

        language = LANGUAGES.get(language_label, None)

        result = None
        for event in transcribe_audio_stream(wav_file, language):
            if event["status"] == "progress":
                progress(event["fraction"], desc=event["desc"])
            elif event["status"] == "done":
                result = event

        # Cap displayed lines for extremely long transcripts
        MAX_DISPLAY_WORDS = 20000
        words = result["words"]
        truncated = len(words) > MAX_DISPLAY_WORDS
        display_words = words[:MAX_DISPLAY_WORDS]

        lines = "\n".join(
            f"{fmt_time(w['start'])} → {fmt_time(w['end'])}   {w['word']}"
            for w in display_words
        )
        if truncated:
            lines += f"\n\n… truncated, {len(words) - MAX_DISPLAY_WORDS} more words not shown …"

        diag = result["diagnostics"]
        warning = ""
        if not words:
            warning = f"\n⚠️ No words detected.\n   {diag['note'] or ''}\n"

        return (
            f"🌐 {result['language']} ({result['language_prob']:.0%})   "
            f"⏱ {result['duration_seconds']/60:.1f} min   "
            f"📦 {file_size/1e6:.0f} MB   "
            f"🧩 {result['segment_count']} segments   "
            f"🚀 {result['decode_seconds']:.1f}s (RTF={result['real_time_factor']:.2f}x)\n"
            f"{warning}\n"
            f"{lines}"
        )

    except MemoryError:
        return "❌ Out of memory — file too large for this Space's RAM. Try a smaller/shorter file."
    except ValueError as e:
        return f"❌ {e}"
    except Exception as e:
        return f"❌ Error: {e}"
    finally:
        _job_finished()


# ══════════════════════════════════════════════════════════════════════════
# API functions — exposed via gr.api(), Gradio's own supported mechanism
# for adding callable endpoints outside the visual UI.
#
# WHY gr.api() INSTEAD OF RAW FASTAPI ROUTES:
# HF Spaces' frontend proxy only recognizes Gradio's own known route
# patterns (/, /queue/, /assets/, named API endpoints registered through
# Gradio itself, etc). Custom routes added directly via
# `demo.app.add_api_route(...)` or a separately-mounted FastAPI app can be
# silently swallowed by that proxy layer. gr.api() registers the endpoint
# through Gradio's OWN routing system, so the same proxy that correctly
# serves the UI also correctly serves this endpoint, with zero extra
# reverse-proxy config.
#
# CALLING CONVENTION:
# gr.api()-registered functions are plain Python functions (not async,
# no FastAPI Request/UploadFile/Form types). File inputs arrive as a
# filesystem path (string) — Gradio handles the upload and saves it to a
# temp file for you, then passes the path in.
# ══════════════════════════════════════════════════════════════════════════

def api_health() -> dict:
    """Health check. Returns service status."""
    return {"status": "ok", "model": "whisper-small", "device": "cpu", "compute_type": "int8"}


def _resolve_file_path(wav_file: Union[str, dict]) -> str:
    """
    gr.api() functions with a plain `str` type hint (no bound gr.File()
    component) do NOT get Gradio's automatic FileData→path preprocessing.
    When called via gradio_client's handle_file(), the client sends a
    FileData-shaped dict instead of a plain path string:
        {"path": "...", "url": ..., "orig_name": ..., "meta": {...}}
    This unwraps that dict if present; passes through a plain string
    unchanged (e.g. when called directly from Python or the Gradio UI).
    """
    if isinstance(wav_file, dict):
        path = wav_file.get("path")
        if not path:
            raise ValueError(f"File dict missing 'path' key: {wav_file!r}")
        return path
    return wav_file


def api_transcribe(wav_file: Union[str, dict], language: str = "") -> dict:
    """
    Upload a WAV file + language code, get back JSON with segments + word
    timestamps.

    THIS IS A GENERATOR FUNCTION. gr.api() automatically streams generator
    yields as Server-Sent Events — each yield becomes one SSE `data:` event
    the client (Android, gradio_client, curl, etc.) can read while polling.
    This is what lets a client show a REAL transcription percentage instead
    of guessing one from elapsed poll attempts.

    Event shapes yielded, in order:
      {"status": "progress", "fraction": 0.05, "desc": "Loading audio..."}
      {"status": "progress", "fraction": 0.42, "desc": "Transcribed 00:08.1 / 00:19.3"}
      ... (one progress event per completed Whisper segment) ...
      {"status": "progress", "fraction": 1.0,  "desc": "Done"}
      {"status": "done", "language": "en", "words": [...], ...}   ← FINAL event

    On error, the FIRST and ONLY event yielded is:
      {"status": "error", "error": "<message>"}

    Python client example (reading every streamed event):
      from gradio_client import Client, handle_file
      client = Client("don0726/Whis")
      job = client.submit(
          wav_file=handle_file("audio.wav"),
          language="en",
          api_name="/api_transcribe",
      )
      for event in job:
          print(event)   # each intermediate + the final result

    Args:
        wav_file: path to an uploaded .wav file (Gradio provides this
                  automatically — pass handle_file("local/path.wav") from
                  the client).
        language: ISO code like 'en', or '' / 'auto' for auto-detect.
    """
    snapshot = _get_capacity_snapshot()
    if snapshot["available_slots"] <= 0:
        # All MAX_CONCURRENT_JOBS slots are busy. Gradio's own queue would
        # normally just make this request wait, but we surface an explicit
        # signal here too so clients can show "server busy" immediately
        # instead of silently waiting with no explanation.
        yield {
            "status": "error",
            "error": (
                f"Server busy: all {snapshot['max_concurrent_jobs']} "
                f"transcription slots are in use. Please try again shortly."
            ),
        }
        return

    _job_started()
    try:
        resolved_path = _resolve_file_path(wav_file)
        file_size = os.path.getsize(resolved_path)
        if file_size > MAX_FILE_BYTES:
            yield {"status": "error",
                   "error": f"File too large: {file_size/1e6:.0f} MB (limit {MAX_FILE_BYTES/1e6:.0f} MB)"}
            return

        try:
            lang = resolve_language(language)
        except ValueError as e:
            yield {"status": "error", "error": str(e)}
            return

        duration_check = get_wav_duration_seconds(resolved_path)
        if duration_check is not None and duration_check > MAX_AUDIO_SECONDS:
            yield {"status": "error",
                   "error": f"Audio too long: {duration_check/60:.1f} min (limit {MAX_AUDIO_SECONDS/60:.0f} min)"}
            return

        yield from transcribe_audio_stream(resolved_path, lang)

    except Exception as e:
        yield {"status": "error", "error": f"Internal error: {e}"}
    finally:
        _job_finished()


def api_status() -> dict:
    """
    Check server capacity — how many transcription "channels"/slots are
    available right now. Useful for clients to decide whether to submit a
    job immediately or show a "server busy, try again" message.

    Returns:
      {
        "max_concurrent_jobs": 3,
        "active_jobs": 1,
        "available_slots": 2,
        "total_jobs_served_since_startup": 47
      }
    """
    return _get_capacity_snapshot()


with gr.Blocks(title="Whisper Word Timestamps") as demo:
    gr.Markdown(f"""
    # 🎙️ Word-Level Timestamps
    Upload a **.wav** file → get word + start/end time.
    Decoded directly via faster_whisper (no custom PCM format needed) and
    accelerated with batched inference for speed.
    Handles up to **{MAX_CONCURRENT_JOBS} simultaneous transcriptions**.

    **API users:** call via `gradio_client.Client("don0726/Whis")`, then
    `job = client.submit(wav_file=handle_file("audio.wav"), language="en",
    api_name="/api_transcribe")` and iterate `for event in job:` to see
    live progress + the final result. Check `/api_status` first to see
    current server capacity. See the "View API" link in this page's
    footer for full request/response details.
    """)
    with gr.Row():
        with gr.Column():
            wav_input  = gr.File(label="🎵 WAV audio", file_types=[".wav"])
            lang_input = gr.Dropdown(label="🌐 Language", choices=list(LANGUAGES.keys()), value="Auto Detect")
            btn = gr.Button("🚀 Transcribe", variant="primary")
        with gr.Column():
            out = gr.Textbox(label="Word timestamps", lines=30)

    btn.click(fn=gradio_decode, inputs=[wav_input, lang_input], outputs=out,
              concurrency_limit=MAX_CONCURRENT_JOBS)

    # gr.api() must be called INSIDE the `with gr.Blocks()` context — it
    # needs the active Blocks context to attach the endpoint to. Calling
    # it after the `with` block exits raises:
    #   "Cannot call api() outside of a gradio.Blocks context."
    gr.api(api_health, api_name="api_health")
    gr.api(api_status, api_name="api_status")

    # concurrency_limit=MAX_CONCURRENT_JOBS lets Gradio's own queue run up
    # to N transcriptions truly in parallel (each on its own worker
    # thread). faster_whisper/ctranslate2 is safe to call concurrently on
    # one shared `model` instance. Without this, gr.api() defaults to
    # concurrency_limit=1 — i.e. every request queues behind the previous
    # one even if the server has spare capacity.
    gr.api(api_transcribe, api_name="api_transcribe", concurrency_limit=MAX_CONCURRENT_JOBS)


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