#!/usr/bin/env python3
"""Set up the faster-qwen3-tts backend for the audiobook generator.

faster-qwen3-tts is an OpenAI-compatible Qwen3-TTS server with CUDA-graph
inference (NVIDIA GPU required). It always uses voice cloning, with the
reference voice configured on the server through a ``voices.json``. This
module sets the whole backend up end-to-end as a TUI: pip-install the
package into its own managed venv (``app/envs/faster`` — separate from the
app venv and from qwen-tts's: both stacks ship a ``qwen_tts`` module whose
transformers requirements conflict), clone the repo (for
``examples/openai_server.py``), build a ``voices.json`` from a directory of
.wav references (transcribed with Whisper), sync ``app/converter/config.py``, and print the launch command. It is
driven by ``audiobook.py``'s hub but can also be run directly with flags.

Usage:
    python app/backends/faster.py [--wavs WAV_DIR] [--output PATH]
        [--language LANG] [--whisper-model NAME] [--force]
        [--voice NAME] [--skip-install] [--skip-clone]

When the target ``voices.json`` already exists, the TUI wizard runs as a
"modify": it loads the existing voices, pre-fills the language, wav
directory and transcription choice from them instead of prompting to
overwrite, and writes back to the same file.
"""

import argparse
import json
import shutil
import sys
from pathlib import Path
from typing import List, Optional

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from backends import (
    BackendStatus,
    ServerSpec,
    common,
    envs,
    format_launch_hint,
    probe,
    servers,
    setup,
)
from backends.common import (
    APP_DIR,
    VOICES_DIR,
    find_wav_files,
    normalize_dir_arg,
)
from converter import config
from converter.clients import (
    normalize_language,
    transcribe_reference_audio,
    whisper_backend_available,
)
from ui import taskview, tui

FASTER_DIR_NAME = "faster-qwen3-tts"
FASTER_GIT_URL = "https://github.com/andimarafioti/faster-qwen3-tts"
FASTER_PIP_PKG = "faster-qwen3-tts[demo]"
# The dedicated venv faster-qwen3-tts is installed into (never the app env
# or the qwen backend's; the wheel pulls its own qwen-tts-hf dependency,
# which ships the same qwen_tts module upstream qwen-tts does).
FASTER_ENV = envs.FASTER_ENV_DIR
WHISPER_MODELS = ("tiny", "base", "small", "medium", "large-v3")


def _checkout() -> Path:
    return APP_DIR / FASTER_DIR_NAME


def _is_installed() -> bool:
    return envs.module_available("faster_qwen3_tts", FASTER_ENV)


def _is_cloned() -> bool:
    return (_checkout() / "examples" / "openai_server.py").is_file()


def _config_port() -> int:
    import urllib.parse
    try:
        return urllib.parse.urlsplit(config.FASTER_API_URL).port or 8000
    except ValueError:
        return 8000


def build_voices(wav_files: list, language: str, whisper_model: str) -> dict:
    """Transcribe each wav file and build the voices mapping."""
    voices = {}
    for wav_file in wav_files:
        name = wav_file.stem
        print(f"[INFO] Transcribing {wav_file.name} (voice '{name}')...")
        text = transcribe_reference_audio(str(wav_file), model_name=whisper_model)
        if text:
            print(f"[OK] {name}: {text}")
        else:
            print(f"[WARNING] No transcript for '{name}'; the faster backend "
                  "strongly recommends an accurate transcript — consider "
                  "editing voices.json by hand before starting the server")
        voices[name] = {
            "ref_audio": str(wav_file.resolve()),
            "ref_text": text or "",
            "language": language,
        }
    return voices


def load_voices(path: Path) -> dict:
    """Read voices.json into a name -> voice-entry dict, or {} when unusable.

    Returns {} for a missing file, unreadable content, or a non-dict
    document. Used by the wizard's modify flow to seed its defaults from an
    existing voices.json instead of prompting to overwrite it.
    """
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}
    if not isinstance(data, dict):
        return {}
    return data


def _write_voices_json(output_path: Path, wav_dir: Path, language: str,
                       whisper_model: str,
                       plan: Optional[dict]) -> Optional[dict]:
    """Transcribe the wav dir and write voices.json; return the voices dict.

    PLAN (from the wizard's transcription toggle, or an "all" plan for a
    fresh/flag run) decides whether every voice is re-transcribed
    ("all"), only the new ones ("missing" — merged into the existing
    entries), or nothing changes ("keep" — the existing file is left
    untouched and returned as-is). None (cancelled) writes nothing.
    """
    if plan is None:
        return None
    if plan["mode"] == "keep":
        return dict(plan["existing"])
    wav_files = find_wav_files(wav_dir)
    if not wav_files:
        print(f"[ERROR] No .wav files found in {wav_dir}")
        return None
    if whisper_backend_available() is None:
        print("[WARNING] Neither faster_whisper nor whisper was found, so "
              "transcripts will be empty — install one or edit voices.json "
              "by hand.")
    if plan["mode"] == "missing":
        voices = dict(plan["existing"])
        voices.update(build_voices(plan["missing"], language, whisper_model))
    else:
        voices = build_voices(wav_files, language, whisper_model)
    with output_path.open("w", encoding="utf-8") as handle:
        json.dump(voices, handle, indent=4, ensure_ascii=False)
        handle.write("\n")
    print(f"[OK] Wrote {output_path} with {len(voices)} voice(s): "
          f"{', '.join(voices)}")
    return voices


def _plan_for(mode: str, wav_dir: Path, existing_voices: dict) -> dict:
    """Build the transcription PLAN for the chosen form MODE.

    The plan dict is what ``_write_voices_json`` consumes: "missing"
    carries the new .wavs (computed here from the final directory choice)
    plus the existing entries; "all"/"keep" only name the mode.
    """
    if mode == "missing":
        missing = [wav for wav in find_wav_files(wav_dir)
                   if wav.stem not in existing_voices]
        return {"mode": mode, "missing": missing,
                "existing": dict(existing_voices)}
    return {"mode": mode, "missing": [], "existing": {}}


def _wizard(stdscr, args: argparse.Namespace) -> Optional[dict]:
    """Single-form setup: every faster-setup decision on one screen.

    The form mirrors the Generate-audiobooks screen: a Voices-directory
    picker, Language, Whisper model, and — on a modify run with an
    existing voices.json — which voices to transcribe. There is no port
    question: the server port lives in app/converter/config.py (edit it
    in the hub's Settings screen). Install and clone happen without
    asking; Esc or Cancel aborts the whole setup.
    """
    _GO_BACK = object()

    # An existing voices.json seeds the defaults (modify flow) instead of
    # an overwrite prompt; its voices also seed the directory field.
    default_output = args.output
    if default_output is None and _is_cloned():
        default_output = _checkout() / "voices.json"
    existing_voices = {}
    if default_output is not None and default_output.exists() \
            and not args.force:
        existing_voices = load_voices(default_output)
    wav_start = VOICES_DIR
    if existing_voices:
        ref_dirs = {Path(voice["ref_audio"]).parent
                    for voice in existing_voices.values()
                    if isinstance(voice, dict) and voice.get("ref_audio")}
        if len(ref_dirs) == 1:
            wav_start = next(iter(ref_dirs))

    # Install and clone happen without asking: when the package or repo is
    # missing (and not skipped by flag), the tail just does it afterwards.
    do_install = (not _is_installed()) and not args.skip_install
    do_clone = (not _is_cloned()) and not args.skip_clone

    default_language = config.LANGUAGE
    for voice in existing_voices.values():
        if isinstance(voice, dict) and voice.get("language"):
            default_language = voice["language"]
            break

    fields: List[dict] = [
        {"key": "wav_dir", "label": "Voices directory", "kind": "dir",
         "value": Path(args.input_dir) if args.input_dir is not None
         else wav_start,
         "info": common.wav_dir_info, "preview": common.wav_dir_preview},
        {"key": "language", "label": "Language", "kind": "text",
         "value": default_language,
         "validate": lambda s: None if _try_language(s)
         else "Unknown language (e.g. English, en)"},
        {"key": "whisper_model", "label": "Whisper model", "kind": "choice",
         "value": args.whisper_model or "base",
         "choices": list(WHISPER_MODELS)},
    ]
    modifying = bool(existing_voices) and not args.force
    if modifying:
        # Modify flow: a plain in-place toggle, always offered.
        fields.append({
            "key": "transcription", "label": "Transcription",
            "kind": "toggle", "value": "missing",
            "choices": [("Transcribe new voices", "missing"),
                        ("Re-transcribe all voices", "all")],
            "note": "An existing voices.json was found.",
        })

    result = tui.form(
        stdscr, "Set Up faster-qwen3-tts", fields,
        buttons=("Continue", "Cancel"),
        start_on_buttons=False, back_value=_GO_BACK)
    if result is _GO_BACK:
        return None

    wav_dir = Path(result["wav_dir"])
    language = normalize_language(result["language"])
    whisper_model = result["whisper_model"]
    if not modifying:
        plan = {"mode": "all", "missing": [], "existing": {}}
    elif find_wav_files(wav_dir):
        mode = result.get("transcription")
        if mode not in ("missing", "all"):
            mode = "missing"
        plan = _plan_for(mode, wav_dir, existing_voices)
    else:
        # Directory without .wavs on a modify run: keep the existing file.
        plan = {"mode": "keep", "missing": [],
                "existing": dict(existing_voices)}

    # Default into the cloned checkout — also when the clone is still
    # pending in this run's steps (do_clone): detect() and the server
    # launch only read voices.json from there, so a fresh install must
    # not leave the file in the wav directory. The wav-directory
    # fallback keeps runs working without any checkout.
    output_path = default_output
    if output_path is None:
        if _is_cloned() or do_clone:
            output_path = _checkout() / "voices.json"
        else:
            output_path = wav_dir / "voices.json"

    return {
        "do_install": do_install,
        "do_clone": do_clone,
        "wav_dir": wav_dir,
        "language": language,
        "whisper_model": whisper_model,
        "output_path": output_path,
        "force": args.force,
        "plan": plan,
    }


def _try_language(value: str) -> bool:
    try:
        normalize_language(value)
        return True
    except ValueError:
        return False


def _execute_steps(settings: dict) -> List[taskview.TaskStep]:
    """Build the ordered setup steps for the in-TUI task view.

    The same work ``_execute`` runs on the console, split into named steps so
    the view can show per-step state and progress. Subprocess steps (pip
    install, git clone) stream through EMIT and abort on CANCEL; print()-based
    steps are captured by the view's stdout redirect.
    """
    steps: List[taskview.TaskStep] = []

    if settings["do_install"]:
        def install(emit, cancel):
            rc = common.pip_install([FASTER_PIP_PKG], emit=emit,
                                    cancel=cancel, env_dir=FASTER_ENV)
            if rc != 0:
                print(f"[WARNING] pip install failed (exit {rc}); install "
                      f"{FASTER_PIP_PKG} manually")
            else:
                print("[OK] faster-qwen3-tts installed")
            return rc
        steps.append(taskview.TaskStep(
            f"Install {FASTER_PIP_PKG}", install))

    if settings["do_clone"]:
        def clone(emit, cancel):
            rc = common.git_clone(FASTER_GIT_URL, _checkout(),
                                  emit=emit, cancel=cancel)
            if rc != 0:
                print(f"[WARNING] git clone failed (exit {rc}); clone "
                      f"manually: git clone {FASTER_GIT_URL} {_checkout()}")
            else:
                print(f"[OK] cloned into {_checkout()}")
            return rc
        steps.append(taskview.TaskStep(
            "Clone faster-qwen3-tts", clone))

    def write(emit, cancel):
        voices = _write_voices_json(settings["output_path"],
                                    settings["wav_dir"],
                                    settings["language"],
                                    settings["whisper_model"],
                                    settings["plan"])
        if voices is None:
            return 1

        # Sync app/converter/config.py default voice. (The server port is
        # not touched here: it lives in FASTER_API_URL, edited in the
        # Settings screen.)
        default_voice = next(iter(voices))
        if default_voice != config.FASTER_VOICE:
            if common.update_config_value("FASTER_VOICE", default_voice):
                print(f"[OK] Updated FASTER_VOICE to {default_voice}")
            else:
                print("[WARNING] Could not update FASTER_VOICE; edit "
                      "app/converter/config.py by hand")

        _print_launch_hint(settings["output_path"])
        return 0
    steps.append(taskview.TaskStep(
        "Write voices.json & sync config", write))

    return steps


def _execute(settings: dict) -> int:
    """Console tail: install, clone, write voices.json, sync, advise."""
    return taskview.run_steps_inline(_execute_steps(settings))


def _print_launch_hint(voices_path: Path) -> None:
    """Remediation only (troubleshooting): what's missing when not cloned.

    The hub starts and stops the server itself, so a working install gets
    no manual launch instructions.
    """
    if _is_cloned():
        return
    print("[INFO] Clone faster-qwen3-tts to get examples/openai_server.py,")
    print(f"        then run it with --voices {voices_path} "
          f"--port {_config_port()}")


def setup_screen(stdscr) -> int:
    """Run the setup wizard on an existing curses screen (the hub's).

    See backends.setup.screen_flow for the shared flow. Returns 0 on
    completion, 1 when the user aborted.
    """
    return setup.screen_flow(stdscr, wizard=_wizard,
                             steps_of=_execute_steps,
                             title="Setting up faster-qwen3-tts",
                             parser_factory=build_parser)


def run_tui(args: Optional[argparse.Namespace] = None) -> int:
    """Run the faster setup wizard end-to-end."""
    if args is None:
        args = build_parser().parse_args([])
    return setup.tui_flow(_wizard, _execute, args=args,
                          aborted_message="[INFO] Aborted")


def _collect_from_flags(args: argparse.Namespace,
                        parser: argparse.ArgumentParser) -> Optional[dict]:
    """Build the settings dict from flags for a non-interactive run."""
    wav_dir = args.input_dir if args.input_dir is not None else VOICES_DIR
    if not wav_dir.is_dir():
        parser.error(f"WAV directory not found: {wav_dir}")
    try:
        language = normalize_language(args.language or config.LANGUAGE)
    except ValueError as exc:
        parser.error(str(exc))
    do_install = (not _is_installed()) and not args.skip_install
    do_clone = (not _is_cloned()) and not args.skip_clone
    # The checkout's voices.json is the canonical location (detect() and
    # the server launch read it there) — including when this run clones
    # the checkout itself. Without a checkout, fall back to the wav dir.
    output_path = args.output
    if output_path is None:
        if _is_cloned() or do_clone:
            output_path = _checkout() / "voices.json"
        else:
            output_path = wav_dir / "voices.json"
    if output_path.exists() and not args.force:
        print("[INFO] Aborted; existing voices.json kept")
        return None
    return {
        "do_install": do_install,
        "do_clone": do_clone,
        "wav_dir": wav_dir,
        "language": language,
        "whisper_model": args.whisper_model or "base",
        "output_path": output_path,
        "force": args.force,
        "plan": {"mode": "all", "missing": [], "existing": {}},
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Set up the faster-qwen3-tts backend: pip install, clone, "
                    "build voices.json, and sync app/converter/config.py.")
    parser.add_argument("input_dir", type=normalize_dir_arg, nargs="?",
                        default=None, metavar="WAV_DIR",
                        help="Directory with .wav reference files "
                             f"(default: {VOICES_DIR}; browsed for in the TUI)")
    parser.add_argument("--output", type=Path, default=None,
                        help="Output path for voices.json (default: "
                             "./app/faster-qwen3-tts/voices.json, or "
                             "WAV_DIR/voices.json when not cloned)")
    parser.add_argument("--language", type=str, default=None,
                        help="Language for all voices (default: English; "
                             "names and short codes accepted)")
    parser.add_argument("--whisper-model", type=str, default=None,
                        choices=WHISPER_MODELS,
                        help="Whisper model size for transcription "
                             "(default: base)")
    parser.add_argument("--force", action="store_true",
                        help="Overwrite an existing voices.json without "
                             "prompting; in the TUI, re-transcribe every "
                             "voice instead of reusing the existing file")
    parser.add_argument("--skip-install", action="store_true",
                        help="Do not pip install faster-qwen3-tts[demo] "
                             "into app/envs/faster")
    parser.add_argument("--skip-clone", action="store_true",
                        help="Do not clone the faster-qwen3-tts repo")
    return parser


def detect() -> BackendStatus:
    """Detect how far faster-qwen3-tts is set up, plus the launch command."""
    installed = _is_installed()
    cloned = _is_cloned()
    voices_json = _checkout() / "voices.json"
    configured = installed and cloned and voices_json.exists()
    details: List[str] = []
    details.append("pip: installed" if installed else
                   "not installed — run setup to pip install")
    details.append(f"checkout: {_checkout()}" if cloned else
                   f"not cloned — run setup to clone ./app/{FASTER_DIR_NAME}")
    details.append(f"voices: {voices_json}" if voices_json.exists() else
                   "no voices.json — run setup to create one")
    launch = ""
    specs: List[ServerSpec] = []
    if cloned and voices_json.exists():
        argv = [str(envs.env_python(FASTER_ENV)),
                str(_checkout() / "examples" / "openai_server.py"),
                "--voices", str(voices_json), "--port", str(_config_port())]
        # identity: /health must report model_loaded before the server is
        # really usable (the model loads after the port opens).
        specs = [ServerSpec("faster", config.FASTER_API_URL, argv,
                            identity=probe.IDENTITY_FASTER)]
        launch = format_launch_hint(specs)
    managed = servers.manages(specs)
    remote_running, remote_urls = _detect_remote(managed)
    return BackendStatus("faster", "faster-qwen3-tts",
                         installed=installed and cloned,
                         configured=configured,
                         running=managed or remote_running,
                         details=details, launch_hint=launch,
                         servers=specs, managed=managed,
                         remote=remote_running, remote_urls=remote_urls)


def _detect_remote(managed: bool = False):
    """Detect an externally-run faster server at the remote URL.

    Returns ``(running, {spec_name: url})``; see audiocpp._detect_remote for
    the shared semantics (empty URL disables, own server not counted twice).
    """
    url = (config.FASTER_REMOTE_URL or "").strip()
    if not url:
        return False, {}
    if managed and probe.same_endpoint(url, config.FASTER_API_URL):
        return False, {}
    if probe.identify_server(url) == probe.IDENTITY_FASTER:
        return True, {"faster": url}
    return False, {}


def update(*, emit=None, cancel=None) -> int:
    """Update the faster-qwen3-tts backend: pip upgrade + checkout refresh.

    A managed server that is running is stopped first (best-effort): the
    server runs ``examples/openai_server.py`` from the checkout being
    reset and imports the package being upgraded. Phases: stop server /
    pip install -U / git update — CANCEL is honored between phases only,
    so a started phase always completes. The pip package (into
    FASTER_ENV) and the cloned checkout are refreshed independently: the
    checkout only holds ``examples/openai_server.py`` (and the untracked
    voices.json, which a hard reset leaves alone), so a failed phase is
    warned about and reflected in the exit code without undoing the
    other. Returns the exit code (130 when cancelled before a remaining
    phase).
    """
    if servers.pid_for("faster") is not None:
        servers.stop("faster")
    if common.cancel_requested(cancel):
        return 130
    rc = common.pip_install([FASTER_PIP_PKG], emit=emit, cancel=cancel,
                            env_dir=FASTER_ENV, upgrade=True)
    if rc != 0:
        print(f"[WARNING] pip install -U failed (exit {rc}); update "
              f"{FASTER_PIP_PKG} manually")
    else:
        print(f"[OK] {FASTER_PIP_PKG} is up to date (or just upgraded).")
    if common.cancel_requested(cancel):
        return 130
    if _is_cloned():
        clone_rc = common.git_update(_checkout(), emit=emit, cancel=cancel)
        if clone_rc != 0:
            print(f"[WARNING] checkout update failed (exit {clone_rc}); "
                  f"run 'git -C {_checkout()} pull' manually")
            return clone_rc
        print(f"[OK] {_checkout()} is at origin's HEAD.")
    return rc


def uninstall(*, emit=None, cancel=None) -> int:
    """Remove the faster-qwen3-tts backend entirely.

    Uninstalls the pip package (``faster-qwen3-tts``) from its managed venv
    (``app/envs/faster``, FASTER_ENV — never the app env or the qwen
    backend's) and deletes the cloned checkout (``app/faster-qwen3-tts``,
    which holds examples/openai_server.py and voices.json). A running server
    this tool started is stopped first (best-effort).

    With EMIT given (the in-TUI task view) pip runs piped, streaming into
    EMIT, so its output never touches the terminal behind curses. CANCEL is
    a ``threading.Event`` honored between phases only (stop server / pip /
    delete checkout) — a started phase always completes, so pip is never
    killed mid-run. Returns the exit code (130 when cancelled before a
    remaining phase).
    """
    # Only stop when a pid file exists: without one this tool never
    # started the server, so the "not started by this tool" notice would
    # be uninstall-time noise.
    if servers.pid_for("faster") is not None:
        servers.stop("faster")
    if common.cancel_requested(cancel):
        return 130
    rc = common.pip_uninstall(["faster-qwen3-tts"], emit=emit,
                              env_dir=FASTER_ENV)
    if rc != 0:
        print("[WARNING] pip uninstall failed (exit "
              f"{rc}); remove faster-qwen3-tts from {FASTER_ENV} manually")
    else:
        print("[OK] faster-qwen3-tts removed.")
    if common.cancel_requested(cancel):
        return 130
    checkout = _checkout()
    if checkout.is_dir():
        print(f"[INFO] Removing checkout {checkout}...")
        shutil.rmtree(checkout, ignore_errors=True)
        print("[OK] checkout removed.")
    return rc


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if setup.interactive():
        return run_tui(args)

    settings = _collect_from_flags(args, parser)
    if settings is None:
        return 1
    return _execute(settings)


if __name__ == "__main__":
    sys.exit(main())