import gradio as gr
import numpy as np
import struct
import time
import os
import json
import threading
from typing import Union
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 = 3

_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 ──────────────────────────────────────────────────
# pcm.bin is 4 bytes/sample at 16kHz → 1 hour of audio ≈ 230 MB.
# Whisper's encoder works in fixed 30s windows regardless of total length,
# so faster_whisper already chunks internally — the real risk for big files
# is RAM (loading the whole array at once) and Gradio's upload size cap.
MAX_AUDIO_SECONDS = 3 * 60 * 60        # 3 hours — adjust to your needs
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 parse_pcm_header(path_or_filelike, file_size: int):
    """
    Reads only the 8-byte header and validates against actual file size.
    Self-correcting: handles endian flips and fully-corrupt headers by
    falling back to deriving n_samples from the real body size.

    Returns (n_samples, audio_duration).
    """
    if file_size < 8:
        raise ValueError("File too small to contain a valid pcm.bin header")

    if hasattr(path_or_filelike, "read"):
        header = path_or_filelike.read(8)
    else:
        with open(path_or_filelike, "rb") as f:
            header = f.read(8)

    n_samples      = struct.unpack_from("<i", header, 0)[0]
    audio_duration = struct.unpack_from("<f", header, 4)[0]

    expected_body = n_samples * 4
    actual_body   = file_size - 8

    if actual_body != expected_body:
        # Strategy 1: simple endian flip
        n_samples_be = struct.unpack_from(">i", header, 0)[0]
        if n_samples_be * 4 == actual_body:
            n_samples = n_samples_be
            audio_duration = struct.unpack_from(">f", header, 4)[0]
            print(f"⚠️  Header was big-endian — corrected. n_samples={n_samples}")
        else:
            # Strategy 2: header unreliable — derive from actual file size
            if actual_body % 4 != 0:
                raise ValueError(
                    f"Body size {actual_body} bytes is not a multiple of 4 "
                    f"(not valid float32 PCM data)"
                )
            n_samples = actual_body // 4
            audio_duration = n_samples / 16000.0
            print(f"⚠️  Header unreliable — derived from file size instead. "
                  f"n_samples={n_samples}, duration={audio_duration:.1f}s")

    if audio_duration > MAX_AUDIO_SECONDS:
        raise ValueError(
            f"Audio too long: {audio_duration/60:.1f} min "
            f"(limit {MAX_AUDIO_SECONDS/60:.0f} min)"
        )

    return n_samples, audio_duration


def transcribe_pcm_stream(pcm, audio_duration: float, 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 (the same shape transcribe_pcm used
    to return directly):
        {"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 Android's
    poll loop can show a real percentage instead of a guessed one).

    pcm may be a real np.ndarray OR a np.memmap — both work transparently
    with faster_whisper/ctranslate2.

    Includes automatic VAD fallback: if VAD filters out all speech
    (a common cause of empty results on quiet/compressed/phone audio),
    automatically retries once with VAD disabled.
    """
    t0 = time.time()

    pcm_max = float(np.max(np.abs(pcm))) if pcm.size > 0 else 0.0
    pcm_rms = float(np.sqrt(np.mean(np.asarray(pcm, dtype=np.float64) ** 2))) if pcm.size > 0 else 0.0

    def _run(vad_filter: bool, vad_params):
        segments_iter, info = model.transcribe(
            audio=pcm,
            language=language,
            beam_size=1,
            word_timestamps=True,
            vad_filter=vad_filter,
            vad_parameters=vad_params,
            condition_on_previous_text=False,
        )
        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))
                yield {
                    "status":   "progress",
                    "fraction": round(frac, 4),
                    "desc":     f"Transcribed {fmt_time(seg.end)} / {fmt_time(audio_duration)}",
                }
        return_value["segments"] = segments
        return_value["words"] = words
        return_value["info"] = info
        return_value["seg_count"] = seg_count

    yield {"status": "progress", "fraction": 0.1, "desc": "Transcribing..."}

    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"]

    vad_fallback_used = False
    if not segments:
        vad_fallback_used = True
        yield {"status": "progress", "fraction": 0.5,
               "desc": "No speech detected with VAD — 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"]

    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 and pcm_max < 0.01:
        note = "Audio may be silent or near-silent (peak amplitude is very low)."
    elif vad_fallback_used and words:
        note = "VAD initially filtered everything out; 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": {
            "pcm_peak_amplitude": round(pcm_max, 4),
            "pcm_rms_amplitude":  round(pcm_rms, 4),
            "vad_fallback_used":  vad_fallback_used,
            "note": note,
        },
    }


# ══════════════════════════════════════════════════════════════════════════
# Gradio UI — big-file safe (memory-mapped read, streaming progress)
# ══════════════════════════════════════════════════════════════════════════

def gradio_decode(pcm_file, language_label, progress=gr.Progress()):
    if pcm_file is None:
        return "❌ Please upload pcm.bin"

    _job_started()
    try:
        file_size = os.path.getsize(pcm_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="Reading header...")
        n_samples, audio_duration = parse_pcm_header(pcm_file, file_size)

        # Memory-map the PCM body instead of read()-ing it all in — avoids
        # holding a full duplicate copy of a multi-hundred-MB file in RAM.
        progress(0.05, desc=f"Loading {audio_duration:.0f}s of audio...")
        pcm = np.memmap(pcm_file, dtype=np.float32, mode="r",
                        offset=8, shape=(n_samples,))

        language = LANGUAGES.get(language_label, None)

        result = None
        for event in transcribe_pcm_stream(pcm, audio_duration, 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"
                f"   peak amplitude = {diag['pcm_peak_amplitude']}  (near 0 = silent/empty audio)\n"
                f"   rms amplitude  = {diag['pcm_rms_amplitude']}\n"
                f"   vad fallback used = {diag['vad_fallback_used']}\n"
                f"   {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 — which is exactly what produced
# the empty 200 response. 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. This actually simplifies
# our code since we can reuse the same memmap-based reading path as the
# Gradio UI function above.
# ══════════════════════════════════════════════════════════════════════════

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


def _resolve_file_path(pcm_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(pcm_file, dict):
        path = pcm_file.get("path")
        if not path:
            raise ValueError(f"File dict missing 'path' key: {pcm_file!r}")
        return path
    return pcm_file


def api_transcribe(pcm_file: Union[str, dict], language: str = "") -> dict:
    """
    Upload pcm.bin + 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 Android show a REAL transcription percentage instead
    of guessing one from elapsed poll attempts.

    Event shapes yielded, in order:
      {"status": "progress", "fraction": 0.1,  "desc": "Transcribing..."}
      {"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(
          pcm_file=handle_file("pcm.bin"),
          language="en",
          api_name="/api_transcribe",
      )
      for event in job:
          print(event)   # each intermediate + the final result

    Args:
        pcm_file: path to an uploaded pcm.bin file (Gradio provides this
                  automatically — pass handle_file("local/path.bin") 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 (like the Android app) 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(pcm_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

        try:
            n_samples, audio_duration = parse_pcm_header(resolved_path, file_size)
        except ValueError as e:
            yield {"status": "error", "error": f"Invalid pcm.bin: {e}"}
            return

        # Memory-mapped read — same big-file-safe path as the Gradio UI
        pcm = np.memmap(resolved_path, dtype=np.float32, mode="r",
                        offset=8, shape=(n_samples,))

        yield from transcribe_pcm_stream(pcm, audio_duration, 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 `pcm.bin` → get word + start/end time.
    Supports large files (memory-mapped, streamed progress).
    Handles up to **{MAX_CONCURRENT_JOBS} simultaneous transcriptions**.

    **API users:** call via `gradio_client.Client("don0726/Whis")`, then
    `job = client.submit(pcm_file=handle_file("pcm.bin"), 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():
            pcm_input  = gr.File(label="📁 pcm.bin", file_types=[".bin"])
            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=[pcm_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()