#!/usr/bin/env python3
"""
prepare_dataset.py
Build a clean, single-speaker voice dataset from a folder of video files.

Pipeline (per video):
  1. ffmpeg      -> extract the audio track
  2. Demucs      -> strip music / background noise, keep vocals only
  3. Whisper     -> transcribe with word timestamps (faster-whisper), cut into sentences
  4. Speaker ID  -> check every sentence against the reference voice (ECAPA embeddings)
  5. Chunking    -> merge consecutive target-speaker sentences into 3-15 s clips,
                    then re-check each finished clip
  6. Quality     -> reject clipped, too-quiet, or near-empty clips
  7. Export      -> wavs/ + metadata.csv (LJSpeech style) + report.csv (every candidate, kept or not)

Resumable: videos already processed are skipped on re-run.
"""

import os
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")  # quiet a harmless Windows warning
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS", "1")

import argparse
import csv
import hashlib
import json
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torchaudio
from tqdm import tqdm

VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v", ".mts", ".wav", ".mp3", ".m4a"}
EMBED_SR = 16000  # speaker model + whisper both work at 16 kHz


# ---------------------------------------------------------------- utilities

def run(cmd):
    """Run a command; raise with its stderr if it fails."""
    result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"Command failed: {' '.join(map(str, cmd))}\n{result.stderr[-2000:]}")


def check_ffmpeg():
    if shutil.which("ffmpeg") is None:
        raise ValueError("ffmpeg not found. Install it first (see README).")


def decode_audio(src, dst, sr, channels):
    run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(src), "-vn",
         "-ac", str(channels), "-ar", str(sr), "-c:a", "pcm_s16le", str(dst)])


def load_mono(path):
    data, sr = sf.read(str(path), dtype="float32", always_2d=True)
    return data.mean(axis=1), sr


def resample(x, sr_in, sr_out):
    if sr_in == sr_out:
        return x
    return torchaudio.functional.resample(torch.from_numpy(x), sr_in, sr_out).numpy()


def video_id(path):
    slug = re.sub(r"[^A-Za-z0-9]+", "_", path.stem).strip("_")[:40] or "video"
    digest = hashlib.md5(str(path.resolve()).encode()).hexdigest()[:6]
    return f"{slug}_{digest}"


def dbfs(x):
    rms = np.sqrt(np.mean(np.square(x)) + 1e-12)
    return 20 * np.log10(rms + 1e-12)


# ---------------------------------------------------------------- vocal separation

def separate_vocals(wav_path, workdir, device):
    run([sys.executable, "-m", "demucs", "--two-stems", "vocals", "-n", "htdemucs",
         "-d", device, "-o", str(workdir), str(wav_path)])
    vocals = workdir / "htdemucs" / wav_path.stem / "vocals.wav"
    if not vocals.exists():
        raise RuntimeError(f"Demucs did not produce {vocals}")
    return vocals


# ---------------------------------------------------------------- speaker filter

class SpeakerFilter:
    """Scores how closely a clip matches the reference speaker (cosine similarity, -1..1)."""

    def __init__(self, reference_paths, device, tmpdir):
        from importlib import import_module

        try:
            EncoderClassifier = import_module("speechbrain.inference.speaker").EncoderClassifier
        except ImportError:  # speechbrain < 1.0
            EncoderClassifier = import_module("speechbrain.pretrained").EncoderClassifier
        self.device = device
        kwargs = dict(
            source="speechbrain/spkrec-ecapa-voxceleb",
            savedir=str(Path.home() / ".cache" / "voiceprep" / "ecapa"),
            run_opts={"device": device},
        )
        # Windows blocks symlinks for normal users, so copy the model files instead.
        try:
            from speechbrain.utils.fetching import LocalStrategy
            kwargs["local_strategy"] = LocalStrategy.COPY
        except ImportError:  # older speechbrain without LocalStrategy
            pass
        self.model = EncoderClassifier.from_hparams(**kwargs)
        embs = []
        for i, ref in enumerate(reference_paths):
            tmp = Path(tmpdir) / f"ref_{i}.wav"
            decode_audio(ref, tmp, EMBED_SR, 1)
            x, _ = load_mono(tmp)
            embs.append(self.embed(x))
        ref = torch.stack(embs).mean(dim=0)
        self.ref = ref / ref.norm()

    @torch.no_grad()
    def embed(self, x16):
        wav = torch.from_numpy(x16).float().unsqueeze(0).to(self.device)
        emb = self.model.encode_batch(wav).squeeze().cpu()
        return emb / emb.norm()

    def score(self, x16):
        return float(torch.dot(self.embed(x16), self.ref))


# ---------------------------------------------------------------- chunking

class Unit:
    """A sentence-sized piece of speech with timing, text, and a speaker verdict."""

    def __init__(self, start, end, text):
        self.start, self.end, self.text = start, end, text
        self.ok = True
        self.sim = None


def split_sentences(segments, pause=0.5):
    """Re-cut Whisper segments into sentence-sized units using word timestamps.
    A new unit starts after sentence-ending punctuation or a pause, so a quick
    back-and-forth between two people ends up in separate units."""
    units = []
    for seg in segments:
        words = [w for w in (seg.words or []) if w.word.strip()]
        if not words:
            if seg.text.strip():
                units.append(Unit(seg.start, seg.end, seg.text.strip()))
            continue
        cur = None
        for w in words:
            if cur and (w.start - cur.end > pause):
                units.append(cur)
                cur = None
            if cur is None:
                cur = Unit(w.start, w.end, w.word.strip())
            else:
                cur.end = w.end
                cur.text += w.word if w.word.startswith(" ") else " " + w.word
            if w.word.strip()[-1:] in ".?!":
                units.append(cur)
                cur = None
        if cur:
            units.append(cur)
    for u in units:
        u.text = re.sub(r"\s+", " ", u.text).strip()
    return units


def build_chunks(segments, min_len, max_len, max_gap):
    """Merge consecutive speech units into clips between min_len and max_len seconds.
    Units are never split, so text stays aligned with audio. A unit marked as
    another speaker (ok=False) always ends the current clip, so voices never mix."""
    chunks, cur = [], None
    for seg in segments:
        text = seg.text.strip()
        if not text:
            continue
        if not getattr(seg, "ok", True):  # someone else talking: never merge across it
            if cur:
                chunks.append(cur)
            cur = None
            continue
        if seg.end - seg.start > max_len:  # a single overlong segment can't be aligned safely
            if cur:
                chunks.append(cur)
            cur = None
            continue
        if cur and (seg.start - cur["end"] <= max_gap) and (seg.end - cur["start"] <= max_len):
            cur["end"] = seg.end
            cur["text"] += " " + text
        else:
            if cur:
                chunks.append(cur)
            cur = {"start": seg.start, "end": seg.end, "text": text}
    if cur:
        chunks.append(cur)
    return [c for c in chunks if c["end"] - c["start"] >= min_len]


# ---------------------------------------------------------------- main per-video work

def process_video(path, args, whisper, spk, writers, tmp_root, report=None):
    vid = video_id(path)
    work = Path(tempfile.mkdtemp(prefix=vid + "_", dir=tmp_root))
    say = report or (lambda msg: tqdm.write(f"  [{path.name}] {msg}"))
    try:
        say("extracting audio...")
        full = work / f"{vid}.wav"
        decode_audio(path, full, 44100, 2)
        say("removing music and background noise (slowest step)...")
        vocals_path = separate_vocals(full, work, args.device)

        vocals, sr = load_mono(vocals_path)
        vocals16 = resample(vocals, sr, EMBED_SR)
        out_audio = resample(vocals, sr, args.sample_rate)
        total = len(vocals16) / EMBED_SR

        say("transcribing...")
        segments, _ = whisper.transcribe(
            vocals16, language=args.language, vad_filter=True, beam_size=5,
            condition_on_previous_text=False, word_timestamps=True,
        )
        units = split_sentences(list(segments))

        # Check the speaker on every sentence BEFORE merging, so mixed clips can't form.
        say(f"checking who is speaking in {len(units)} sentences...")
        other = 0
        for j, u in enumerate(units):
            dur = u.end - u.start
            if dur < args.min_unit:  # too short to judge reliably ("Yep.") -> treat as a break
                u.ok = False
                continue
            s, e = u.start, u.end
            if dur < 1.5:  # give the voice model at least 1.5 s of context
                mid = (s + e) / 2
                s, e = max(0.0, mid - 0.75), min(total, mid + 0.75)
            u.sim = spk.score(vocals16[int(s * EMBED_SR):int(e * EMBED_SR)])
            u.ok = u.sim >= args.sentence_threshold
            if not u.ok:
                other += 1
                writers["report"].writerow([f"{vid}_s{j:04d}", path.name, f"{u.start:.2f}",
                                            f"{u.end:.2f}", f"{dur:.2f}", f"{u.sim:.3f}", "",
                                            "rejected", "other_speaker_sentence", u.text])
        say(f"{other} sentences from other speakers removed, building clips...")
        chunks = build_chunks(units, args.min_len, args.max_len, args.max_gap)

        kept = 0
        pad = args.pad
        for i, c in enumerate(chunks):
            start = max(0.0, c["start"] - pad)
            end = min(len(vocals16) / EMBED_SR, c["end"] + pad)
            text = re.sub(r"\s+", " ", c["text"]).strip()
            clip_id = f"{vid}_{i:04d}"

            seg16 = vocals16[int(start * EMBED_SR):int(end * EMBED_SR)]
            seg_out = out_audio[int(start * args.sample_rate):int(end * args.sample_rate)]

            sim = spk.score(seg16)
            level = dbfs(seg_out)
            clipped = float(np.mean(np.abs(seg_out) >= 0.999))

            reason = ""
            if len(text.split()) < args.min_words:
                reason = "too_few_words"
            elif sim < args.threshold:
                reason = "other_speaker"
            elif level < args.min_dbfs:
                reason = "too_quiet"
            elif clipped > 0.001:
                reason = "clipped"

            status = "rejected" if reason else "kept"
            if not reason:
                peak = np.max(np.abs(seg_out)) + 1e-9
                seg_out = seg_out * (0.89 / peak)  # peak-normalize to about -1 dBFS
                sf.write(str(args.out / "wavs" / f"{clip_id}.wav"), seg_out, args.sample_rate,
                         subtype="PCM_16")
                writers["meta"].writerow([clip_id, text, text])
                kept += 1

            writers["report"].writerow([clip_id, path.name, f"{start:.2f}", f"{end:.2f}",
                                        f"{end - start:.2f}", f"{sim:.3f}", f"{level:.1f}",
                                        status, reason, text])
        return kept, len(chunks)
    finally:
        if args.keep_temp:
            print(f"  temp files kept in {work}")
        else:
            shutil.rmtree(work, ignore_errors=True)


def open_csv(path, header, delimiter=","):
    new = not path.exists()
    f = open(path, "a", newline="", encoding="utf-8")
    w = csv.writer(f, delimiter=delimiter)
    if new and header:
        w.writerow(header)
    return f, w


def run_dataset(args, progress=None):
    report_progress = progress or (lambda message, value: None)
    check_ffmpeg()
    videos = sorted(f for f in args.input.rglob("*") if f.suffix.lower() in VIDEO_EXTS)
    if not videos:
        raise ValueError(f"No video/audio files found in {args.input}")

    (args.out / "wavs").mkdir(parents=True, exist_ok=True)
    progress_file = args.out / "progress.json"
    done = set(json.loads(progress_file.read_text())) if progress_file.exists() else set()
    todo = [v for v in videos if str(v.resolve()) not in done]
    print(f"{len(videos)} files found, {len(todo)} left to process, device={args.device}")
    if args.device == "cpu":
        print("Warning: no GPU detected. This will be very slow on a large archive.")
    if not todo:
        result = {
            "videos_found": len(videos), "videos_processed": 0,
            "videos_skipped": len(videos), "clips_kept": 0,
            "failures": [], "dataset": str(args.out.resolve()),
        }
        report_progress("All uploaded videos have already been processed.", 100)
        return result

    tmp_root = Path(tempfile.mkdtemp(prefix="voiceprep_"))
    meta_f = rep_f = None
    try:
        report_progress("Loading transcription and speaker-recognition models…", 0)
        from faster_whisper import WhisperModel
        whisper = WhisperModel(args.whisper_model, device=args.device,
                               compute_type="float16" if args.device == "cuda" else "int8")
        spk = SpeakerFilter(args.reference, args.device, tmp_root)
        meta_f, meta_w = open_csv(args.out / "metadata.csv", None, delimiter="|")
        rep_f, rep_w = open_csv(
            args.out / "report.csv",
            ["clip_id", "source", "start", "end", "duration", "similarity",
             "dbfs", "status", "reason", "text"],
        )
        writers = {"meta": meta_w, "report": rep_w}
        total_kept = 0
        failures = []
        for index, v in enumerate(todo, 1):
            base_progress = (index - 1) / max(len(todo), 1) * 100

            def report_video(message, current=index, current_path=v, base=base_progress):
                report_progress(
                    f"{current_path.name}: {message}",
                    min(99, round(base + 5 / max(len(todo), 1))),
                )

            report_progress(f"Processing video {index} of {len(todo)}: {v.name}", round(base_progress))
            try:
                kept, cands = process_video(v, args, whisper, spk, writers, tmp_root,
                                            report=report_video)
                total_kept += kept
                tqdm.write(f"{v.name}: kept {kept}/{cands} clips")
            except Exception as e:  # one bad file shouldn't stop the whole run
                tqdm.write(f"{v.name}: FAILED - {e}")
                failures.append({"video": v.name, "error": str(e)})
                report_progress(f"{v.name} failed: {e}", round(index / max(len(todo), 1) * 100))
                continue
            meta_f.flush()
            rep_f.flush()
            done.add(str(v.resolve()))
            progress_file.write_text(json.dumps(sorted(done)))
        result = {
            "videos_found": len(videos),
            "videos_processed": len(todo) - len(failures),
            "videos_skipped": len(videos) - len(todo),
            "clips_kept": total_kept,
            "failures": failures,
            "dataset": str(args.out.resolve()),
        }
    finally:
        if meta_f is not None:
            meta_f.close()
        if rep_f is not None:
            rep_f.close()
        shutil.rmtree(tmp_root, ignore_errors=True)

    report_progress(
        f"Dataset ready: {total_kept} new clips from {result['videos_processed']} videos.",
        100,
    )
    print(f"\nDone. {total_kept} new clips this run. Dataset: {args.out.resolve()}")
    print("Review report.csv to tune --threshold before training.")
    return result


def main():
    p = argparse.ArgumentParser(description="Build a single-speaker voice dataset from videos.")
    p.add_argument("input", type=Path, help="Folder of videos (searched recursively)")
    p.add_argument("--reference", type=Path, nargs="+", required=True,
                   help="1-5 short clips (10-30 s each) of ONLY the target voice")
    p.add_argument("--out", type=Path, default=Path("dataset"))
    p.add_argument("--threshold", type=float, default=0.60,
                   help="Speaker similarity cutoff for finished clips")
    p.add_argument("--sentence-threshold", type=float, default=0.55,
                   help="Speaker cutoff for each sentence before merging (raise if other "
                        "voices slip into clips)")
    p.add_argument("--min-unit", type=float, default=0.6,
                   help="Sentences shorter than this (s) are too short to judge and split clips")
    p.add_argument("--sample-rate", type=int, default=44100, help="Output clip sample rate")
    p.add_argument("--min-len", type=float, default=3.0)
    p.add_argument("--max-len", type=float, default=15.0)
    p.add_argument("--max-gap", type=float, default=0.6, help="Max pause (s) when merging segments")
    p.add_argument("--pad", type=float, default=0.10, help="Silence padding (s) around each clip")
    p.add_argument("--min-words", type=int, default=2)
    p.add_argument("--min-dbfs", type=float, default=-40.0)
    p.add_argument("--whisper-model", default="large-v3",
                   help="large-v3 (best), medium, small (faster on weak GPUs)")
    p.add_argument("--language", default="en")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--keep-temp", action="store_true", help="Keep extracted/separated audio")
    args = p.parse_args()
    try:
        run_dataset(args)
    except ValueError as e:
        sys.exit(str(e))


if __name__ == "__main__":
    main()