import io
import os
import time
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 optimum.onnxruntime import ORTModelForAudioXVector
from transformers import Wav2Vec2FeatureExtractor

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

N_CPUS = os.cpu_count() or 2

DATASETS_SERVER = "https://datasets-server.huggingface.co"
ONNX_MODEL_ID = "fosters/wavlm-base-plus-sv-onnx"
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(ONNX_MODEL_ID)
            # ONNX Runtime: thread-safe, no GIL concerns
            _model = ORTModelForAudioXVector.from_pretrained(ONNX_MODEL_ID)
    return _feature_extractor, _model


def _embed(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
    fe, mdl = _load_model()
    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)
    waveform = waveform[: max_sec * TARGET_SR]
    inputs = fe(waveform.numpy(), sampling_rate=TARGET_SR, return_tensors="pt")
    out = mdl(**inputs)
    return out.embeddings.squeeze().numpy()


def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> tuple[list[str], str]:
    """Get direct audio URLs via datasets-server. Returns (urls, debug_info)."""
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    resp = requests.get(
        f"{DATASETS_SERVER}/rows",
        params={"dataset": repo_id, "config": "default", "split": "train", "offset": 0, "length": n},
        headers=headers,
        timeout=30,
    )
    if not resp.ok:
        return [], f"datasets-server {resp.status_code}"

    rows = resp.json().get("rows", [])
    if not rows:
        return [], "datasets-server: empty rows"

    # Inspect actual audio field structure for debugging
    sample_audio = rows[0]["row"].get("audio", {})
    audio_keys = list(sample_audio.keys()) if isinstance(sample_audio, dict) else type(sample_audio).__name__

    urls = []
    for row in rows:
        audio = row["row"].get("audio", {})
        # audio can be a dict {"src": ...} or a list [{"src": ...}, ...]
        if isinstance(audio, list):
            audio = audio[0] if audio else {}
        if isinstance(audio, dict) and "src" in audio:
            urls.append(audio["src"])

    if not urls:
        return [], f"datasets-server: no src in audio field (keys={audio_keys})"
    return urls, f"datasets-server ok (audio keys={audio_keys})"


def _download_one(url: str, token: str | None, max_sec: int) -> tuple[np.ndarray, dict]:
    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 = (time.time() - t0) * 1000

    t1 = time.time()
    audio_array, sr = sf.read(io.BytesIO(resp.content))
    decode_ms = (time.time() - t1) * 1000

    t2 = time.time()
    emb = _embed(audio_array, sr, max_sec)
    embed_ms = (time.time() - t2) * 1000

    return emb, {
        "size_kb": len(resp.content) // 1024,
        "dl_ms": int(dl_ms),
        "decode_ms": int(decode_ms),
        "embed_ms": int(embed_ms),
    }


def _streaming_fallback(repo: str, n_samples: int, audio_sec: int, token: str | None) -> list[np.ndarray]:
    """Fallback: use datasets streaming when datasets-server is unavailable."""
    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))
    embs = []
    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()
        audio_array, sr = sf.read(io.BytesIO(audio_bytes))
        embs.append(_embed(audio_array, sr, audio_sec))
    return embs


def _process_repo(
    repo: str, n_samples: int, audio_sec: int, token: str | None
) -> tuple[str, np.ndarray | None, str]:
    t0 = time.time()
    try:
        urls, api_debug = _fetch_audio_urls(repo, n_samples, token)
        fetch_ms = int((time.time() - t0) * 1000)

        if urls:
            # Fast path: parallel downloads via datasets-server URLs
            with concurrent.futures.ThreadPoolExecutor(max_workers=len(urls)) as ex:
                futures = [ex.submit(_download_one, u, token, audio_sec) for u in urls]
                results = [f.result() for f in futures]
            embs = [r[0] for r in results]
            s = results[0][1]
            total_ms = int((time.time() - t0) * 1000)
            log = (
                f"{repo.split('/')[-1]}: "
                f"api={fetch_ms}ms dl={s['dl_ms']}ms "
                f"decode={s['decode_ms']}ms embed={s['embed_ms']}ms "
                f"total={total_ms}ms ({s['size_kb']}KB/file) [{api_debug}]"
            )
        else:
            # Slow fallback: streaming (downloads parquet shard)
            t1 = time.time()
            embs = _streaming_fallback(repo, n_samples, audio_sec, token)
            total_ms = int((time.time() - t0) * 1000)
            log = (
                f"{repo.split('/')[-1]}: streaming fallback "
                f"total={total_ms}ms [reason: {api_debug}]"
            )

        if not embs:
            return repo, None, f"{repo}: no audio samples extracted"

        return repo, np.mean(embs, axis=0), log

    except Exception as exc:
        return repo, None, f"{repo}: ERROR {exc}"


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.", ""

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

    progress(0, desc="Loading model…")
    _load_model()
    progress(0.02, desc=f"Processing {len(repos)} datasets in parallel…")

    embeddings: dict[str, np.ndarray] = {}
    logs: list[str] = []
    errors: list[str] = []
    done = 0

    # Process repos in parallel — capped at N_CPUS since embed is the bottleneck
    with concurrent.futures.ThreadPoolExecutor(max_workers=N_CPUS) as ex:
        future_to_repo = {
            ex.submit(_process_repo, repo, int(samples_per_book), int(audio_sec), token): repo
            for repo in repos
        }
        for future in concurrent.futures.as_completed(future_to_repo):
            repo, emb, log = future.result()
            done += 1
            progress(done / len(repos), desc=f"[{done}/{len(repos)}] {repo.split('/')[-1]}")
            if emb is not None:
                embeddings[repo] = emb
                logs.append(log)
            else:
                errors.append(log)

    if not embeddings:
        return pd.DataFrame(), "No embeddings extracted.", "\n".join(errors)

    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))
    summary = f"✅ {len(repo_names)} books → {n_speakers} unique speakers"

    debug = "\n".join(logs)
    if errors:
        debug += "\n\nERRORS:\n" + "\n".join(errors)
    return df, summary, debug


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), processes datasets in parallel, and clusters by voice similarity.

**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**. Datasets are processed in parallel — expect ~30–60 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 breakdown (API fetch / download / decode / embed) — 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,
    )
    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],
    )

demo.launch()