#!/usr/bin/env python3
"""VoiceGrab — cut a time range from a video file and export a clean MP3 voice sample.

- Pick an input video (mp4, mkv, webm, avi, mov, ...)
- Set start / stop timestamps
- Optional voice processing (one of, mutually exclusive):
    * fast noise reduction (CPU DSP via noisereduce — no downloads)
    * AI voice separation (Demucs htdemucs_ft — GPU if available, optional install)
    * no processing (loudness-normalize only)
- Export MP3 (320 kbps, loudness-normalized to -16 LUFS — good for AI voice training)
- Waveform preview with the selected range highlighted
"""
from __future__ import annotations

import math
import os
import platform
import shutil
import subprocess
import sys
import tempfile
import threading
import wave

APP_NAME = "VoiceGrab"
APP_VERSION = "1.1.0"

# ---------------------------------------------------------------------------
# Resource / executable helpers
# ---------------------------------------------------------------------------


def app_dir() -> str:
    """Directory that contains our files (bundled data dir in frozen builds)."""
    if getattr(sys, "frozen", False):
        return sys._MEIPASS  # type: ignore[attr-defined]
    return os.path.dirname(os.path.abspath(__file__))


def appdata_dir() -> str:
    """Per-user directory for cached AI models (and any saved profiles)."""
    if platform.system() == "Windows":
        base = os.environ.get("APPDATA", os.path.expanduser("~"))
    elif platform.system() == "Darwin":
        base = os.path.expanduser("~/Library/Application Support")
    else:
        base = os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
    d = os.path.join(base, APP_NAME)
    os.makedirs(d, exist_ok=True)
    return d


# ---------------------------------------------------------------------------
# ffmpeg helpers (bundled binary in frozen builds, system ffmpeg otherwise)
# ---------------------------------------------------------------------------


def ffmpeg_exe() -> str:
    if getattr(sys, "frozen", False):
        exe = os.path.join(app_dir(), "ffmpeg.exe" if platform.system() == "Windows" else "ffmpeg")
        if os.path.exists(exe):
            # ensure executable bit
            try:
                os.chmod(exe, 0o755)
            except OSError:
                pass
            return exe
    return "ffmpeg"


def ffprobe_exe() -> str:
    if getattr(sys, "frozen", False):
        exe = os.path.join(app_dir(), "ffprobe.exe" if platform.system() == "Windows" else "ffprobe")
        if os.path.exists(exe):
            try:
                os.chmod(exe, 0o755)
            except OSError:
                pass
            return exe
    return "ffprobe"


class FfmpegError(RuntimeError):
    pass


def probe_duration(path: str) -> float:
    try:
        out = subprocess.run(
            [ffprobe_exe(), "-hide_banner", "-v", "error",
             "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", path],
            capture_output=True, text=True, timeout=60,
        )
        return float(out.stdout.strip().splitlines()[0])
    except (ValueError, IndexError, subprocess.SubprocessError) as exc:
        raise FfmpegError(f"Could not read media file: {path}") from exc


def extract_wav(path: str, start: float, stop: float, workdir: str) -> str:
    """Extract the selected range as mono 48 kHz WAV (PCM 16-bit)."""
    out = os.path.join(workdir, "clip.wav")
    cmd = [
        ffmpeg_exe(), "-hide_banner", "-v", "error", "-y",
        "-ss", f"{start:.3f}", "-to", f"{stop:.3f}",
        "-i", path,
        "-vn", "-ac", "1", "-ar", "48000",
        "-c:a", "pcm_s16le", out,
    ]
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
    if r.returncode != 0 or not os.path.exists(out):
        raise FfmpegError(f"Audio extraction failed:\n{r.stderr.strip()[:1500]}")
    return out


def wav_stats(path: str):
    """Return (nframes, samplerate, rms) of a 16-bit WAV."""
    with wave.open(path, "rb") as w:
        n = w.getnframes()
        sr = w.getframerate()
        data = w.readframes(n)
    import array
    a = array.array("h")
    a.frombytes(data[: len(data) // 2 * 2])
    if len(a) == 0:
        return n, sr, 0.0
    rms = (sum((x / 32768.0) ** 2 for x in a) / len(a)) ** 0.5
    return n, sr, float(rms)


def peaks(path: str, buckets: int = 2000):
    """Downsample a 16-bit WAV to (max, min) per bucket for plotting."""
    with wave.open(path, "rb") as w:
        n = w.getnframes()
        data = w.readframes(n)
    import array
    a = array.array("h")
    a.frombytes(data[: len(data) // 2 * 2])
    m = len(a)
    if m == 0:
        return [], []
    step = max(1, m // buckets)
    maxs, mins = [], []
    for i in range(0, m, step):
        chunk = a[i:i + step]
        maxs.append(max(chunk) / 32768.0)
        mins.append(min(chunk) / 32768.0)
    return maxs, mins


# ---------------------------------------------------------------------------
# Voice isolation (lightweight, CPU-only)
# ---------------------------------------------------------------------------

try:
    import numpy as np
    import noisereduce as nr
    HAVE_NR = True
    NR_IMPORT_ERROR = ""
except Exception as _exc:  # pragma: no cover
    HAVE_NR = False
    NR_IMPORT_ERROR = str(_exc)


class VoiceReducer:
    """STFT noise reduction with an optional noise reference clip."""

    def __init__(self, intensity: float, stationary: bool):
        self.intensity = float(intensity)          # 0.0 .. 1.0
        self.stationary = stationary
        self.noise_ref = None                        # optional numpy array of noise

    def set_noise_ref(self, noise):
        self.noise_ref = noise

    def reduce(self, y: "np.ndarray", sr: int) -> "np.ndarray":
        kwargs = dict(
            y=y, sr=sr,
            prop_decrease=min(0.95, self.intensity),
            stationary=self.stationary,
            use_tqdm=False,
        )
        if self.noise_ref is not None:
            kwargs["y_noise"] = self.noise_ref
        return nr.reduce_noise(**kwargs)


# ---------------------------------------------------------------------------
# Separation backends — one mode at a time (DSP and AI are alternatives, not
# stacked: running noisereduce after Demucs adds double-processing artifacts)
# ---------------------------------------------------------------------------

MODE_NONE = "none"   # loudness-normalize only
MODE_DSP = "dsp"     # fast STFT noise reduction (noisereduce)
MODE_AI = "ai"       # deep-learning separation (Demucs htdemucs_ft)

AI_MODEL_NAME = "htdemucs_ft"   # MIT-licensed weights; we keep its "vocals" stem


class Separator:
    """Common interface for the voice-processing modes."""

    name = "base"

    def available(self) -> tuple[bool, str]:
        """Cheap check (no heavy imports). Returns (ok, human-readable reason)."""
        raise NotImplementedError

    def separate(self, y: "np.ndarray", sr: int, log=print) -> tuple["np.ndarray", int]:
        """Return (float32 mono audio, sample_rate)."""
        raise NotImplementedError

    def cancel(self):
        """Best-effort cancel of a running separation (no-op for fast modes)."""


def _ai_deps_present() -> bool:
    """Fast presence check for torch+demucs without importing them (importing
    torch takes ~1-2 s and ~400 MB RAM — too expensive for a startup UI check)."""
    import importlib.util
    try:
        return (importlib.util.find_spec("torch") is not None
                and importlib.util.find_spec("demucs") is not None)
    except Exception:
        return False


class DSSeparator(Separator):
    """Current noisereduce STFT path — unchanged behavior."""

    name = MODE_DSP

    def __init__(self, intensity: float = 0.8, stationary: bool = False):
        self.intensity = float(intensity)
        self.stationary = stationary

    def available(self) -> tuple[bool, str]:
        if not HAVE_NR:
            return False, (f"noisereduce is not installed "
                           f"({NR_IMPORT_ERROR or 'unknown import error'})")
        return True, ""

    def separate(self, y, sr, log=print) -> tuple["np.ndarray", int]:
        log("Isolating voice — fast noise reduction (CPU)…")
        red = VoiceReducer(self.intensity, self.stationary)
        return red.reduce(y, sr), sr


class DemucsSeparator(Separator):
    """AI voice separation with Demucs (htdemucs_ft). Lazy imports; GPU-aware.

    Only imported/used when the AI mode is actually selected, so the core app
    stays a hard-dependency-free ~100 MB tool.
    """

    name = MODE_AI

    def __init__(self, model_name: str = AI_MODEL_NAME, device: str | None = None):
        self.model_name = model_name
        self.device = device          # None → auto-pick: cuda → mps → cpu
        self._api = None              # demucs.api.Separator (lazy, worker thread)
        self._device_used = None
        self._cancel = threading.Event()

    # -- cheap checks (safe to call from the GUI thread) -------------------
    def available(self) -> tuple[bool, str]:
        if not _ai_deps_present():
            return False, ("demucs/torch not installed — "
                           "pip install -r requirements-ai.txt "
                           "(pulls torch, ~2 GB with CUDA wheels)")
        return True, ""

    def cancel(self):
        self._cancel.set()

    # -- heavy work (worker thread only) ------------------------------------
    def _pick_device(self) -> str:
        if self.device:
            return self.device
        import torch
        if torch.cuda.is_available():
            return "cuda"
        mps = getattr(torch.backends, "mps", None)
        if mps is not None and mps.is_available():
            return "mps"
        return "cpu"

    def _get_api(self, log):
        if self._api is None:
            # Cache downloads under <appdata>/VoiceGrab/models. Must be set
            # before torch/huggingface_hub are first imported in this thread.
            root = os.path.join(appdata_dir(), "models")
            os.makedirs(root, exist_ok=True)
            os.environ.setdefault("TORCH_HOME", os.path.join(root, "torch"))
            os.environ.setdefault("HF_HOME", os.path.join(root, "huggingface"))
            from demucs.api import Separator as _DemucsApi
            self._device_used = self._pick_device()
            log(f"Loading AI model '{self.model_name}' "
                f"(first use downloads ~90 MB, cached afterwards)…")
            self._api = _DemucsApi(self.model_name, device=self._device_used)
            log("AI model ready.")
        return self._api

    @staticmethod
    def _model_segment(model) -> float:
        try:
            for sub in model.models:      # BagOfModels
                seg = getattr(sub, "segment", None)
                if seg:
                    return float(seg)
        except AttributeError:
            pass
        seg = getattr(model, "segment", None)
        return float(seg) if seg else 8.0

    @staticmethod
    def _estimate_total_chunks(d: dict, api) -> int:
        """Expected 'end' callback count = submodels × shifts × ceil(len/stride)."""
        try:
            length = int(d["audio_length"])
            seg = DemucsSeparator._model_segment(api.model)
            n_subs = int(d.get("models", 1) or 1)
            stride = int(0.75 * seg * api.samplerate)   # overlap=0.25
            if stride <= 0:
                return 0
            return max(1, math.ceil(length / stride)) * n_subs
        except Exception:
            return 0   # unknown → log raw chunk counts instead

    def separate(self, y, sr, log=print) -> tuple["np.ndarray", int]:
        import numpy as np
        import torch

        api = self._get_api(log)
        device = self._device_used
        if device == "cuda":
            log("AI voice separation on GPU (CUDA)…")
        elif device == "mps":
            log("AI voice separation on GPU (Apple MPS)…")
        else:
            log("AI voice separation on CPU (no GPU detected — expect a wait)…")

        # (C, T) float32 at the original rate; demucs resamples to 44.1 kHz
        # and duplicates mono→stereo for us (see demucs.audio.convert_audio).
        wav = torch.from_numpy(np.ascontiguousarray(y, dtype=np.float32))
        if wav.ndim == 1:
            wav = wav[None, :]

        state = {"done": 0, "total": None, "last_pct": -1}

        def _progress(d):
            if self._cancel.is_set():
                raise KeyboardInterrupt()   # demucs' documented way to abort
            if state["total"] is None:
                state["total"] = self._estimate_total_chunks(d, api)
            if d.get("state") == "end":
                state["done"] += 1
                total = state["total"]
                if total:
                    pct = min(99, int(state["done"] * 100 / total))
                    if state["done"] >= total or pct >= state["last_pct"] + 10:
                        state["last_pct"] = pct
                        log(f"AI separation: {state['done']}/{total} chunks ({pct}%)…")
                else:
                    log(f"AI separation: chunk {state['done']} done…")

        self._cancel.clear()
        try:
            api.update_parameter(callback=_progress)
            _ref, stems = api.separate_tensor(wav, sr)
        except KeyboardInterrupt:
            raise FfmpegError("Cancelled by user.")

        key = ("vocals" if "vocals" in stems
               else "speech" if "speech" in stems
               else next(iter(stems)))
        v = stems[key]
        while v.dim() > 1:
            v = v[0]                       # first channel → mono
        out = v.detach().cpu().numpy().astype(np.float32)
        out = np.where(np.isfinite(out), out, 0.0)   # defensive: never emit NaN/Inf
        return out, int(api.samplerate)


# ---------------------------------------------------------------------------
# Audio IO helpers (numpy/soundfile)
# ---------------------------------------------------------------------------


def read_wav_f32(path: str):
    import soundfile as sf
    data, sr = sf.read(path, dtype="float32")
    return data, sr


def write_wav_f32(path: str, y, sr: int):
    import soundfile as sf
    sf.write(path, y, sr, subtype="PCM_16")


# ---------------------------------------------------------------------------
# GUI
# ---------------------------------------------------------------------------


def _import_pyside():
    from PySide6 import QtCore, QtGui, QtWidgets  # noqa: F401
    return QtCore, QtGui, QtWidgets


def _peak_normalize(y):
    """Peak-normalize (99.5th percentile) to ~ -1 dBFS reference."""
    import numpy as _np
    p = _np.percentile(_np.abs(y), 99.5)
    if p < 1e-6:
        return y
    return (y / p * 0.891)  # ~ -1 dBFS reference peak


def run_gui() -> int:
    QtCore, QtGui, QtWidgets = _import_pyside()

    class Worker(QtCore.QThread):
        log = QtCore.Signal(str)
        done = QtCore.Signal(object, str)   # (success, message)

        def __init__(self, job):
            super().__init__()
            self.job = job
            self._sep: Separator | None = None

        def cancel(self):
            if self._sep is not None:
                self._sep.cancel()

        def run(self):
            import numpy as _np
            try:
                self.log.emit("Extracting audio range…")
                with tempfile.TemporaryDirectory(prefix="voicegrab-") as wd:
                    wav = extract_wav(self.job["input"], self.job["start"], self.job["stop"], wd)
                    y, sr = read_wav_f32(wav)

                    mode = self.job.get("mode", MODE_DSP)
                    if mode == MODE_AI:
                        self._sep = DemucsSeparator()
                    elif mode == MODE_DSP:
                        self._sep = DSSeparator(self.job["intensity"], self.job["stationary"])
                    else:
                        self._sep = None
                    if self._sep is not None:
                        ok, reason = self._sep.available()
                        if not ok:
                            raise FfmpegError(f"Voice processing unavailable: {reason}")
                        y, sr = self._sep.separate(y, sr, self.log.emit)
                        if mode == MODE_AI:
                            # Separated vocals sit a few dB below the mix; bring
                            # peaks up before loudnorm (post-gain step).
                            self.log.emit("Leveling separated vocals…")
                            y = _peak_normalize(y)
                    else:
                        self.log.emit("Normalizing loudness…")
                        y = _peak_normalize(y)

                    work_wav = os.path.join(wd, "work.wav")
                    write_wav_f32(work_wav, y, sr)

                    self.log.emit("Encoding MP3 (320 kbps, loudness-normalized)…")
                    enc = [
                        ffmpeg_exe(), "-hide_banner", "-v", "error", "-y",
                        "-i", work_wav,
                        "-af", f"loudnorm=I={self.job['lufts']}:TP=-1.5:LRA=11",
                        "-c:a", "libmp3lame", "-b:a", "320k", self.job["output"],
                    ]
                    r = subprocess.run(enc, capture_output=True, text=True, timeout=300)
                    if r.returncode != 0 or not os.path.exists(self.job["output"]):
                        raise FfmpegError(f"MP3 encoding failed:\n{r.stderr.strip()[:1500]}")
                    self.log.emit("Done.")
                    self.done.emit(True, self.job["output"])
            except Exception as exc:
                self.done.emit(False, str(exc))
    return _main_loop(QtCore, QtGui, QtWidgets, Worker)


def _main_loop(QtCore, QtGui, QtWidgets, Worker) -> int:
    app = QtWidgets.QApplication(sys.argv)
    app.setApplicationName(APP_NAME)
    app.setApplicationVersion(APP_VERSION)

    class MainWindow(QtWidgets.QMainWindow):
        def __init__(self):
            super().__init__()
            self.setWindowTitle(f"{APP_NAME} — MP3 voice sample cutter")
            self.resize(920, 640)
            self.input_path = None
            self.duration = 0.0
            self.worker = None
            self._build()

        # ---------- UI ----------
        def _build(self):
            central = QtWidgets.QWidget()
            self.setCentralWidget(central)
            lay = QtWidgets.QVBoxLayout(central)
            lay.setContentsMargins(14, 14, 14, 14)
            lay.setSpacing(12)

            # Row 1: input file
            row1 = QtWidgets.QHBoxLayout()
            row1.addWidget(QtWidgets.QLabel("Video / audio file:"))
            self.in_edit = QtWidgets.QLineEdit()
            self.in_edit.setPlaceholderText("Choose an .mp4 / .mkv / .mov / .webm … file")
            row1.addWidget(self.in_edit, 1)
            btn_open = QtWidgets.QPushButton("Open…")
            btn_open.clicked.connect(self.choose_input)
            row1.addWidget(btn_open)
            self.lbl_info = QtWidgets.QLabel("")
            self.lbl_info.setStyleSheet("color:#64748b;")
            row1.addWidget(self.lbl_info)
            lay.addLayout(row1)

            # Row 2: times
            row2 = QtWidgets.QHBoxLayout()
            row2.addWidget(QtWidgets.QLabel("Start:"))
            self.t_start = QtWidgets.QLineEdit("0:00.00")
            self.t_start.setFixedWidth(110)
            row2.addWidget(self.t_start)
            row2.addSpacing(12)
            row2.addWidget(QtWidgets.QLabel("Stop:"))
            self.t_stop = QtWidgets.QLineEdit("0:00.00")
            self.t_stop.setFixedWidth(110)
            row2.addWidget(self.t_stop)
            self.lbl_range = QtWidgets.QLabel("")
            self.lbl_range.setStyleSheet("color:#64748b;")
            row2.addWidget(self.lbl_range)
            row2.addStretch(1)
            btn_preview = QtWidgets.QPushButton("Load preview")
            btn_preview.setToolTip("Extract the selected range so you can listen before saving")
            btn_preview.clicked.connect(self.load_preview)
            row2.addWidget(btn_preview)
            btn_listen = QtWidgets.QPushButton("▶ Play range")
            btn_listen.clicked.connect(self.play_range)
            row2.addWidget(btn_listen)
            lay.addLayout(row2)

            # Row 3: processing mode (one of: fast DSP / AI / none)
            proc_box = QtWidgets.QGroupBox("Voice processing (choose one)")
            proc_lay = QtWidgets.QVBoxLayout(proc_box)
            self.rb_dsp = QtWidgets.QRadioButton(
                "Fast noise reduction (CPU) — fan / hum / room tone; no downloads")
            self.rb_ai = QtWidgets.QRadioButton(
                "AI voice separation (Demucs) — background music & second speakers; "
                "uses GPU if available")
            self.rb_none = QtWidgets.QRadioButton(
                "No processing — just loudness-normalize")
            self.rb_dsp.setChecked(True)
            self.rb_ai.setToolTip(
                "Deep-learning vocal separation. First use downloads ~90 MB of models. "
                "GPU (NVIDIA CUDA / Apple) is strongly recommended; on CPU a 60 s clip "
                "takes a few minutes. Long AI runs can be cancelled between chunks.")
            for rb in (self.rb_dsp, self.rb_ai, self.rb_none):
                proc_lay.addWidget(rb)
            lay.addWidget(proc_box)
            self.rb_dsp.toggled.connect(self._proc_mode_changed)
            self.rb_ai.toggled.connect(self._proc_mode_changed)

            iso_row = QtWidgets.QHBoxLayout()
            iso_row.addWidget(QtWidgets.QLabel("Noise reduction strength:"))
            self.spn_intensity = QtWidgets.QSpinBox()
            self.spn_intensity.setRange(0, 100)
            self.spn_intensity.setValue(80)
            self.spn_intensity.setSuffix("%")
            iso_row.addWidget(self.spn_intensity)
            iso_row.addSpacing(16)
            self.chk_stationary = QtWidgets.QCheckBox(
                "Stationary noise (fan / hum — better if constant)")
            iso_row.addWidget(self.chk_stationary)
            iso_row.addStretch(1)
            lay.addLayout(iso_row)

            self.lbl_proc_status = QtWidgets.QLabel("")
            self.lbl_proc_status.setStyleSheet("color:#64748b;")
            self.lbl_proc_status.setWordWrap(True)
            lay.addWidget(self.lbl_proc_status)

            out_row = QtWidgets.QHBoxLayout()
            out_row.addWidget(QtWidgets.QLabel("Output MP3:"))
            self.out_edit = QtWidgets.QLineEdit()
            self.out_edit.setPlaceholderText("Defaults to <input>_voicegrab.mp3 next to the source")
            out_row.addWidget(self.out_edit, 1)
            btn_out = QtWidgets.QPushButton("Browse…")
            btn_out.clicked.connect(self.choose_output)
            out_row.addWidget(btn_out)
            lay.addLayout(out_row)

            # Waveform
            lay.addWidget(QtWidgets.QLabel("Waveform of the selected range:"))
            self.wave = WaveView()
            self.wave.setMinimumHeight(160)
            lay.addWidget(self.wave, 1)

            # Bottom bar
            bar = QtWidgets.QHBoxLayout()
            self.btn_cancel = QtWidgets.QPushButton("Cancel")
            self.btn_cancel.setMinimumHeight(40)
            self.btn_cancel.setEnabled(False)
            self.btn_cancel.setToolTip(
                "Cancels an in-progress AI separation at the next chunk boundary. "
                "Fast DSP runs finish in seconds and are not cancelable.")
            self.btn_cancel.clicked.connect(self._cancel)
            bar.addWidget(self.btn_cancel)
            self.btn_export = QtWidgets.QPushButton("Export MP3")
            self.btn_export.setMinimumHeight(40)
            self.btn_export.setStyleSheet(
                "font-size:14px; font-weight:600; background:#0ea5e9; color:white;"
                "border-radius:6px; padding:4px 18px;")
            self.btn_export.clicked.connect(self.export)
            bar.addWidget(self.btn_export)
            self.lbl_status = QtWidgets.QLabel("Ready.")
            self.lbl_status.setStyleSheet("color:#64748b;")
            bar.addWidget(self.lbl_status, 1)
            lay.addLayout(bar)

            # ----- availability / default mode (cheap checks only) -----
            dsp_ok, dsp_reason = DSSeparator().available()
            ai_ok, ai_reason = DemucsSeparator().available()
            self.spn_intensity.setEnabled(dsp_ok)
            self.chk_stationary.setEnabled(dsp_ok)
            if not dsp_ok:
                self.rb_dsp.setEnabled(False)
                self.rb_none.setChecked(True)
            self.rb_ai.setEnabled(ai_ok)
            self._proc_mode_changed(self.rb_dsp.isChecked())
            hints = []
            if not ai_ok:
                hints.append(f"⚠ AI separation unavailable — {ai_reason}")
            else:
                hints.append(
                    "AI mode ready — GPU (NVIDIA/Apple) used if present, otherwise CPU "
                    "(slower); first use downloads ~90 MB of models, cached afterwards.")
            if not dsp_ok:
                hints.append(f"⚠ fast noise reduction unavailable — {dsp_reason}")
            self.lbl_proc_status.setText(" ".join(hints))

        # ---------- helpers ----------
        def _proc_mode_changed(self, *_):
            dsp = self.rb_dsp.isChecked()
            self.spn_intensity.setEnabled(dsp)
            self.chk_stationary.setEnabled(dsp)

        @staticmethod
        def _parse_time(s: str) -> float:
            """Accept 125, 1:05, 1:05.5, 1:05:00, 01:00:00.500."""
            s = s.strip()
            parts = s.split(":")
            if len(parts) > 3:
                raise ValueError("Invalid time format")
            t = 0.0
            for i, p in enumerate(parts):
                v = float(p)
                t = t * 60 + v
            return max(0.0, t)

        def _fmt_time(self, t: float) -> str:
            m, s = divmod(t, 60)
            h, m = divmod(int(m), 60)
            if h:
                return f"{h}:{int(m):02d}:{s:05.2f}"
            return f"{int(m)}:{s:05.2f}"

        def _times(self):
            start = self._parse_time(self.t_start.text() or "0")
            stop = self._parse_time(self.t_stop.text() or "0")
            if start >= stop:
                raise ValueError("Start must be before stop.")
            return start, stop

        def refresh_info(self):
            if self.input_path and self.duration:
                self.lbl_info.setText(
                    f"Duration {self._fmt_time(self.duration)}   Range {self._range_text()}")
            else:
                self.lbl_info.setText("")

        def _range_text(self):
            try:
                s, e = self._times()
                return f"{s:.2f}s – {e:.2f}s  (clip {e - s:.2f}s)"
            except ValueError:
                return "— invalid range —"

        # ---------- slots ----------
        def choose_input(self):
            path, _ = QtWidgets.QFileDialog.getOpenFileName(
                self, "Choose video or audio file", "",
                "Media files (*.mp4 *.mkv *.mov *.webm *.avi *.m4a *.mp3 *.wav *.ogg *.flv *.ts);;All files (*)")
            if not path:
                return
            self.input_path = path
            self.in_edit.setText(path)
            self.t_stop.setText(self._fmt_time(self.duration))  # after probe below
            try:
                self.duration = probe_duration(path)
                self.t_stop.setText(self._fmt_time(self.duration))
            except FfmpegError as exc:
                self.status(str(exc)); return
            self.refresh_info()
            self.load_preview()

        def choose_output(self):
            if not self.input_path:
                return
            base = os.path.splitext(os.path.basename(self.input_path))[0]
            start = os.path.join(os.path.dirname(self.input_path), base + "_voicegrab.mp3")
            path, _ = QtWidgets.QFileDialog.getSaveFileName(
                self, "Save MP3 as", start, "MP3 files (*.mp3)")
            if path:
                self.out_edit.setText(path)

        def load_preview(self):
            if not self.input_path:
                return
            try:
                s, e = self._times()
            except ValueError as exc:
                self.status(str(exc)); return
            try:
                with tempfile.TemporaryDirectory(prefix="voicegrab-") as wd:
                    wav = extract_wav(self.input_path, s, e, wd)
                    maxs, mins = peaks(wav)
                    n, sr, rms = wav_stats(wav)
            except FfmpegError as exc:
                self.status(str(exc)); return
            self.wave.set_peaks(maxs, mins)
            self.lbl_range.setText(f"clip {n / sr:.2f}s @ {sr} Hz   RMS {rms:.3f}")

        def play_range(self):
            if not self.input_path:
                return
            try:
                s, e = self._times()
            except ValueError as exc:
                self.status(str(exc)); return
            player = os.path.join(tempfile.gettempdir(), "voicegrab_preview.mp3")
            try:
                cmd = [
                    ffmpeg_exe(), "-hide_banner", "-v", "error", "-y",
                    "-ss", f"{s:.3f}", "-to", f"{e:.3f}", "-i", self.input_path,
                    "-vn", "-c:a", "libmp3lame", "-b:a", "128k", player,
                ]
                subprocess.run(cmd, check=True, capture_output=True, text=True, timeout=300)
                if platform.system() == "Windows":
                    os.startfile(player)  # type: ignore[attr-defined]
                elif platform.system() == "Darwin":
                    subprocess.Popen(["open", player])
                else:
                    subprocess.Popen(["xdg-open", player])
                self.status("Playing range in system player…")
            except (subprocess.SubprocessError, OSError) as exc:
                self.status(f"Could not play range: {exc}")

        def status(self, msg: str):
            self.lbl_status.setText(msg)

        def export(self):
            if self.worker is not None:
                self.status("Already working…"); return
            if not self.input_path:
                QtWidgets.QMessageBox.warning(self, APP_NAME, "Choose an input file first.")
                return
            try:
                s, e = self._times()
            except ValueError as exc:
                QtWidgets.QMessageBox.warning(self, APP_NAME, str(exc)); return
            out = self.out_edit.text().strip()
            if not out:
                base = os.path.splitext(os.path.basename(self.input_path))[0]
                out = os.path.join(os.path.dirname(os.path.abspath(self.input_path)),
                                   base + "_voicegrab.mp3")
            self.btn_export.setEnabled(False)
            self.btn_cancel.setEnabled(True)
            self.status("Working… (see log in status area)")
            if self.rb_ai.isChecked():
                mode = MODE_AI
            elif self.rb_none.isChecked():
                mode = MODE_NONE
            else:
                mode = MODE_DSP
            self.worker = Worker({
                "input": self.input_path, "start": s, "stop": e,
                "output": out,
                "mode": mode,
                "intensity": self.spn_intensity.value() / 100.0,
                "stationary": self.chk_stationary.isChecked(),
                "lufts": -16,
            })
            self.worker.log.connect(self.status)
            self.worker.done.connect(self._export_done)
            self.worker.start()

        def _cancel(self):
            if self.worker is None:
                return
            self.worker.cancel()
            self.status("Cancelling… (AI separation stops at the next chunk boundary)")

        def _export_done(self, ok, msg):
            self.btn_export.setEnabled(True)
            self.btn_cancel.setEnabled(False)
            self.worker = None
            if ok:
                self.status(f"✔ Saved: {msg}")
                QtWidgets.QMessageBox.information(
                    self, APP_NAME,
                    f"MP3 saved to:\n{msg}\n\nTip: 30–120 s of clean speech is usually "
                    f"enough for high-quality AI voice cloning.")
            else:
                self.status(f"✖ {msg.splitlines()[0]}")
                QtWidgets.QMessageBox.critical(self, APP_NAME, msg)

    # ---------------- wave view ----------------
    class WaveView(QtWidgets.QWidget):
        def __init__(self):
            super().__init__()
            self.maxs: list = []
            self.mins: list = []
            self.setMinimumSize(120, 80)

        def set_peaks(self, maxs, mins):
            self.maxs = list(maxs)
            self.mins = list(mins)
            self.update()

        def paintEvent(self, ev):
            with QtGui.QPainter(self) as p:
                w, h = self.width(), self.height()
                mid = h / 2
                p.fillRect(0, 0, w, h, QtGui.QColor("#0f172a"))
                p.setPen(QtGui.QColor("#334155"))
                p.drawLine(0, int(mid), w, int(mid))
                if not self.maxs:
                    p.setPen(QtGui.QColor("#94a3b8"))
                    p.drawText(0, 0, w, h, int(QtCore.Qt.AlignCenter), "Open a file to see the waveform")
                    return
                n = len(self.maxs)
                color = QtGui.QColor("#38bdf8")
                for i, (mx, mn) in enumerate(zip(self.maxs, self.mins)):
                    x = i * w / n
                    y1 = mid - mx * (mid - 4)
                    y2 = mid - mn * (mid - 4)
                    p.setPen(color)
                    p.drawLine(int(x), int(y1), int(x), max(int(y1), int(y2)))

    win = MainWindow()
    win.show()
    return app.exec()

if __name__ == "__main__":
    if "--cli" in sys.argv:
        # Simple headless mode:
        #   voicegrab --cli <input> <start> <stop> [output.mp3] [--mode none|dsp|ai]
        sys.argv = [a for a in sys.argv if a != "--cli"]
        mode = MODE_DSP
        if "--mode" in sys.argv:
            i = sys.argv.index("--mode")
            try:
                mode = sys.argv[i + 1]
            except IndexError:
                print("usage: --mode needs a value: none | dsp | ai")
                raise SystemExit(2)
            del sys.argv[i:i + 2]
        if mode not in (MODE_NONE, MODE_DSP, MODE_AI):
            print(f"unknown --mode {mode!r} (use: none | dsp | ai)")
            raise SystemExit(2)
        if len(sys.argv) < 4:
            print("usage: voicegrab --cli <input> <start> <stop> [output.mp3] [--mode none|dsp|ai]")
            raise SystemExit(2)
        inp, s, e = sys.argv[1], float(sys.argv[2]), float(sys.argv[3])
        out = sys.argv[4] if len(sys.argv) > 4 else os.path.join(
            os.path.dirname(os.path.abspath(inp)),
            os.path.splitext(os.path.basename(inp))[0] + "_voicegrab.mp3")
        import numpy as _np
        with tempfile.TemporaryDirectory(prefix="voicegrab-") as wd:
            wav = extract_wav(inp, s, e, wd)
            y, sr = read_wav_f32(wav)
            if mode == MODE_AI:
                sep = DemucsSeparator()
                ok, reason = sep.available()
                if not ok:
                    print(f"AI separation unavailable: {reason}")
                    raise SystemExit(3)
                y, sr = sep.separate(y, sr, print)
                y = _peak_normalize(y)   # post-gain before loudnorm
            elif mode == MODE_DSP:
                sep = DSSeparator(intensity=0.8, stationary=True)
                ok, reason = sep.available()
                if not ok:
                    print(f"Fast noise reduction unavailable: {reason}")
                    raise SystemExit(3)
                y, sr = sep.separate(y, sr, print)
            red_wav = os.path.join(wd, "reduced.wav")
            write_wav_f32(red_wav, y, sr)
            subprocess.run([
                ffmpeg_exe(), "-hide_banner", "-v", "error", "-y",
                "-i", red_wav, "-af", "loudnorm=I=-16:TP=-1.5:LRA=11",
                "-c:a", "libmp3lame", "-b:a", "320k", out,
            ], check=True)
        print(f"Saved {out}")
    else:
        raise SystemExit(run_gui())