app.py

"""
Podcast Transcription Labeling Tool
HuggingFace Space — Gradio 6.6+, huggingface_hub 0.24+
SQLite backend, pre-assigned chunks (no race conditions).
465 files split: labeler1=116, labeler2=116, labeler3=116, labeler4=117
"""

import os
import sqlite3
import json
import shutil
import threading
from datetime import datetime
from pathlib import Path
from contextlib import contextmanager

import gradio as gr
from huggingface_hub import hf_hub_download, upload_file, list_repo_files

# ─────────────────────────────────────────────────────────────
# Config  (set these as HF Space Secrets)
# ─────────────────────────────────────────────────────────────
HF_TOKEN      = os.environ.get("HF_TOKEN") or None   # None = no auth (not empty string)
HF_REPO_ID    = os.environ.get("HF_REPO_ID", "hams-ai/stt-dataset-v2-human-hand-labeled")
BATCH_PATH    = "data/batch_0000"

# Read labeler credentials from secrets.
# Keys are the actual usernames; values are passwords.
LABELER_ACCOUNTS = {
    os.environ.get("LABELER1_USER"): os.environ.get("LABELER1_PASS"),
    os.environ.get("LABELER2_USER"): os.environ.get("LABELER2_PASS"),
    os.environ.get("LABELER3_USER"): os.environ.get("LABELER3_PASS"),
    os.environ.get("LABELER4_USER"): os.environ.get("LABELER4_PASS"),
}
# Remove any entries where the env var wasn't set (key is None)
LABELER_ACCOUNTS = {k: v for k, v in LABELER_ACCOUNTS.items() if k}

# Startup check: warn if any two labelers share the same username
if len(LABELER_ACCOUNTS) != len(set(LABELER_ACCOUNTS.keys())):
    print("[startup] WARNING: Duplicate labeler usernames detected — check your Secrets!")

# Ordered list of labelers — order determines round-robin assignment
LABELERS = list(LABELER_ACCOUNTS.keys())

# Debug: printed at import time so it always shows up in Spaces logs
print(f"[startup] Labeler accounts loaded: {LABELERS}")
print(f"[startup] Number of accounts: {len(LABELER_ACCOUNTS)}")
if not LABELER_ACCOUNTS:
    print("[startup] ERROR: No labeler accounts found! Check your Space Secrets.")

# ─────────────────────────────────────────────────────────────
# Paths  (/data persists across Space restarts on HF Spaces)
# ─────────────────────────────────────────────────────────────
DATA_DIR  = Path("/data")
AUDIO_DIR = DATA_DIR / "audio"
DB_PATH   = DATA_DIR / "labeling.db"

DATA_DIR.mkdir(exist_ok=True)
AUDIO_DIR.mkdir(exist_ok=True)

# Lock so concurrent HF pushes don't interleave writes to the JSONL file
_hf_push_lock = threading.Lock()


# ─────────────────────────────────────────────────────────────
# Database helpers
# ─────────────────────────────────────────────────────────────
@contextmanager
def get_db():
    """Thread-safe DB connection. WAL mode lets multiple readers coexist with one writer."""
    conn = sqlite3.connect(str(DB_PATH), timeout=30, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")   # retry for up to 5 s if locked
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def _column_exists(conn, table: str, column: str) -> bool:
    cols = conn.execute(f"PRAGMA table_info({table})").fetchall()
    return any(c["name"] == column for c in cols)


def init_db():
    """
    Create table (or migrate old schema) and populate from HF repo file listing on first run.
    Idempotent — safe to call on every startup.
    Migration handled:
      Old schema had: claimed_by, claimed_at, status='in_progress'
      New schema has: assigned_to  (no claiming needed)
    """
    if not LABELERS:
        raise RuntimeError(
            "No labeler accounts found! "
            "Set LABELER1_USER / LABELER1_PASS (etc.) as Space Secrets."
        )

    with get_db() as conn:
        # Create table if brand new
        conn.execute("""
            CREATE TABLE IF NOT EXISTS chunks (
                filename     TEXT PRIMARY KEY,
                podcast_id   TEXT,
                duration_sec REAL,
                assigned_to  TEXT,
                status       TEXT DEFAULT 'unlabeled',
                transcript   TEXT DEFAULT '',
                labeled_by   TEXT DEFAULT NULL,
                labeled_at   TEXT DEFAULT NULL
            )
        """)

        # Migration: add assigned_to column if upgrading from old schema
        if not _column_exists(conn, "chunks", "assigned_to"):
            print("[init_db] Migrating: adding 'assigned_to' column ...")
            conn.execute("ALTER TABLE chunks ADD COLUMN assigned_to TEXT")

        # Migration: reset any stale 'in_progress' rows from old claiming system
        if _column_exists(conn, "chunks", "claimed_by"):
            stale = conn.execute(
                "UPDATE chunks SET status='unlabeled', claimed_by=NULL, claimed_at=NULL "
                "WHERE status='in_progress'"
            ).rowcount
        else:
            stale = conn.execute(
                "UPDATE chunks SET status='unlabeled' WHERE status='in_progress'"
            ).rowcount
        if stale:
            print(f"[init_db] Reset {stale} stale 'in_progress' rows to 'unlabeled'.")

        count = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
        if count > 0:
            unassigned = conn.execute(
                "SELECT COUNT(*) FROM chunks WHERE assigned_to IS NULL"
            ).fetchone()[0]
            if unassigned == 0:
                print(f"[init_db] DB has {count} chunks, all assigned. Skipping population.")
                return
            else:
                print(f"[init_db] {unassigned} unassigned chunks found — assigning now.")
                _assign_existing_chunks(conn)
                return

    # ── Fresh DB — list .wav files directly from the HF repo ──────────────
    print("[init_db] Fresh DB. Listing .wav files from HF repo ...")
    all_files = list_repo_files(
        repo_id=HF_REPO_ID,
        repo_type="dataset",
        token=HF_TOKEN,
    )
    rows = [
        (Path(f).name, "batch_0000", 0.0)
        for f in sorted(all_files)
        if f.startswith(f"{BATCH_PATH}/") and f.endswith(".wav")
    ]

    if not rows:
        raise RuntimeError(
            f"[init_db] No .wav files found under {BATCH_PATH}/ — check repo path and token."
        )

    print(f"[init_db] Found {len(rows)} .wav files.")

    # Round-robin assignment: 465 files → 116, 116, 116, 117
    assigned_rows = [
        (filename, podcast_id, duration_sec, LABELERS[i % len(LABELERS)])
        for i, (filename, podcast_id, duration_sec) in enumerate(rows)
    ]

    with get_db() as conn:
        conn.executemany(
            "INSERT OR IGNORE INTO chunks (filename, podcast_id, duration_sec, assigned_to) "
            "VALUES (?,?,?,?)",
            assigned_rows,
        )

    print(f"[init_db] Populated {len(assigned_rows)} chunks across {len(LABELERS)} labelers.")
    _log_assignment_counts()


def _assign_existing_chunks(conn):
    """Round-robin assign chunks that have assigned_to = NULL (migration path)."""
    rows = conn.execute(
        "SELECT filename FROM chunks WHERE assigned_to IS NULL ORDER BY rowid ASC"
    ).fetchall()
    updates = [
        (LABELERS[i % len(LABELERS)], row["filename"])
        for i, row in enumerate(rows)
    ]
    conn.executemany("UPDATE chunks SET assigned_to=? WHERE filename=?", updates)
    print(f"[init_db] Assigned {len(updates)} previously-unassigned chunks.")


def _log_assignment_counts():
    for labeler in LABELERS:
        with get_db() as conn:
            c = conn.execute(
                "SELECT COUNT(*) FROM chunks WHERE assigned_to=?", (labeler,)
            ).fetchone()[0]
        print(f"  {labeler}: {c} chunks")


# ─────────────────────────────────────────────────────────────
# Core data functions
# ─────────────────────────────────────────────────────────────
def get_next_for_user(username: str):
    """Return the next unlabeled chunk assigned to this user, or None if all done."""
    with get_db() as conn:
        return conn.execute("""
            SELECT filename, podcast_id, duration_sec, transcript
            FROM chunks
            WHERE assigned_to = ? AND status = 'unlabeled'
            ORDER BY rowid ASC
            LIMIT 1
        """, (username,)).fetchone()


def get_chunk_by_filename(filename: str):
    """Fetch a chunk by filename (used for Previous navigation)."""
    with get_db() as conn:
        return conn.execute("""
            SELECT filename, podcast_id, duration_sec, transcript, status
            FROM chunks WHERE filename = ?
        """, (filename,)).fetchone()


def submit_transcript(filename: str, transcript: str, username: str) -> bool:
    """
    Persist the transcript to SQLite, then fire a background HF push.
    Returns False if transcript is empty (caller should guard before calling).
    Data is committed to SQLite before this function returns — HF push is best-effort.
    """
    if not transcript.strip():
        return False
    with get_db() as conn:
        conn.execute("""
            UPDATE chunks
            SET status     = 'done',
                transcript = ?,
                labeled_by = ?,
                labeled_at = ?
            WHERE filename = ?
        """, (transcript.strip(), username, datetime.utcnow().isoformat(), filename))

    # Fire HF push in background — UI responds immediately, push happens async
    threading.Thread(target=push_results_to_hf, daemon=True).start()
    return True


def get_progress(username: str = None) -> dict:
    """Return overall + per-labeler progress counts."""
    with get_db() as conn:
        row = conn.execute("""
            SELECT
                COUNT(*) as total,
                SUM(CASE WHEN status='done'      THEN 1 ELSE 0 END) as done,
                SUM(CASE WHEN status='unlabeled' THEN 1 ELSE 0 END) as unlabeled
            FROM chunks
        """).fetchone()

        per_labeler = conn.execute("""
            SELECT assigned_to,
                   COUNT(*) as total,
                   SUM(CASE WHEN status='done' THEN 1 ELSE 0 END) as done
            FROM chunks
            GROUP BY assigned_to
            ORDER BY assigned_to
        """).fetchall()

    return {
        "total":        row["total"],
        "done":         row["done"] or 0,
        "unlabeled":    row["unlabeled"] or 0,
        "per_labeler":  [(r["assigned_to"], r["done"] or 0, r["total"]) for r in per_labeler],
        "current_user": username,
    }


# ─────────────────────────────────────────────────────────────
# HuggingFace push  (always runs in a background daemon thread)
# ─────────────────────────────────────────────────────────────
def push_results_to_hf():
    """
    Export all 'done' rows to JSONL and push to the HF dataset repo.
    - Threading lock ensures concurrent submits don't interleave file writes.
    - Writes to .tmp first then atomic rename so file is never half-written.
    - Push failures are logged but non-fatal — data is already safe in SQLite.

    Output file on HF repo: data/transcripts.jsonl
    Each line: {"filename": "...", "podcast_id": "batch_0000", "duration_sec": 0.0,
                 "transcript": "...", "labeled_by": "...", "labeled_at": "..."}
    """
    with _hf_push_lock:
        try:
            with get_db() as conn:
                rows = conn.execute("""
                    SELECT filename, podcast_id, duration_sec,
                           transcript, labeled_by, labeled_at
                    FROM chunks WHERE status = 'done'
                    ORDER BY rowid ASC
                """).fetchall()

            tmp_path = DATA_DIR / "transcripts.jsonl.tmp"
            out_path = DATA_DIR / "transcripts.jsonl"

            with open(tmp_path, "w", encoding="utf-8") as f:
                for r in rows:
                    f.write(json.dumps(dict(r), ensure_ascii=False) + "\n")

            # Atomic rename — readers never see a partially written file
            tmp_path.replace(out_path)

            upload_file(
                path_or_fileobj=str(out_path),
                path_in_repo="data/transcripts.jsonl",
                repo_id=HF_REPO_ID,
                repo_type="dataset",
                token=HF_TOKEN,
            )
            print(f"[HF push] ✓ {len(rows)} transcripts uploaded successfully.")
        except Exception as e:
            print(f"[HF push] ✗ Push failed: {e}. Data is safe in SQLite.")


# ─────────────────────────────────────────────────────────────
# Audio helpers
# ─────────────────────────────────────────────────────────────
def get_audio_path(filename: str) -> tuple:
    """
    Return (local_path_or_None, error_message_or_empty_string).
    Files are cached in AUDIO_DIR (/data/audio) which persists across Space restarts.
    """
    local_path = AUDIO_DIR / filename
    if local_path.exists():
        return str(local_path), ""

    try:
        print(f"[audio] Downloading {filename} from HF ...")
        downloaded = hf_hub_download(
            repo_id=HF_REPO_ID,
            filename=f"{BATCH_PATH}/{filename}",
            repo_type="dataset",
            token=HF_TOKEN,
        )
        shutil.copy(downloaded, local_path)
        return str(local_path), ""
    except Exception as e:
        msg = f"⚠️ Could not load audio `{filename}`: {e}"
        print(f"[audio] ERROR: {msg}")
        return None, msg


# ─────────────────────────────────────────────────────────────
# UI helpers
# ─────────────────────────────────────────────────────────────
def build_progress_md(p: dict) -> str:
    done  = p["done"]
    total = p["total"]
    pct   = int(100 * done / total) if total else 0
    filled = int(pct / 5)
    bar   = "█" * filled + "░" * (20 - filled)

    labeler_lines = []
    for name, ldone, ltotal in p["per_labeler"]:
        lpct   = int(100 * ldone / ltotal) if ltotal else 0
        marker = " ← you" if name == p.get("current_user") else ""
        labeler_lines.append(f"  - **{name}**: {ldone}/{ltotal} ({lpct}%){marker}")
    labeler_block = "\n".join(labeler_lines) or "  - No data yet"

    return f"""
### 📊 Progress Dashboard
`{bar}` **{pct}%** — {done} of {total} done
**Per labeler:**
{labeler_block}
"""


def build_chunk_info_md(row) -> str:
    if row is None:
        return ""
    return (
        f"**File:** `{row['filename']}`  |  "
        f"**Batch:** {row['podcast_id']}  |  "
        f"**Duration:** {row['duration_sec']:.1f}s"
    )


# ─────────────────────────────────────────────────────────────
# Gradio app
# ─────────────────────────────────────────────────────────────
APP_CSS = """
    #submit-btn { background: #16a34a !important; color: white !important; }
    #submit-btn:hover { background: #15803d !important; }
    #prev-btn { background: #6b7280 !important; color: white !important; }
    #audio-player { min-height: 80px; }
"""


def build_app():
    with gr.Blocks(
        title="Podcast Transcription Tool",
        theme=gr.themes.Soft(),
        css=APP_CSS,
    ) as demo:

        # Per-session state: filenames visited this session + current index.
        # Resets on every page load/refresh — intentional.
        history  = gr.State([])
        hist_pos = gr.State(-1)

        gr.Markdown("# 🎙️ Podcast Transcription Tool")
        gr.Markdown(
            "Listen to each audio clip and type **exactly** what you hear in Saudi Arabic. "
            "Press **Submit & Next** when done, or **Previous** to fix your last answer."
        )

        progress_display = gr.Markdown()
        gr.Markdown("---")

        chunk_info = gr.Markdown(value="*Loading your first chunk…*")

        audio_player = gr.Audio(
            label="Audio",
            elem_id="audio-player",
            interactive=False,
            autoplay=True,
        )

        transcript_box = gr.Textbox(
            label="اكتب ما تسمعه بالضبط بالعربية السعودية",
            placeholder="اكتب كل شيء تسمعه، بما في ذلك كلمات التردد مثل: آآ، أم، إيه، هاه ...",
            lines=3,
            rtl=True,
            interactive=True,
        )
        status_msg = gr.Markdown(value="")

        with gr.Row():
            prev_btn   = gr.Button("⬅️ Previous",     elem_id="prev-btn",   scale=1)
            submit_btn = gr.Button("✅ Submit & Next", elem_id="submit-btn", scale=2)

        # Every handler returns exactly these 7 values
        outputs = [
            history, hist_pos,
            chunk_info, audio_player,
            transcript_box, status_msg,
            progress_display,
        ]

        # ── Handlers ─────────────────────────────────────────────────────

        def load_next(history, hist_pos, request: gr.Request):
            """Load the next unlabeled chunk for this user."""
            user = request.username
            row  = get_next_for_user(user)
            p    = get_progress(user)

            if row is None:
                return (
                    history, hist_pos,
                    "🎉 **You've finished all your chunks! Great work.**",
                    None, "",
                    "✅ All done!",
                    build_progress_md(p),
                )

            filename        = row["filename"]
            audio_path, err = get_audio_path(filename)
            new_history     = history + [filename]

            return (
                new_history, len(new_history) - 1,
                build_chunk_info_md(row),
                audio_path,
                row["transcript"] or "",
                err,
                build_progress_md(p),
            )

        def load_previous(history, hist_pos, request: gr.Request):
            """
            Go one step back in session history.
            If already at the first chunk, show a warning but keep current state visible.
            """
            user = request.username

            if hist_pos <= 0:
                if hist_pos == 0 and history:
                    cur_row       = get_chunk_by_filename(history[0])
                    cur_audio, _  = get_audio_path(history[0])
                    return (
                        history, hist_pos,
                        build_chunk_info_md(cur_row),
                        cur_audio,
                        cur_row["transcript"] if cur_row else "",
                        "⚠️ This is the first chunk in your session — no previous chunk.",
                        build_progress_md(get_progress(user)),
                    )
                return (
                    history, hist_pos,
                    "*No chunk loaded yet.*",
                    None, "",
                    "⚠️ No previous chunk in this session.",
                    build_progress_md(get_progress(user)),
                )

            new_pos         = hist_pos - 1
            filename        = history[new_pos]
            row             = get_chunk_by_filename(filename)
            audio_path, err = get_audio_path(filename)
            status          = err if err else f"↩️ Loaded: `{filename}`. Edit and resubmit to update."

            return (
                history, new_pos,
                build_chunk_info_md(row),
                audio_path,
                row["transcript"] if row else "",
                status,
                build_progress_md(get_progress(user)),
            )

        def submit(transcript, history, hist_pos, request: gr.Request):
            """
            Save transcript for the current chunk, then load the next unlabeled one.
            Handles both first-time labeling and corrections via Previous + resubmit.
            """
            user = request.username

            if hist_pos < 0 or not history:
                return (
                    history, hist_pos,
                    "*No chunk loaded.*",
                    None, "",
                    "⚠️ No chunk is loaded. Please wait for the page to finish loading.",
                    build_progress_md(get_progress(user)),
                )

            filename = history[hist_pos]

            if not transcript.strip():
                row           = get_chunk_by_filename(filename)
                audio_path, _ = get_audio_path(filename)
                return (
                    history, hist_pos,
                    build_chunk_info_md(row),
                    audio_path,
                    transcript,
                    "⚠️ Transcript cannot be empty. Please type what you hear.",
                    build_progress_md(get_progress(user)),
                )

            submit_transcript(filename, transcript, user)

            row = get_next_for_user(user)
            p   = get_progress(user)

            if row is None:
                return (
                    history, hist_pos,
                    "🎉 **You've finished all your chunks! Amazing work.**",
                    None, "",
                    f"✅ Saved `{filename}`. You're all done!",
                    build_progress_md(p),
                )

            new_filename    = row["filename"]
            audio_path, err = get_audio_path(new_filename)
            new_history     = history + [new_filename]
            status          = f"✅ Saved `{filename}`" + (f" — but: {err}" if err else "")

            return (
                new_history, len(new_history) - 1,
                build_chunk_info_md(row),
                audio_path,
                row["transcript"] or "",
                status,
                build_progress_md(p),
            )

        # ── Wire events ───────────────────────────────────────────────────
        demo.load(
            fn=load_next,
            inputs=[history, hist_pos],
            outputs=outputs,
            concurrency_limit=None,
        )
        submit_btn.click(
            fn=submit,
            inputs=[transcript_box, history, hist_pos],
            outputs=outputs,
            concurrency_limit=None,
        )
        prev_btn.click(
            fn=load_previous,
            inputs=[history, hist_pos],
            outputs=outputs,
            concurrency_limit=None,
        )
        transcript_box.submit(
            fn=submit,
            inputs=[transcript_box, history, hist_pos],
            outputs=outputs,
            concurrency_limit=None,
        )

    return demo


# ─────────────────────────────────────────────────────────────
# Module-level startup
# NOTE: Do NOT wrap in `if __name__ == "__main__"` on HF Spaces —
# the runner imports the file, so that block would be silently skipped.
# ─────────────────────────────────────────────────────────────
print("[startup] Initializing database...")
init_db()

print("[startup] Building Gradio app...")
demo = build_app()

print("[startup] Launching...")
demo.launch(
    auth=list(LABELER_ACCOUNTS.items()),
    auth_message="🎙️ Podcast Transcription Tool — Enter your credentials",
    show_error=True,
    allowed_paths=[str(AUDIO_DIR)],
)