import io
import os
import tempfile
import time
import datetime
import threading
import concurrent.futures
import gradio as gr
import numpy as np
import pandas as pd
import requests
import soundfile as sf
import torchaudio
import torch
from sklearn.cluster import AgglomerativeClustering
from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector

os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")

N_CPUS = os.cpu_count() or 2
N_WORKERS = max(1, N_CPUS // 2)  # 8 workers on 16-CPU Space
EMBED_THREADS = 2                  # per-worker thread count — NUMA sweet spot
API_TIMEOUT = 12
STREAMING_TIMEOUT = 120

# Set thread count before model load so ORT/MKL picks it up
torch.set_num_threads(EMBED_THREADS)


def _ts() -> str:
    return datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]

DATASETS_SERVER = "https://datasets-server.huggingface.co"
MODEL_ID = "microsoft/wavlm-base-plus-sv"
TARGET_SR = 16000

_feature_extractor = None
_model = None
_init_lock = threading.Lock()


def _load_model():
    global _feature_extractor, _model
    with _init_lock:
        if _model is None:
            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
            _model = WavLMForXVector.from_pretrained(MODEL_ID).eval()
    return _feature_extractor, _model


def _to_array(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
    waveform = torch.tensor(audio_array, dtype=torch.float32)
    if waveform.ndim == 2:
        waveform = waveform.mean(0)
    if sr != TARGET_SR:
        waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)
    return waveform[: max_sec * TARGET_SR].numpy().astype(np.float32)


def _embed_one(array: np.ndarray) -> np.ndarray:
    """Embed a single clip. Thread-safe: eval()+no_grad(), GIL released in C++."""
    fe, mdl = _load_model()
    inputs = fe(array, sampling_rate=TARGET_SR, return_tensors="pt")
    with torch.no_grad():
        out = mdl(**inputs)
    return out.embeddings.squeeze().numpy()


def _parallel_embed(arrays: list[np.ndarray], log: list[str]) -> np.ndarray:
    """Embed all clips using N_WORKERS parallel threads (threads=2 each)."""
    t0 = time.time()
    log.append(f"[{_ts()}] --- Phase 2: embed {len(arrays)} clips "
               f"({N_WORKERS} workers × {EMBED_THREADS} threads) ---")
    with concurrent.futures.ThreadPoolExecutor(max_workers=N_WORKERS) as ex:
        futures = [ex.submit(_embed_one, arr) for arr in arrays]
        results = [f.result() for f in futures]
    ms = int((time.time() - t0) * 1000)
    log.append(f"[{_ts()}] embed done: {ms}ms total, {ms//len(arrays)}ms/clip avg")
    return np.stack(results)


# ── Audio fetching ────────────────────────────────────────────────────────────

def _parse_audio_urls(rows: list) -> list[str]:
    urls = []
    for row in rows:
        audio = row["row"].get("audio", {})
        if isinstance(audio, list):
            audio = audio[0] if audio else {}
        if isinstance(audio, dict) and "src" in audio:
            urls.append(audio["src"])
    return urls


def _try_endpoint(endpoint: str, params: dict, headers: dict) -> tuple[list[str], str]:
    """Single request attempt. Returns (urls, status_str)."""
    try:
        t0 = time.time()
        resp = requests.get(f"{DATASETS_SERVER}{endpoint}", params=params,
                            headers=headers, timeout=API_TIMEOUT)
        elapsed = int((time.time() - t0) * 1000)
        if not resp.ok:
            return [], f"{endpoint} HTTP {resp.status_code} ({elapsed}ms)"
        rows = resp.json().get("rows", [])
        if not rows:
            return [], f"{endpoint} empty ({elapsed}ms)"
        urls = _parse_audio_urls(rows)
        if urls:
            return urls, f"{endpoint} {elapsed}ms"
        return [], f"{endpoint} no src ({elapsed}ms)"
    except requests.Timeout:
        return [], f"{endpoint} timeout>{API_TIMEOUT}s"
    except Exception as exc:
        return [], f"{endpoint} {exc}"


def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> tuple[list[str], str]:
    """Fire /rows and /first-rows in parallel, return first successful result.
    Uses shutdown(wait=False) so the losing request doesn't block the caller."""
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    calls = [
        ("/rows",       {"dataset": repo_id, "config": "default", "split": "train", "offset": 0, "length": n}),
        ("/first-rows", {"dataset": repo_id, "config": "default", "split": "train"}),
    ]
    ex = concurrent.futures.ThreadPoolExecutor(max_workers=2)
    futs = {ex.submit(_try_endpoint, ep, params, headers): ep for ep, params in calls}
    errs = []
    try:
        for fut in concurrent.futures.as_completed(futs, timeout=API_TIMEOUT + 2):
            urls, status = fut.result()
            if urls:
                ex.shutdown(wait=False, cancel_futures=True)
                return urls[:n], status
            errs.append(status)
    except concurrent.futures.TimeoutError:
        errs.append("both endpoints timed out")
    finally:
        ex.shutdown(wait=False, cancel_futures=True)
    return [], " | ".join(errs)


def _download_audio(url: str, token: str | None, max_sec: int) -> tuple[np.ndarray, int, int]:
    """Download one audio file. Returns (array, size_kb, dl_ms)."""
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    t0 = time.time()
    resp = requests.get(url, headers=headers, timeout=30)
    resp.raise_for_status()
    dl_ms = int((time.time() - t0) * 1000)
    audio_array, sr = sf.read(io.BytesIO(resp.content))
    return _to_array(audio_array, sr, max_sec), len(resp.content) // 1024, dl_ms


def _fetch_repo_audio(
    repo: str, n_samples: int, audio_sec: int, token: str | None,
    log: list[str],
) -> tuple[str, list[np.ndarray]]:
    """Fetch + download audio for one repo. Appends timestamped entries to log."""
    short = repo.split("/")[-1]
    t0 = time.time()

    try:
        log.append(f"[{_ts()}] {short}: fetching URLs from datasets-server…")
        urls, api_status = _fetch_audio_urls(repo, n_samples, token)
        fetch_ms = int((time.time() - t0) * 1000)

        if urls:
            log.append(f"[{_ts()}] {short}: {api_status} → {len(urls)} URLs")
            with concurrent.futures.ThreadPoolExecutor(max_workers=len(urls)) as ex:
                futures = [ex.submit(_download_audio, u, token, audio_sec) for u in urls]
                results = [f.result() for f in futures]
            arrays = [r[0] for r in results]
            sizes  = [r[1] for r in results]
            dls    = [r[2] for r in results]
            total_ms = int((time.time() - t0) * 1000)
            log.append(
                f"[{_ts()}] {short}: ✓ {len(arrays)} clips "
                f"dl=[{', '.join(str(d)+'ms' for d in dls)}] "
                f"size=[{', '.join(str(s)+'KB' for s in sizes)}] "
                f"total={total_ms}ms"
            )
            return repo, arrays

        # Streaming fallback with timeout
        log.append(f"[{_ts()}] {short}: ⚠ datasets-server failed ({api_status}) → streaming fallback (timeout={STREAMING_TIMEOUT}s)")
        t1 = time.time()

        def _stream() -> list[np.ndarray]:
            from datasets import load_dataset, Audio as HFAudio
            ds = load_dataset(repo, split="train", streaming=True, token=token)
            ds = ds.cast_column("audio", HFAudio(decode=False))
            result = []
            for j, row in enumerate(ds):
                if j >= n_samples:
                    break
                raw = row["audio"]
                audio_bytes = raw.get("bytes") or open(raw["path"], "rb").read()
                a, sr = sf.read(io.BytesIO(audio_bytes))
                result.append(_to_array(a, sr, audio_sec))
                elapsed = int((time.time() - t1) * 1000)
                log.append(f"[{_ts()}] {short}: streaming clip {j+1}/{n_samples} ({len(audio_bytes)//1024}KB, {elapsed}ms so far)")
            return result

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            fut = ex.submit(_stream)
            try:
                arrays = fut.result(timeout=STREAMING_TIMEOUT)
                total_ms = int((time.time() - t0) * 1000)
                log.append(f"[{_ts()}] {short}: ✓ streaming done, {len(arrays)} clips, total={total_ms}ms")
                return repo, arrays
            except concurrent.futures.TimeoutError:
                total_ms = int((time.time() - t0) * 1000)
                log.append(f"[{_ts()}] {short}: ✗ streaming timeout after {STREAMING_TIMEOUT}s — skipping")
                return repo, []

    except Exception as exc:
        log.append(f"[{_ts()}] {short}: ✗ {exc} (total={int((time.time()-t0)*1000)}ms)")
        return repo, []


# ── Main ──────────────────────────────────────────────────────────────────────

def identify_speakers(
    repo_ids_text: str,
    samples_per_book: int,
    audio_sec: int,
    threshold: float,
    hf_token: str,
    progress=gr.Progress(),
):
    repos = [r.strip() for r in repo_ids_text.strip().splitlines() if r.strip()]
    if not repos:
        return pd.DataFrame(), "No repos provided.", "", "", None

    token = hf_token.strip() or os.environ.get("HF_TOKEN") or None

    log: list[str] = []
    t_total = time.time()

    log.append(f"[{_ts()}] === START: {len(repos)} repos, {samples_per_book} sample/book, {audio_sec}s/clip ===")
    log.append(f"[{_ts()}] CPU count: {N_CPUS}, workers: {N_WORKERS}, embed_threads: {EMBED_THREADS}")

    progress(0, desc="Loading model…")
    t_model = time.time()
    _load_model()
    log.append(f"[{_ts()}] model loaded in {int((time.time()-t_model)*1000)}ms")

    # ── Phase 1: download all audio in parallel (I/O bound) ──────────────────
    log.append(f"[{_ts()}] --- Phase 1: download ({N_CPUS*2} workers) ---")
    progress(0.05, desc=f"Downloading audio from {len(repos)} datasets…")
    repo_arrays: dict[str, list[np.ndarray]] = {}

    with concurrent.futures.ThreadPoolExecutor(max_workers=N_CPUS * 2) as ex:
        futures = {
            ex.submit(_fetch_repo_audio, repo, int(samples_per_book), int(audio_sec), token, log): repo
            for repo in repos
        }
        done = 0
        for future in concurrent.futures.as_completed(futures):
            repo, arrays = future.result()
            done += 1
            progress(0.05 + 0.55 * done / len(repos), desc=f"[{done}/{len(repos)}] {repo.split('/')[-1]}")
            if arrays:
                repo_arrays[repo] = arrays

    ok = len(repo_arrays)
    failed = len(repos) - ok
    log.append(f"[{_ts()}] Phase 1 done: {ok} ok, {failed} failed, phase_total={int((time.time()-t_total)*1000)}ms")

    if not repo_arrays:
        return pd.DataFrame(), "No audio downloaded.", "\n".join(log), "", None

    # ── Phase 2: parallel embed (N_WORKERS workers × EMBED_THREADS each) ───────
    progress(0.60, desc="Embedding clips…")
    all_arrays: list[np.ndarray] = []
    repo_slices: dict[str, tuple[int, int]] = {}

    for repo in repos:
        if repo in repo_arrays:
            start = len(all_arrays)
            all_arrays.extend(repo_arrays[repo])
            repo_slices[repo] = (start, len(all_arrays))

    all_embeddings = _parallel_embed(all_arrays, log)

    # ── Phase 3: average per repo, cluster ───────────────────────────────────
    log.append(f"[{_ts()}] --- Phase 3: cluster ({len(repo_slices)} repos) ---")
    progress(0.95, desc="Clustering…")
    embeddings: dict[str, np.ndarray] = {}
    for repo, (start, end) in repo_slices.items():
        embeddings[repo] = all_embeddings[start:end].mean(axis=0)

    repo_names = list(embeddings.keys())
    emb_matrix = np.stack([embeddings[r] for r in repo_names])
    emb_matrix = emb_matrix / np.linalg.norm(emb_matrix, axis=1, keepdims=True)
    sim_matrix = np.clip(emb_matrix @ emb_matrix.T, -1.0, 1.0)
    dist_matrix = 1.0 - sim_matrix
    np.fill_diagonal(dist_matrix, 0.0)

    n = len(repo_names)
    if n == 1:
        labels = [0]
    else:
        labels = AgglomerativeClustering(
            n_clusters=None,
            distance_threshold=1.0 - float(threshold),
            metric="precomputed",
            linkage="average",
        ).fit_predict(dist_matrix).tolist()

    rows = []
    for i, repo in enumerate(repo_names):
        cluster = labels[i]
        same_idx = [j for j, l in enumerate(labels) if l == cluster and j != i]
        intra_sim = float(np.mean([sim_matrix[i][j] for j in same_idx])) if same_idx else 1.0
        other_sorted = sorted([j for j in range(n) if j != i], key=lambda j: -sim_matrix[i][j])
        closest = (
            f"{repo_names[other_sorted[0]].split('/')[-1]} ({sim_matrix[i][other_sorted[0]]:.2f})"
            if other_sorted else "-"
        )
        rows.append({
            "dataset": repo.split("/")[-1],
            "speaker_id": f"speaker_{cluster + 1:02d}",
            "books_with_speaker": sum(1 for l in labels if l == cluster),
            "intra_sim": round(intra_sim, 3),
            "closest_match": closest,
        })

    df = pd.DataFrame(rows).sort_values(["speaker_id", "dataset"]).reset_index(drop=True)
    n_speakers = len(set(labels))
    total_ms = int((time.time() - t_total) * 1000)
    summary = f"✅ {len(repo_names)} books → {n_speakers} unique speakers  ({total_ms/1000:.1f}s total)"
    log.append(f"[{_ts()}] === DONE: {total_ms}ms total ===")

    # Plain-text copy-friendly output (space-separated, matches log format)
    text_lines = []
    for _, r in df.iterrows():
        text_lines.append(
            f"{r['dataset']}  {r['speaker_id']}  {r['books_with_speaker']}  {r['intra_sim']}  {r['closest_match']}"
        )
    plain_text = "\n".join(text_lines)

    # CSV file for download
    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False, encoding="utf-8")
    df.to_csv(tmp.name, index=False)
    tmp.close()

    return df, summary, "\n".join(log), plain_text, tmp.name


DESCRIPTION = """
# 🎙️ Speaker Identifier

Finds unique speakers across multiple HF audio datasets. Each dataset is assumed to have
**one speaker** (e.g. an audiobook). The app fetches audio directly via the datasets-server API
(no full parquet download), downloads in parallel, then embeds clips across 8 parallel workers.

**Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) —
language-agnostic speaker embeddings, works for any language.

---

## How to use

1. **Paste dataset repo IDs** (one per line, `owner/name` format) into the left box.
2. **Adjust parameters** if needed (defaults work well for audiobooks):
   - *Samples per book* — audio chunks to average per dataset. More = more robust, slower. 3 is usually enough.
   - *Audio length (sec)* — seconds of each chunk to use for embedding. 5 sec is sufficient for a clear voice.
   - *Same-speaker threshold* — cosine similarity cutoff. Raise if too many books merge into one speaker; lower if one person gets split across IDs.
   - *HF Token* — only needed for **private** repos.
3. Click **Identify Speakers**. Downloads run in parallel, then all clips are embedded in one batch — expect ~20–40 sec for 35 books.

## Output columns

| Column | Meaning |
|---|---|
| `dataset` | Repo name (short) |
| `speaker_id` | Cluster label — same ID = same voice |
| `books_with_speaker` | How many books share this speaker |
| `intra_sim` | Avg cosine similarity within cluster (1.0 = only one book; lower = cluster is less tight) |
| `closest_match` | Most similar other book and similarity score |

**Tip:** Sort by `speaker_id` to see all books by the same narrator grouped together.

The **Errors / Timing** box shows per-dataset timing and batch embed stats — useful for diagnosing slow datasets.
"""

with gr.Blocks(title="Speaker Identifier") as demo:
    gr.Markdown(DESCRIPTION)
    gr.Markdown("---")
    with gr.Row():
        with gr.Column(scale=2):
            repo_input = gr.Textbox(
                label="Dataset repo IDs (one per line)",
                placeholder="fosters/some_audiobook_output\nfosters/another_audiobook_output",
                lines=16,
            )
        with gr.Column(scale=1):
            samples = gr.Slider(1, 10, value=3, step=1, label="Samples per book")
            audio_sec = gr.Slider(
                2, 30, value=5, step=1,
                label="Audio length per sample (sec)",
                info="5 sec is usually enough; longer = more accurate but slower",
            )
            threshold = gr.Slider(
                0.60, 0.98, value=0.82, step=0.01,
                label="Same-speaker threshold",
                info="Higher = stricter matching → more clusters",
            )
            hf_token = gr.Textbox(
                label="HF Token (private repos only)",
                type="password",
                placeholder="hf_…",
            )
            run_btn = gr.Button("Identify Speakers", variant="primary", size="lg")

    summary_out = gr.Textbox(label="Summary", interactive=False)
    table_out = gr.Dataframe(
        label="Results — sorted by speaker_id",
        headers=["dataset", "speaker_id", "books_with_speaker", "intra_sim", "closest_match"],
        wrap=True,
    )
    with gr.Row():
        text_out = gr.Textbox(
            label="Plain text (copy-friendly)",
            interactive=False,
            lines=12,
            info="dataset  speaker_id  n_books  intra_sim  closest_match",
        )
        csv_out = gr.File(label="Download CSV", file_types=[".csv"])
    errors_out = gr.Textbox(label="Errors / Timing", interactive=False)

    run_btn.click(
        identify_speakers,
        inputs=[repo_input, samples, audio_sec, threshold, hf_token],
        outputs=[table_out, summary_out, errors_out, text_out, csv_out],
    )

demo.launch()