"""
Whisper Transcription Space
- Accepts YouTube or Google Drive links
- Runs in background (survives page close)
- CPU capped at 50% via threading limits + cpulimit
- UUID seed saved locally for job recovery
- Outputs .txt with speaker labels + timestamps
- Optionally pushes to HF Dataset
"""

import os
import re
import uuid
import json
import time
import sqlite3
import threading
import subprocess
import tempfile
import shutil
from pathlib import Path
from datetime import datetime

import gradio as gr
import urllib.request
import urllib.error

# Compatibility shim — HfFolder removed in huggingface_hub>=0.24
try:
    from huggingface_hub import HfFolder  # noqa: F401
except ImportError:
    import huggingface_hub as _hfhub
    class _HfFolder:
        @staticmethod
        def get_token(): return _hfhub.utils.get_token()
    _hfhub.HfFolder = _HfFolder
    import sys, types
    sys.modules.setdefault("huggingface_hub.hf_api", types.ModuleType("huggingface_hub.hf_api"))
    sys.modules["huggingface_hub.hf_api"].HfFolder = _HfFolder

# ── Paths ─────────────────────────────────────────────────────────────────────
DB_PATH   = Path("/data/jobs.db")          # persisted HF /data volume
OUT_DIR   = Path("/data/outputs")
AUDIO_DIR = Path("/data/audio")

DB_PATH.parent.mkdir(parents=True, exist_ok=True)
OUT_DIR.mkdir(parents=True, exist_ok=True)
AUDIO_DIR.mkdir(parents=True, exist_ok=True)

# ── CPU cap: limit Whisper threads to ~50% of available cores ─────────────────
import psutil
_TOTAL_CORES   = psutil.cpu_count(logical=True) or 2
_WHISPER_THREADS = max(1, _TOTAL_CORES // 2)   # 50% of logical cores

# ── Database ──────────────────────────────────────────────────────────────────
def get_db():
    conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    with get_db() as db:
        db.execute("""
            CREATE TABLE IF NOT EXISTS jobs (
                id          TEXT PRIMARY KEY,
                url         TEXT NOT NULL,
                status      TEXT NOT NULL DEFAULT 'queued',
                created_at  TEXT NOT NULL,
                updated_at  TEXT NOT NULL,
                log         TEXT DEFAULT '',
                output_path TEXT DEFAULT '',
                error       TEXT DEFAULT ''
            )
        """)
        db.commit()

init_db()

# ── Job helpers ───────────────────────────────────────────────────────────────
def create_job(job_id: str, url: str):
    now = datetime.utcnow().isoformat()
    with get_db() as db:
        db.execute(
            "INSERT OR IGNORE INTO jobs (id, url, status, created_at, updated_at) VALUES (?,?,?,?,?)",
            (job_id, url, "queued", now, now)
        )
        db.commit()

def update_job(job_id: str, **kwargs):
    kwargs["updated_at"] = datetime.utcnow().isoformat()
    cols = ", ".join(f"{k}=?" for k in kwargs)
    vals = list(kwargs.values()) + [job_id]
    with get_db() as db:
        db.execute(f"UPDATE jobs SET {cols} WHERE id=?", vals)
        db.commit()

def get_job(job_id: str):
    with get_db() as db:
        row = db.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
        return dict(row) if row else None

def append_log(job_id: str, msg: str):
    job = get_job(job_id)
    if job:
        log = (job["log"] or "") + f"[{datetime.utcnow().strftime('%H:%M:%S')}] {msg}\n"
        update_job(job_id, log=log)

# ── URL helpers ───────────────────────────────────────────────────────────────
def is_youtube(url: str) -> bool:
    return bool(re.search(r"(youtube\.com|youtu\.be)", url))

def is_gdrive(url: str) -> bool:
    return bool(re.search(r"drive\.google\.com", url))

def extract_youtube_id(url: str) -> str:
    m = re.search(r"(?:v=|youtu\.be/|embed/)([A-Za-z0-9_-]{11})", url)
    if not m:
        raise ValueError(f"Cannot extract YouTube video ID from: {url}")
    return m.group(1)

def gdrive_direct(url: str) -> str:
    """Convert GDrive share link → direct download URL."""
    m = re.search(r"/file/d/([^/]+)", url)
    if m:
        fid = m.group(1)
        return f"https://drive.google.com/uc?export=download&id={fid}"
    m2 = re.search(r"id=([^&]+)", url)
    if m2:
        return f"https://drive.google.com/uc?export=download&id={m2.group(1)}"
    raise ValueError("Could not parse Google Drive file ID from URL")

# ── Apify downloader ──────────────────────────────────────────────────────────
APIFY_BASE = "https://api.apify.com/v2"

def apify_get_youtube_stream(video_url: str, apify_token: str, job_id: str) -> str:
    """
    Run Apify's YouTube downloader actor (apify/youtube-scraper style).
    We use the synchronous run endpoint so we wait for the result inline.
    Returns a direct audio/video stream URL we can pipe to ffmpeg.

    Actor: marcelo.leal/youtube-downloader   (accepts videoUrl, returns streamUrl)
    Docs: https://apify.com/marcelo.leal/youtube-downloader
    """
    import json as _json

    actor_id = "marcelo.leal~youtube-downloader"
    run_url  = f"{APIFY_BASE}/acts/{actor_id}/run-sync-get-dataset-items"

    payload  = _json.dumps({
        "videoUrl": video_url,
        "quality": "lowest",          # audio-only quality → smallest file
    }).encode()

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {apify_token}",
    }

    append_log(job_id, f"📡 Calling Apify actor: {actor_id}")
    req  = urllib.request.Request(run_url, data=payload, headers=headers, method="POST")

    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            body = _json.loads(resp.read())
    except urllib.error.HTTPError as e:
        raise RuntimeError(f"Apify HTTP {e.code}: {e.read().decode()}")

    # Dataset items → first item with a streamUrl or videoUrl
    items = body if isinstance(body, list) else body.get("items", [])
    if not items:
        raise RuntimeError("Apify returned no dataset items")

    item = items[0]
    stream_url = (
        item.get("streamUrl")
        or item.get("videoUrl")
        or item.get("url")
    )
    if not stream_url:
        raise RuntimeError(f"No stream URL in Apify response: {item}")

    append_log(job_id, f"✅ Apify returned stream URL")
    return stream_url

def stream_to_audio(stream_url: str, out_path: Path, job_id: str):
    """Use ffmpeg to pull stream URL → mp3 audio file."""
    append_log(job_id, "🎵 Converting to audio with ffmpeg...")
    cmd = [
        "ffmpeg", "-y",
        "-i", stream_url,
        "-vn",                  # drop video
        "-ar", "16000",         # 16 kHz — optimal for Whisper
        "-ac", "1",             # mono
        "-b:a", "64k",
        str(out_path),
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(f"ffmpeg failed:\n{proc.stderr[-800:]}")
    append_log(job_id, f"Audio ready: {out_path.name} ({out_path.stat().st_size // 1024} KB)")

# ── Download audio (Apify for YouTube, direct for Drive) ──────────────────────
def download_audio(url: str, job_id: str, apify_token: str = "") -> Path:
    audio_path = AUDIO_DIR / f"{job_id}.mp3"

    if is_youtube(url):
        if not apify_token or not apify_token.strip():
            raise ValueError(
                "An Apify API token is required to download YouTube videos. "
                "Get a free token at https://apify.com and paste it in the Apify Token field."
            )
        append_log(job_id, "🎬 Fetching YouTube video via Apify...")
        stream_url = apify_get_youtube_stream(url, apify_token.strip(), job_id)
        stream_to_audio(stream_url, audio_path, job_id)

    elif is_gdrive(url):
        append_log(job_id, "⬇️ Downloading from Google Drive...")
        direct = gdrive_direct(url)
        tmp_path = AUDIO_DIR / f"{job_id}_raw"
        urllib.request.urlretrieve(direct, str(tmp_path))
        # Convert whatever Drive gave us → wav/mp3 via ffmpeg
        stream_to_audio(str(tmp_path), audio_path, job_id)
        try:
            tmp_path.unlink()
        except Exception:
            pass

    else:
        raise ValueError("URL must be a YouTube or Google Drive link")

    if not audio_path.exists():
        raise FileNotFoundError("Audio conversion produced no file")
    return audio_path

# ── Transcribe with Whisper (CPU-capped) ──────────────────────────────────────
def run_whisper(audio_path: Path, job_id: str, model_size: str, push_dataset: bool, hf_token: str, hf_repo: str):
    """Run whisper in a subprocess with CPU affinity + nice."""
    append_log(job_id, f"Starting Whisper ({model_size}) with {_WHISPER_THREADS} threads (≤50% CPU)...")

    out_base = OUT_DIR / job_id
    out_base.mkdir(parents=True, exist_ok=True)

    # Use whisper CLI (installed via openai-whisper)
    cmd = [
        "nice", "-n", "10",          # lower process priority
        "whisper", str(audio_path),
        "--model", model_size,
        "--language", "auto",
        "--output_format", "all",     # txt, vtt, srt, json
        "--output_dir", str(out_base),
        "--threads", str(_WHISPER_THREADS),
        "--verbose", "False",
    ]

    # Optionally wrap with cpulimit if available
    if shutil.which("cpulimit"):
        cpu_pct = _WHISPER_THREADS * 100 // _TOTAL_CORES  # e.g. 50
        cmd = ["cpulimit", "--limit", str(min(cpu_pct, 50)), "--"] + cmd

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    for line in proc.stdout:
        line = line.strip()
        if line:
            append_log(job_id, line)

    proc.wait()
    if proc.returncode != 0:
        raise RuntimeError(f"Whisper exited with code {proc.returncode}")

    # ── Post-process: build speaker-labelled transcript ───────────────────────
    json_files = list(out_base.glob("*.json"))
    txt_out    = out_base / "transcript_with_speakers.txt"

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

        segments = data.get("segments", [])
        lines = []
        for i, seg in enumerate(segments):
            start = _fmt_time(seg["start"])
            end   = _fmt_time(seg["end"])
            text  = seg["text"].strip()
            # Whisper doesn't do diarization natively; label as SPEAKER_XX by gap heuristic
            speaker = _detect_speaker_change(segments, i)
            lines.append(f"[{start} → {end}] {speaker}: {text}")

        txt_out.write_text("\n".join(lines), encoding="utf-8")
        append_log(job_id, f"Transcript saved: {txt_out.name}")
    else:
        # fallback: use the plain .txt whisper produced
        plain = list(out_base.glob("*.txt"))
        if plain:
            txt_out = plain[0]

    # ── Optionally push to HF Dataset ─────────────────────────────────────────
    if push_dataset and hf_token and hf_repo and txt_out.exists():
        _push_to_hf(txt_out, job_id, hf_token, hf_repo)

    return txt_out

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_change(segments, idx, gap_threshold=1.5):
    """Heuristic: new speaker when gap > threshold seconds."""
    if idx == 0:
        return "SPEAKER_01"
    gap = segments[idx]["start"] - segments[idx - 1]["end"]
    prev_tag = _detect_speaker_change(segments, idx - 1) if idx > 1 else "SPEAKER_01"
    if gap > gap_threshold:
        # toggle between SPEAKER_01 and SPEAKER_02 on large gaps
        return "SPEAKER_02" if prev_tag == "SPEAKER_01" else "SPEAKER_01"
    return prev_tag

def _push_to_hf(txt_path: Path, job_id: str, token: str, repo: str):
    from huggingface_hub import HfApi
    api = HfApi(token=token)
    try:
        api.upload_file(
            path_or_fileobj=str(txt_path),
            path_in_repo=f"{job_id}/{txt_path.name}",
            repo_id=repo,
            repo_type="dataset",
        )
        append_log(job_id, f"Pushed to HF Dataset: {repo}")
    except Exception as e:
        append_log(job_id, f"HF push failed: {e}")

# ── Background worker ─────────────────────────────────────────────────────────
_worker_lock = threading.Lock()

def process_job(job_id: str, url: str, model_size: str, push_dataset: bool,
                hf_token: str, hf_repo: str, apify_token: str = ""):
    try:
        update_job(job_id, status="downloading")
        audio_path = download_audio(url, job_id, apify_token=apify_token)

        update_job(job_id, status="transcribing")
        txt_out = run_whisper(audio_path, job_id, model_size, push_dataset, hf_token, hf_repo)

        update_job(job_id, status="done", output_path=str(txt_out))
        append_log(job_id, "✅ Job complete!")

        try:
            audio_path.unlink()
        except Exception:
            pass

    except Exception as e:
        import traceback
        err = traceback.format_exc()
        update_job(job_id, status="error", error=str(e))
        append_log(job_id, f"❌ Error: {e}\n{err}")

def submit_job(job_id, url, model_size, push_dataset, hf_token, hf_repo, apify_token=""):
    t = threading.Thread(
        target=process_job,
        args=(job_id, url, model_size, push_dataset, hf_token, hf_repo, apify_token),
        daemon=True,
    )
    t.start()

# ── Gradio UI ─────────────────────────────────────────────────────────────────
CUSTOM_CSS = """
:root {
    --bg: #0a0a0a;
    --surface: #111111;
    --border: #222222;
    --accent: #f0f0f0;
    --muted: #555555;
    --green: #4ade80;
    --yellow: #facc15;
    --red: #f87171;
    --blue: #60a5fa;
    --font: 'JetBrains Mono', 'Courier New', monospace;
}
body, .gradio-container { background: var(--bg) !important; color: var(--accent) !important; font-family: var(--font) !important; }
.gr-box, .gr-form, .gr-panel, .block { background: var(--surface) !important; border: 1px solid var(--border) !important; border-radius: 4px !important; }
.gr-button-primary { background: var(--accent) !important; color: #000 !important; font-family: var(--font) !important; font-weight: 700 !important; border: none !important; border-radius: 2px !important; letter-spacing: 0.1em !important; }
.gr-button-secondary { background: transparent !important; color: var(--accent) !important; font-family: var(--font) !important; border: 1px solid var(--border) !important; border-radius: 2px !important; }
input, textarea, select { background: #0d0d0d !important; color: var(--accent) !important; border: 1px solid var(--border) !important; font-family: var(--font) !important; border-radius: 2px !important; }
label, .gr-label { color: var(--muted) !important; font-size: 0.75rem !important; letter-spacing: 0.12em !important; text-transform: uppercase !important; }
#status_badge { font-size: 0.85rem; letter-spacing: 0.08em; padding: 4px 10px; border-radius: 2px; display: inline-block; }
"""

JS_SEED = """
function() {
    // Generate or retrieve persistent UUID seed
    let seed = localStorage.getItem('whisper_job_seed');
    if (!seed) {
        seed = crypto.randomUUID();
        localStorage.setItem('whisper_job_seed', seed);
    }
    return seed;
}
"""

JS_SAVE_SEED = """
function(seed) {
    localStorage.setItem('whisper_job_seed', seed);
    return seed;
}
"""

def submit_handler(seed, url, model_size, push_ds, hf_token, hf_repo, apify_token):
    if not url or not url.strip():
        return seed, "❌ Please enter a URL", "", None, gr.update(visible=False)
    if not is_youtube(url) and not is_gdrive(url):
        return seed, "❌ Only YouTube or Google Drive URLs are supported", "", None, gr.update(visible=False)
    if is_youtube(url) and not (apify_token or "").strip():
        return seed, "❌ Apify token required for YouTube links — get one free at apify.com", "", None, gr.update(visible=False)

    job = get_job(seed)
    if job and job["status"] in ("queued", "downloading", "transcribing"):
        return seed, f"⏳ Already running — Job `{seed[:8]}…`", job["log"], None, gr.update(visible=False)

    if job and job["status"] == "done":
        out = Path(job["output_path"]) if job["output_path"] else None
        return seed, "✅ Already complete", job["log"], str(out) if out and out.exists() else None, gr.update(visible=out is not None and out.exists())

    create_job(seed, url)
    submit_job(seed, url, model_size, push_ds, hf_token, hf_repo, apify_token)
    return seed, f"🚀 Job started — ID: `{seed[:8]}…`", "", None, gr.update(visible=False)

def poll_status(seed):
    if not seed:
        return "—", "", None, gr.update(visible=False)
    job = get_job(seed)
    if not job:
        return "No job found", "", None, gr.update(visible=False)

    status = job["status"]
    icons = {"queued": "🕐", "downloading": "⬇️", "transcribing": "🔊", "done": "✅", "error": "❌"}
    label = f"{icons.get(status, '?')} {status.upper()}"

    out_path = job["output_path"]
    file_val = None
    show_dl  = False
    if status == "done" and out_path and Path(out_path).exists():
        file_val = out_path
        show_dl  = True

    return label, job["log"] or "", file_val, gr.update(visible=show_dl)

def reset_seed():
    new_seed = str(uuid.uuid4())
    return new_seed, "—", "", None, gr.update(visible=False)

def recover_handler(seed):
    return poll_status(seed)

with gr.Blocks(css=CUSTOM_CSS, title="Whisper Space") as demo:

    # ── Persistent seed stored in browser localStorage ─────────────────────
    seed_state = gr.State("")

    gr.HTML("""
    <div style="padding:28px 0 8px 0">
        <div style="font-size:0.7rem;letter-spacing:0.2em;color:#444;text-transform:uppercase;margin-bottom:6px">LEXICAL SPACE</div>
        <div style="font-size:1.6rem;font-weight:700;letter-spacing:0.05em;color:#f0f0f0">WHISPER TRANSCRIPTION</div>
        <div style="font-size:0.78rem;color:#555;margin-top:4px;letter-spacing:0.08em">Background processing · Survives page close · Speaker + timestamp output</div>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=3):
            url_input = gr.Textbox(
                placeholder="https://youtube.com/watch?v=... or https://drive.google.com/file/d/...",
                label="Video / Audio URL",
                lines=1,
            )
            with gr.Row():
                model_dd = gr.Dropdown(
                    ["tiny", "base", "small", "medium", "large"],
                    value="base",
                    label="Whisper Model",
                )
                push_ds_cb = gr.Checkbox(label="Push to HF Dataset", value=False)

            with gr.Accordion("YouTube via Apify (required for YouTube)", open=True):
                apify_token_box = gr.Textbox(
                    label="Apify API Token",
                    type="password",
                    placeholder="apify_api_xxxx...",
                    info="Free at apify.com → Settings → Integrations → API tokens",
                )

            with gr.Accordion("HuggingFace Dataset push (optional)", open=False):
                hf_token_box = gr.Textbox(label="HF Token", type="password", placeholder="hf_...")
                hf_repo_box  = gr.Textbox(label="Dataset Repo", placeholder="username/my-transcripts")

            with gr.Row():
                submit_btn = gr.Button("▶  TRANSCRIBE", variant="primary")
                new_btn    = gr.Button("↺  NEW JOB", variant="secondary")

        with gr.Column(scale=2):
            gr.HTML('<div style="font-size:0.7rem;letter-spacing:0.15em;color:#444;margin-bottom:8px">JOB IDENTITY</div>')
            seed_display = gr.Textbox(label="Session Seed (auto-saved to browser)", interactive=False)
            recover_btn  = gr.Button("⟳  Recover Job", variant="secondary", size="sm")

    gr.HTML('<hr style="border-color:#1a1a1a;margin:16px 0">')

    status_box  = gr.Textbox(label="Status", interactive=False, lines=1)
    log_box     = gr.Textbox(label="Processing Log", interactive=False, lines=12, max_lines=30)
    dl_file     = gr.File(label="Download Transcript", visible=False)

    # ── Auto-poll every 3 s while in progress ─────────────────────────────
    timer = gr.Timer(value=3, active=True)

    # ── On page load: read seed from localStorage → display + recover ─────
    demo.load(
        fn=None,
        js=JS_SEED,
        outputs=[seed_state],
    ).then(
        fn=lambda s: s,
        inputs=[seed_state],
        outputs=[seed_display],
    ).then(
        fn=poll_status,
        inputs=[seed_state],
        outputs=[status_box, log_box, dl_file, dl_file],
    )

    # ── Submit ────────────────────────────────────────────────────────────
    submit_btn.click(
        fn=submit_handler,
        inputs=[seed_state, url_input, model_dd, push_ds_cb, hf_token_box, hf_repo_box, apify_token_box],
        outputs=[seed_state, status_box, log_box, dl_file, dl_file],
    ).then(
        fn=None,
        inputs=[seed_state],
        js=JS_SAVE_SEED,
    )

    # ── Auto-poll ─────────────────────────────────────────────────────────
    timer.tick(
        fn=poll_status,
        inputs=[seed_state],
        outputs=[status_box, log_box, dl_file, dl_file],
    )

    # ── Recover ───────────────────────────────────────────────────────────
    recover_btn.click(
        fn=recover_handler,
        inputs=[seed_display],
        outputs=[status_box, log_box, dl_file, dl_file],
    ).then(
        fn=lambda s: s,
        inputs=[seed_display],
        outputs=[seed_state],
    )

    # ── New job ───────────────────────────────────────────────────────────
    new_btn.click(
        fn=reset_seed,
        outputs=[seed_state, status_box, log_box, dl_file, dl_file],
    ).then(
        fn=lambda s: s,
        inputs=[seed_state],
        outputs=[seed_display],
    ).then(
        fn=None,
        inputs=[seed_state],
        js=JS_SAVE_SEED,
    )

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