#!/usr/bin/env python3
from __future__ import annotations

import os
import warnings

# Chatterbox pulls transformers/diffusers/tqdm; their progress + deprecations spam stderr and
# bloat ~/.mugi/logs/chatterbox_worker.log. Set before any ML stack import (lazy in _run).
os.environ.setdefault("TQDM_DISABLE", "1")
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)

import argparse
import importlib
import io
import json
import sys
import tempfile
from typing import Any


_MODEL_CACHE: dict[tuple[str, str], Any] = {}


def _load_first_available(candidates: list[tuple[str, str]]):
    for module_name, class_name in candidates:
        try:
            mod = importlib.import_module(module_name)
            klass = getattr(mod, class_name)
            return klass
        except Exception:
            continue
    names = ", ".join([f"{m}.{c}" for m, c in candidates])
    raise RuntimeError(f"No compatible Chatterbox class found ({names})")


def _pick_device() -> str:
    forced = (os.environ.get("CHATTERBOX_DEVICE") or "").strip().lower()
    if forced in {"cuda", "mps", "cpu"}:
        return forced
    try:
        import torch  # type: ignore

        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
    except Exception:
        pass
    return "cpu"


def _write_wav(data, sample_rate: int, out_path: str) -> None:
    # Chatterbox commonly returns torch tensor-like objects. Support tensor/ndarray/list.
    try:
        import numpy as np
    except Exception as exc:
        raise RuntimeError(f"numpy missing: {exc}") from exc

    if hasattr(data, "detach"):
        data = data.detach()
    if hasattr(data, "cpu"):
        data = data.cpu()
    if hasattr(data, "numpy"):
        arr = data.numpy()
    else:
        arr = np.asarray(data)

    if arr.ndim == 1:
        arr = arr[None, :]
    if arr.ndim != 2:
        raise RuntimeError(f"unexpected wav shape: {arr.shape}")

    # soundfile expects (frames, channels).
    arr = arr.T
    try:
        import soundfile as sf  # type: ignore

        sf.write(out_path, arr, sample_rate)
    except Exception as exc:
        raise RuntimeError(f"failed to write wav: {exc}") from exc


def _load_payload(arg: str) -> dict:
    if arg == "-":
        return json.load(sys.stdin)
    if arg.startswith("{"):
        return json.loads(arg)
    with open(arg, "r", encoding="utf-8") as f:
        return json.load(f)


def _perth_watermarker_error() -> str | None:
    """Chatterbox TTS constructs perth.PerthImplicitWatermarker(); import can fail silently (None)."""
    try:
        import perth  # type: ignore
    except Exception as exc:
        return f"resemble-perth import failed: {exc}"
    ctor = getattr(perth, "PerthImplicitWatermarker", None)
    if ctor is None:
        return (
            "resemble-perth did not load PerthImplicitWatermarker (often setuptools≥82 removed "
            "pkg_resources). In ~/.mugi/chatterbox-venv run: bin/pip install 'setuptools>=69,<82' "
            "or use Settings → Software → Chatterbox → Reinstall."
        )
    return None


def _get_cached_model(variant: str, device: str):
    key = (variant, device)
    model = _MODEL_CACHE.get(key)
    if model is not None:
        return model
    if variant == "multilingual":
        klass = _load_first_available(
            [
                ("chatterbox.mtl_tts", "ChatterboxMultilingualTTS"),
                ("chatterbox.multilingual", "ChatterboxMultilingualTTS"),
            ]
        )
    elif variant == "classic":
        klass = _load_first_available(
            [
                ("chatterbox.tts", "ChatterboxTTS"),
            ]
        )
    else:
        klass = _load_first_available(
            [
                ("chatterbox.tts_turbo", "ChatterboxTurboTTS"),
                ("chatterbox.tts", "ChatterboxTurboTTS"),
                ("chatterbox.tts", "ChatterboxTTS"),
            ]
        )
    model = klass.from_pretrained(device=device)
    _MODEL_CACHE[key] = model
    return model


def _run_quiet(payload: dict) -> dict:
    # Hugging Face / transformers sometimes prints progress to stdout, which would prefix our
    # JSON and break parent parsing. Suppress any incidental stdout from the ML stack.
    _junk = io.StringIO()
    _real_out = sys.stdout
    try:
        sys.stdout = _junk
        return _run(payload)
    finally:
        sys.stdout = _real_out


def _handle_request_quiet(payload: dict) -> dict:
    _junk = io.StringIO()
    _real_out = sys.stdout
    try:
        sys.stdout = _junk
        return _handle_request(payload)
    finally:
        sys.stdout = _real_out


def _run(payload: dict) -> dict:
    text = str(payload.get("text") or "").strip()
    if not text:
        return {"ok": False, "error": "text is required"}
    ref = str(payload.get("audio_prompt_path") or "").strip()
    if not ref:
        return {"ok": False, "error": "audio_prompt_path is required"}
    if not os.path.isfile(ref):
        return {"ok": False, "error": "audio_prompt_path not found"}

    perth_err = _perth_watermarker_error()
    if perth_err:
        return {"ok": False, "error": perth_err}

    variant = str(payload.get("variant") or "turbo").strip().lower()
    language_id = str(payload.get("language_id") or "en").strip() or "en"
    device = _pick_device()

    sample_rate = 24_000
    try:
        if variant == "multilingual":
            model = _get_cached_model(variant, device)
            wav = model.generate(text, audio_prompt_path=ref, language_id=language_id)
            sample_rate = int(getattr(model, "sr", 24_000))
        elif variant == "classic":
            model = _get_cached_model(variant, device)
            wav = model.generate(text, audio_prompt_path=ref)
            sample_rate = int(getattr(model, "sr", 24_000))
        else:
            model = _get_cached_model("turbo", device)
            wav = model.generate(text, audio_prompt_path=ref)
            sample_rate = int(getattr(model, "sr", 24_000))
    except Exception as exc:
        return {"ok": False, "error": f"chatterbox generation failed: {exc}", "device": device}

    fd, out_path = tempfile.mkstemp(suffix=".wav")
    os.close(fd)
    try:
        _write_wav(wav, sample_rate, out_path)
    except Exception as exc:
        try:
            os.unlink(out_path)
        except Exception:
            pass
        return {"ok": False, "error": str(exc), "device": device}
    return {"ok": True, "wav_path": out_path, "sample_rate": sample_rate, "device": device}


def _handle_request(payload: dict) -> dict:
    cmd = str(payload.get("cmd") or "").strip().lower()
    if cmd == "prewarm":
        variant = str(payload.get("variant") or "turbo").strip().lower() or "turbo"
        device = _pick_device()
        try:
            _get_cached_model(variant, device)
            return {"ok": True, "prewarmed": True, "variant": variant, "device": device}
        except Exception as exc:
            return {"ok": False, "error": f"chatterbox prewarm failed: {exc}", "device": device}
    return _run(payload)


def _serve_forever() -> int:
    for line in sys.stdin:
        raw = (line or "").strip()
        if not raw:
            continue
        try:
            payload = json.loads(raw)
            if not isinstance(payload, dict):
                raise ValueError("payload must be JSON object")
            out = _run_quiet(payload) if str(payload.get("cmd") or "").strip().lower() == "" else _handle_request_quiet(payload)
        except Exception as exc:
            out = {"ok": False, "error": str(exc)}
        print(json.dumps(out, ensure_ascii=True), flush=True)
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Chatterbox synthesis worker")
    parser.add_argument("--payload", help="JSON string, file path, or '-' for stdin")
    parser.add_argument(
        "--serve",
        action="store_true",
        help="Read newline-delimited JSON payloads from stdin and emit one JSON response per line",
    )
    args = parser.parse_args()
    if args.serve:
        return _serve_forever()
    if not args.payload:
        print(json.dumps({"ok": False, "error": "--payload is required unless --serve is used"}))
        return 1
    try:
        payload = _load_payload(args.payload)
        if not isinstance(payload, dict):
            raise ValueError("payload must be JSON object")
        out = _run_quiet(payload) if str(payload.get("cmd") or "").strip().lower() == "" else _handle_request_quiet(payload)
    except Exception as exc:
        out = {"ok": False, "error": str(exc)}
    print(json.dumps(out, ensure_ascii=True))
    return 0 if out.get("ok") else 1


if __name__ == "__main__":
    raise SystemExit(main())
