"""Chunk Classifier — HF Space.

Classifies audio chunks in one or more HF datasets using the AST AudioSet model.
Writes a ``<input>_classified`` dataset with id + audio + classification columns only.
"""

from __future__ import annotations

import datetime
import json
import os
import tempfile
import time
from pathlib import Path
from typing import Any, Generator

# Must be set before any HF import
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "120")

import gradio as gr
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import requests
import torch
from huggingface_hub import HfApi, hf_hub_download

import music_detector
from music_detector import _get_ast_runtime, judge_chunk_files_batched

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

try:
    N_CPUS = len(os.sched_getaffinity(0))   # container-aware (Linux cgroup)
except AttributeError:
    N_CPUS = os.cpu_count() or 2            # macOS / Windows fallback
torch.set_num_threads(N_CPUS)

BATCH_SIZE      = 32
DEFAULT_THRESHOLD = 0.25
DEFAULT_NAMESPACE = "fosters"
LOG_MAX_LINES   = 600
DATASETS_SERVER = "https://datasets-server.huggingface.co"

_CLASSIFIER_FIELDS = [
    pa.field("music_score",     pa.float32()),
    pa.field("has_music",       pa.bool_()),
    pa.field("contaminated",    pa.bool_()),
    pa.field("flags",           pa.string()),
    pa.field("top_labels",      pa.string()),
    pa.field("rms_db",          pa.float32()),
    pa.field("clipping_ratio",  pa.float32()),
    pa.field("max_silence_sec", pa.float32()),
]

_HF_FEATURES = {
    "id":              {"_type": "Value", "dtype": "string"},
    "audio":           {"_type": "Audio"},
    "music_score":     {"_type": "Value", "dtype": "float32"},
    "has_music":       {"_type": "Value", "dtype": "bool"},
    "contaminated":    {"_type": "Value", "dtype": "bool"},
    "flags":           {"_type": "Value", "dtype": "string"},
    "top_labels":      {"_type": "Value", "dtype": "string"},
    "rms_db":          {"_type": "Value", "dtype": "float32"},
    "clipping_ratio":  {"_type": "Value", "dtype": "float32"},
    "max_silence_sec": {"_type": "Value", "dtype": "float32"},
}

_AUDIO_PA_TYPE = pa.struct([pa.field("bytes", pa.binary()), pa.field("path", pa.string())])

_OUTPUT_SCHEMA = pa.schema(
    [
        pa.field("id", pa.string()),
        pa.field("audio", _AUDIO_PA_TYPE),
        *_CLASSIFIER_FIELDS,
    ],
    metadata={"huggingface": json.dumps({"info": {"features": _HF_FEATURES}})},
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


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


def _default_out_repo(in_repo: str, suffix: str) -> str:
    parts = in_repo.split("/")
    ns    = parts[0] if len(parts) > 1 else DEFAULT_NAMESPACE
    name  = parts[-1]
    suf   = suffix.strip()
    if not suf.startswith("_"):
        suf = "_" + suf
    return f"{ns}/{name}{suf}"


def _list_parquet_shards(repo_id: str, token: str | None) -> list[str]:
    api = HfApi(token=token)
    return sorted(
        f.rfilename
        for f in api.list_repo_tree(repo_id, repo_type="dataset", recursive=True)
        if hasattr(f, "rfilename")
        and f.rfilename.startswith("data/train-")
        and f.rfilename.endswith(".parquet")
    )


def _fetch_total_rows(repo_id: str, token: str | None) -> int | None:
    """Get total row count via datasets-server (no download required)."""
    try:
        headers = {"Authorization": f"Bearer {token}"} if token else {}
        r = requests.get(
            f"{DATASETS_SERVER}/info",
            params={"dataset": repo_id, "config": "default"},
            headers=headers,
            timeout=20,
        )
        if r.ok:
            splits = r.json().get("dataset_info", {}).get("default", {}).get("splits", {})
            return splits.get("train", {}).get("num_examples")
    except Exception:
        pass
    return None


def _rows_to_table(rows: list[dict[str, Any]]) -> pa.Table:
    return pa.table(
        {
            "id":              pa.array([r["id"]              for r in rows], type=pa.string()),
            "audio":           pa.array(
                [{"bytes": r["audio_bytes"], "path": r["audio_path"]} for r in rows],
                type=_AUDIO_PA_TYPE,
            ),
            "music_score":     pa.array([r["music_score"]     for r in rows], type=pa.float32()),
            "has_music":       pa.array([r["has_music"]       for r in rows], type=pa.bool_()),
            "contaminated":    pa.array([r["contaminated"]    for r in rows], type=pa.bool_()),
            "flags":           pa.array([r["flags"]           for r in rows], type=pa.string()),
            "top_labels":      pa.array([r["top_labels"]      for r in rows], type=pa.string()),
            "rms_db":          pa.array([r["rms_db"]          for r in rows], type=pa.float32()),
            "clipping_ratio":  pa.array([r["clipping_ratio"]  for r in rows], type=pa.float32()),
            "max_silence_sec": pa.array([r["max_silence_sec"] for r in rows], type=pa.float32()),
        },
        schema=_OUTPUT_SCHEMA,
    )


def _fmt_chunk(chunk_id: str, v: Any) -> str:  # v: ChunkVerdict
    if v.has_music:
        icon = "🔴"
    elif v.flags:
        icon = "⚠️"
    else:
        icon = "✅"
    labels = ", ".join(v.top_labels[:2])
    flags  = f"  [{', '.join(v.flags)}]" if v.flags else ""
    return f"  {icon} {chunk_id}  score={v.music_score:.3f}  {labels}{flags}"


# ---------------------------------------------------------------------------
# Per-repo processing generator
# ---------------------------------------------------------------------------

# Yields: (new_log_lines: list[str], status: str, chunks_done: int, stats: dict | None)
# chunks_done counts chunks processed so far in this repo (for caller progress math)

def _process_repo(
    repo_id: str,
    out_repo: str,
    music_threshold: float,
    private: bool,
    token: str | None,
) -> Generator[tuple[list[str], str, int, dict[str, Any] | None], None, None]:

    api   = HfApi(token=token)
    short = repo_id.split("/")[-1]

    # List shards
    yield [f"[{_ts()}] {short}: listing shards…"], "Listing shards…", 0, None
    try:
        shard_names = _list_parquet_shards(repo_id, token)
    except Exception as exc:
        err = str(exc)
        yield [f"[{_ts()}] {short}: ✗ {err}"], f"Error: {err}", 0, {"repo": repo_id, "error": err}
        return

    if not shard_names:
        msg = "no data/train-*.parquet shards found"
        yield [f"[{_ts()}] {short}: ✗ {msg}"], f"Error: {msg}", 0, {"repo": repo_id, "error": msg}
        return

    # Get total row count (for progress display, not critical)
    total_rows = _fetch_total_rows(repo_id, token)
    total_str  = str(total_rows) if total_rows else "?"

    yield [
        f"[{_ts()}] {short}: {len(shard_names)} shard(s), {total_str} rows total"
    ], f"Preparing… {len(shard_names)} shards, {total_str} chunks", 0, None

    api.create_repo(repo_id=out_repo, repo_type="dataset", exist_ok=True, private=private)

    stats: dict[str, Any] = {
        "repo": repo_id, "out_repo": out_repo,
        "total": 0, "contaminated": 0, "has_music": 0, "technical": 0,
    }

    with tempfile.TemporaryDirectory() as tmp:
        out_data_dir = Path(tmp) / "data"
        out_data_dir.mkdir()
        out_parts: list[Path] = []

        for shard_idx, shard_name in enumerate(shard_names):
            yield [
                f"[{_ts()}] {short}: shard {shard_idx+1}/{len(shard_names)} — downloading…"
            ], f"[{shard_idx+1}/{len(shard_names)}] downloading shard…", stats["total"], None

            t_dl = time.time()
            try:
                shard_path = hf_hub_download(
                    repo_id=repo_id, filename=shard_name,
                    repo_type="dataset", token=token,
                )
            except Exception as exc:
                yield [f"[{_ts()}] {short}: ✗ download failed: {exc}"], "Download error", stats["total"], None
                continue

            in_table  = pq.read_table(shard_path)
            n_rows    = len(in_table)
            audio_col = in_table.column("audio")
            id_col    = in_table.column("id").to_pylist()
            n_batches = (n_rows + BATCH_SIZE - 1) // BATCH_SIZE

            yield [
                f"[{_ts()}] {short}: shard {shard_idx+1}/{len(shard_names)} — "
                f"{n_rows} rows, downloaded in {time.time()-t_dl:.1f}s"
            ], (
                f"[{shard_idx+1}/{len(shard_names)}] {n_rows} rows — starting inference…"
            ), stats["total"], None

            shard_rows: list[dict[str, Any]] = []

            for batch_idx in range(n_batches):
                b_start = batch_idx * BATCH_SIZE
                b_end   = min(b_start + BATCH_SIZE, n_rows)
                t_batch = time.time()

                batch_ids: list[str] = id_col[b_start:b_end]
                batch_audio: list[dict] = [(audio_col[i].as_py() or {}) for i in range(b_start, b_end)]
                batch_bytes: list[bytes] = [a.get("bytes") or b"" for a in batch_audio]
                batch_paths: list[str]   = [a.get("path")  or "chunk.mp3" for a in batch_audio]

                with tempfile.TemporaryDirectory() as audio_tmp:
                    files: list[Path] = []
                    for k, (ab, ph) in enumerate(zip(batch_bytes, batch_paths)):
                        p = Path(audio_tmp) / f"c{k:04d}{Path(ph).suffix or '.mp3'}"
                        p.write_bytes(ab)
                        files.append(p)

                    verdicts = judge_chunk_files_batched(
                        files,
                        music_threshold=music_threshold,
                        batch_size=len(files),
                        device="cpu",
                    )

                for cid, ab, ph, v in zip(batch_ids, batch_bytes, batch_paths, verdicts):
                    shard_rows.append({
                        "id": cid, "audio_bytes": ab, "audio_path": ph,
                        "music_score":     v.music_score,
                        "has_music":       v.has_music,
                        "contaminated":    v.contaminated,
                        "flags":           ", ".join(v.flags),
                        "top_labels":      ", ".join(v.top_labels[:3]),
                        "rms_db":          v.rms_db,
                        "clipping_ratio":  v.clipping_ratio,
                        "max_silence_sec": v.max_silence_sec,
                    })

                n_cont  = sum(v.contaminated for v in verdicts)
                n_music = sum(v.has_music    for v in verdicts)
                n_tech  = sum(bool(v.flags) and not v.has_music for v in verdicts)
                stats["total"]        += b_end - b_start
                stats["contaminated"] += n_cont
                stats["has_music"]    += n_music
                stats["technical"]    += n_tech

                elapsed = time.time() - t_batch
                done    = stats["total"]
                cont_pct = 100 * stats["contaminated"] / max(done, 1)

                tm = music_detector._LAST_TIMINGS
                phase_str = (
                    f" · decode {tm.get('decode', 0):.1f}s "
                    f"/ fbank {tm.get('fbank', 0):.1f}s "
                    f"/ infer {tm.get('infer', 0):.1f}s"
                    if tm else ""
                )

                # Build log lines: one per chunk + batch summary
                chunk_lines = [_fmt_chunk(cid, v) for cid, v in zip(batch_ids, verdicts)]
                summary_line = (
                    f"[{_ts()}] batch {batch_idx+1}/{n_batches} "
                    f"· shard {shard_idx+1}/{len(shard_names)} "
                    f"· {elapsed:.1f}s"
                    f"{phase_str} "
                    f"· total {done}/{total_str} chunks "
                    f"· contaminated {stats['contaminated']} ({cont_pct:.1f}%)"
                )

                status = (
                    f"[{shard_idx+1}/{len(shard_names)}] "
                    f"batch {batch_idx+1}/{n_batches} "
                    f"· {done}/{total_str} chunks "
                    f"· {cont_pct:.1f}% contaminated"
                )

                yield [""] + chunk_lines + [summary_line], status, done, None

            # Write shard
            out_table = _rows_to_table(shard_rows)
            out_part  = out_data_dir / f"part_{shard_idx:05d}.parquet"
            pq.write_table(out_table, out_part, row_group_size=1, compression="snappy")
            out_parts.append(out_part)
            yield [
                f"[{_ts()}] {short}: shard {shard_idx+1} written "
                f"({len(shard_rows)} rows, {out_part.stat().st_size // 1024} KB)"
            ], f"Shard {shard_idx+1} written, uploading…", stats["total"], None

        # Rename + upload
        n_shards = len(out_parts)
        for i, p in enumerate(out_parts):
            p.rename(out_data_dir / f"train-{i:05d}-of-{n_shards:05d}.parquet")

        yield [f"[{_ts()}] {short}: uploading {n_shards} shard(s) → {out_repo}…"], "Uploading…", stats["total"], None
        t_up = time.time()
        api.upload_folder(
            folder_path=str(out_data_dir),
            repo_id=out_repo,
            repo_type="dataset",
            path_in_repo="data/",
            delete_patterns=["data/*.parquet"],
            commit_message=(
                f"Classify {stats['total']} chunks "
                f"(threshold={music_threshold:.2f}, "
                f"contaminated={stats['contaminated']})"
            ),
        )

        url = f"https://huggingface.co/datasets/{out_repo}"
        yield [
            f"[{_ts()}] {short}: ✓ uploaded in {time.time()-t_up:.1f}s → {url}"
        ], f"Done ✓  →  {out_repo}", stats["total"], stats


# ---------------------------------------------------------------------------
# Gradio generator (no gr.Progress — avoids overlay)
# ---------------------------------------------------------------------------

# Outputs: log_text, status_text, summary_df, links_md
def classify_repos(
    repos_text: str,
    out_suffix: str,
    music_threshold: float,
    private: bool,
    hf_token_input: str,
) -> Generator[tuple[str, str, pd.DataFrame, str], None, None]:

    token = hf_token_input.strip() or os.environ.get("HF_TOKEN") or None
    repos = [
        r.strip()
        for line in repos_text.replace(",", "\n").splitlines()
        for r in [line.strip()]
        if r and "/" in r
    ]
    if not repos:
        yield "No valid repo IDs (expected owner/name format).", "", pd.DataFrame(), ""
        return

    log:  list[str]            = []
    rows: list[dict[str, Any]] = []
    links: list[str]           = []

    def _snapshot(status: str) -> tuple[str, str, pd.DataFrame, str]:
        log_text = "\n".join(log[-LOG_MAX_LINES:])
        df       = pd.DataFrame(rows) if rows else pd.DataFrame()
        links_md = "## Output datasets\n" + "\n".join(links) if links else ""
        return log_text, status, df, links_md

    log.append(f"[{_ts()}] {len(repos)} repo(s) · threshold={music_threshold:.2f} · CPUs={N_CPUS}")
    log.append(f"[{_ts()}] HF_XET_HIGH_PERFORMANCE={os.environ.get('HF_XET_HIGH_PERFORMANCE','off')}")
    log.append(f"[{_ts()}] Loading AST model…")
    yield _snapshot("Loading AST model…")

    try:
        _get_ast_runtime("cpu")
        log.append(f"[{_ts()}] Model ready ✓")
    except Exception as exc:
        log.append(f"[{_ts()}] ✗ model load failed: {exc}")
        yield _snapshot(f"Error: {exc}")
        return

    yield _snapshot("Model ready")
    t_total = time.time()

    for repo_idx, repo_id in enumerate(repos):
        out_repo = _default_out_repo(repo_id, out_suffix)
        log.append(f"\n[{_ts()}] {'─'*56}")
        log.append(f"[{_ts()}] [{repo_idx+1}/{len(repos)}] {repo_id}  →  {out_repo}")

        t_repo      = time.time()
        final_stats: dict[str, Any] | None = None

        for new_lines, status, chunks_done, maybe_stats in _process_repo(
            repo_id=repo_id,
            out_repo=out_repo,
            music_threshold=music_threshold,
            private=private,
            token=token,
        ):
            log.extend(new_lines)
            if maybe_stats is not None:
                final_stats = maybe_stats
            yield _snapshot(f"[{repo_idx+1}/{len(repos)}] {status}")

        # Summary row
        if final_stats and "error" not in final_stats:
            n     = final_stats["total"]
            cont  = final_stats["contaminated"]
            rows.append({
                "repo":      repo_id.split("/")[-1],
                "chunks":    n,
                "contam_%":  f"{100*cont/max(n,1):.1f}",
                "music":     final_stats["has_music"],
                "technical": final_stats["technical"],
                "time_s":    f"{time.time()-t_repo:.0f}",
                "output":    out_repo,
            })
            links.append(f"- [{out_repo}](https://huggingface.co/datasets/{out_repo})")
            log.append(
                f"[{_ts()}] DONE  total={n}  contaminated={cont} ({100*cont/max(n,1):.1f}%)  "
                f"music={final_stats['has_music']}  technical={final_stats['technical']}  "
                f"elapsed={time.time()-t_repo:.0f}s"
            )
        elif final_stats:
            rows.append({
                "repo": repo_id.split("/")[-1], "chunks": 0,
                "contam_%": "—", "music": 0, "technical": 0,
                "time_s": "—", "output": f"ERROR: {final_stats.get('error')}",
            })

        yield _snapshot(f"[{repo_idx+1}/{len(repos)}] done")

    log.append(f"\n[{_ts()}] {'='*56}")
    log.append(f"[{_ts()}] All done in {time.time()-t_total:.0f}s")
    yield _snapshot("All done ✓")


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

DESCRIPTION = """
# 🎵 Chunk Classifier

Scores audio chunks for **music contamination** and **technical defects** using
[MIT/ast-finetuned-audioset-10-10-0.4593](https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593)
(527 AudioSet classes, multi-label sigmoid).

**Input** — chunk dataset repo IDs (one per line).
**Output** — `<input>_classified` with `id · audio · music_score · has_music · contaminated · flags · top_labels · rms_db · clipping_ratio · max_silence_sec`.
Sort by `music_score ↓` in the HF viewer and press play to spot-check.
"""

with gr.Blocks(title="Chunk Classifier") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=2):
            repos_input = gr.Textbox(
                label="Dataset repo IDs (one per line)",
                placeholder="fosters/my-book-chunks\nfosters/another-book-chunks",
                lines=5,
            )
        with gr.Column(scale=1):
            out_suffix = gr.Textbox(
                label="Output suffix", value="_classified",
                info="Appended to each input repo name",
            )
            threshold = gr.Slider(
                0.10, 0.60, value=DEFAULT_THRESHOLD, step=0.05,
                label="Music threshold",
                info="Lower = more sensitive (recall-biased). Default 0.25.",
            )
            private_toggle = gr.Checkbox(label="Private output", value=True)
            token_input    = gr.Textbox(
                label="HF Token (or set HF_TOKEN secret)",
                type="password", placeholder="hf_…",
            )
            run_btn = gr.Button("▶ Classify", variant="primary", size="lg")

    # Status line — replaces gr.Progress() which caused overlay issues
    status_out = gr.Textbox(
        label="Status",
        interactive=False,
        lines=1,
        max_lines=1,
    )

    with gr.Row():
        with gr.Column(scale=3):
            log_out = gr.Textbox(
                label="Processing log  (updates every batch ~30s)",
                interactive=False,
                lines=28,
                max_lines=60,
            )
        with gr.Column(scale=2):
            summary_out = gr.Dataframe(
                label="Summary",
                headers=["repo", "chunks", "contam_%", "music", "technical", "time_s", "output"],
                wrap=True,
            )
            links_out = gr.Markdown()

    run_btn.click(
        classify_repos,
        inputs=[repos_input, out_suffix, threshold, private_toggle, token_input],
        outputs=[log_out, status_out, summary_out, links_out],
    )

demo.queue()
demo.launch(server_name="0.0.0.0")