#!/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, 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]
        [--port PORT] [--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 and pre-fills the language and
wav directory from them instead of prompting to overwrite, asks whether
to only transcribe new voices or re-transcribe everything, 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,
)
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]"
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")


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 _decide_faster_transcription(wav_files: list, existing_voices: dict,
                                 confirm) -> Optional[dict]:
    """Decide which voices to transcribe when a voices.json already exists.

    CONFIRM asks the yes/no question (returning True/False, or None when the
    user backs out). With new .wavs present it offers to transcribe only
    those (default Yes); otherwise — and always, per the modify design — it
    offers to re-transcribe everything (default No), so a stale transcript
    can be refreshed even when every voice is already known. Returns a plan
    dict: ``{"mode": "missing"|"all"|"keep", "missing": [...], "existing":
    {...}}``, or None when CONFIRM cancelled.
    """
    existing = dict(existing_voices)
    new_wavs = [wav for wav in wav_files if wav.stem not in existing]
    if new_wavs:
        choice = confirm("Existing voices.json found. Only transcribe the "
                         "new voices?", True)
        if choice is None:
            return None
        if choice:
            return {"mode": "missing", "missing": new_wavs,
                    "existing": existing}
        return {"mode": "all", "missing": [], "existing": existing}
    choice = confirm("All voices already in voices.json. Re-transcribe "
                     "anyway?", False)
    if choice is None:
        return None
    if choice:
        return {"mode": "all", "missing": [], "existing": existing}
    return {"mode": "keep", "missing": [], "existing": existing}


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 (built by ``_decide_faster_transcription`` in the wizard, 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 _wizard(stdscr, args: argparse.Namespace) -> Optional[dict]:
    """Linear TUI wizard collecting every faster-setup decision.

    Driven by ``tui.Wizard`` as a stack of screen closures: each screen
    shows one widget and returns the next screen, ``Wizard.BACK`` (Esc/q —
    pop to the previous screen), or the settings dict. Steps whose value is
    already provided by a flag (``--wavs``, ``--language``,
    ``--whisper-model``, ``--port``, ``--skip-install``, ``--skip-clone``)
    or that do not apply (the transcription plan when there is nothing to
    decide) are folded into the ``_after_*`` guards and never become
    screens, so Esc always lands on the previous real screen. Esc on the
    first screen aborts the wizard.
    """
    _GO_BACK = object()
    s: dict = {}

    def _confirm(question: str, default: bool = True) -> Optional[bool]:
        res = tui.confirm(stdscr, question, default=default,
                          cancel_value=_GO_BACK)
        return None if res is _GO_BACK else res

    # An existing voices.json seeds the defaults (modify flow) instead of an
    # overwrite prompt; its voices also seed the wav-directory browser.
    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)
    s["default_output"] = default_output
    s["existing_voices"] = existing_voices
    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))
    s["wav_start"] = wav_start

    # Install and clone happen without asking: when the package or repo is
    # missing (and not skipped by flag), the wizard just does it and moves
    # to the next screen.
    s["do_install"] = (not _is_installed()) and not args.skip_install
    s["do_clone"] = (not _is_cloned()) and not args.skip_clone

    def _after_clone():
        if args.input_dir is None:
            return screen_wav
        s["wav_dir"] = args.input_dir
        return _after_wav()

    def screen_wav():
        wav_dir = tui.browse_directory(
            stdscr, "Select the directory with your .wav voices",
            info=common.wav_dir_info, preview=common.wav_dir_preview,
            start=s["wav_start"], back_value=_GO_BACK)
        if wav_dir is _GO_BACK:
            return tui.Wizard.BACK
        s["wav_dir"] = wav_dir
        return _after_wav()

    def _after_wav():
        if args.language is None:
            return screen_language
        s["language"] = args.language
        return _after_language()

    def screen_language():
        default_language = config.LANGUAGE
        for voice in s["existing_voices"].values():
            if isinstance(voice, dict) and voice.get("language"):
                default_language = voice["language"]
                break
        lang_text = tui.line_edit(
            stdscr, "Language", default_language,
            validate=lambda s: None if _try_language(s)
            else "Unknown language (e.g. English, en)",
            help_lines=["Language for every voice, as passed to the TTS "
                        "model (names or short codes accepted)"],
            back_value=_GO_BACK)
        if lang_text is _GO_BACK:
            return tui.Wizard.BACK
        s["language"] = lang_text
        return _after_language()

    def _after_language():
        if args.whisper_model is None:
            return screen_whisper
        s["whisper_model"] = args.whisper_model
        return _after_whisper()

    def screen_whisper():
        whisper_model = tui.menu(
            stdscr, "Whisper model for transcription",
            [(m, m) for m in WHISPER_MODELS],
            default_index=WHISPER_MODELS.index("base"),
            back_value=_GO_BACK)
        if whisper_model is _GO_BACK:
            return tui.Wizard.BACK
        s["whisper_model"] = whisper_model
        return _after_whisper()

    def _after_whisper():
        # 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 flag-only runs working without any checkout.
        s["output_path"] = args.output
        if s["output_path"] is None:
            if _is_cloned() or s.get("do_clone"):
                s["output_path"] = _checkout() / "voices.json"
            else:
                s["output_path"] = s["wav_dir"] / "voices.json"
        wav_files = find_wav_files(s["wav_dir"])
        if wav_files and s["existing_voices"] and not args.force:
            return screen_transcription
        s["plan"] = {"mode": "all", "missing": [], "existing": {}}
        return _after_transcription()

    def screen_transcription():
        # Re-transcribe only new voices (or all of them) — the
        # "re-transcribe anyway?" offer appears even when nothing is new.
        wav_files = find_wav_files(s["wav_dir"])
        plan = _decide_faster_transcription(
            wav_files, s["existing_voices"], _confirm)
        if plan is None:
            return tui.Wizard.BACK
        s["plan"] = plan
        return _after_transcription()

    def _after_transcription():
        if args.port is None:
            return screen_port
        s["port"] = args.port
        return _finalize()

    def screen_port():
        port_text = tui.line_edit(
            stdscr, "Server port", str(_config_port()),
            validate=lambda s: None if (s.isdigit() and 1 <= int(s) <= 65535)
            else "Enter a port number between 1 and 65535",
            back_value=_GO_BACK)
        if port_text is _GO_BACK:
            return tui.Wizard.BACK
        s["port"] = int(port_text)
        return _finalize()

    def _finalize() -> dict:
        return {
            "do_install": s.get("do_install", False),
            "do_clone": s.get("do_clone", False),
            "wav_dir": s["wav_dir"],
            "language": s["language"],
            "whisper_model": s["whisper_model"],
            "output_path": s["output_path"],
            "port": s["port"],
            "force": args.force,
            "plan": s["plan"],
        }

    return tui.Wizard().run(_after_clone())


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)
            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 port + default voice.
        port = settings["port"]
        new_url = common.url_with_port(config.FASTER_API_URL, port)
        if new_url != config.FASTER_API_URL:
            if common.update_config_value("FASTER_API_URL", new_url):
                print(f"[OK] Updated FASTER_API_URL to {new_url}")
            else:
                print("[WARNING] Could not update FASTER_API_URL; edit "
                      "app/converter/config.py by hand")
        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"], port)
        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, port: int) -> 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} --port {port}")


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

    The hub drives this as one screen of its own ``tui.Wizard`` stack, so
    Esc on the wizard's first screen simply returns here and the hub pops
    back to the menu that launched it. The setup tail (install/clone/
    transcribe/write) runs inside the TUI task view on this same screen, so
    the hub's curses session stays intact and the user sees per-step status
    instead of being dropped to the console. Returns 0 on completion, 1 when
    the user aborted.
    """
    args = build_parser().parse_args([])
    settings = _wizard(stdscr, args)
    if settings is None:
        return 1
    return taskview.run_steps(stdscr, "Setting up faster-qwen3-tts",
                              _execute_steps(settings))


def run_tui(args: Optional[argparse.Namespace] = None) -> int:
    """Run the faster setup wizard end-to-end."""
    import curses
    if args is None:
        args = build_parser().parse_args([])
    try:
        settings = curses.wrapper(_wizard, args)
    except tui.WizardCancelled:
        print("\n[INFO] Cancelled; nothing was written")
        return 1
    try:
        curses.curs_set(1)
    except curses.error:
        pass
    if settings is None:
        print("[INFO] Aborted")
        return 1
    return _execute(settings)


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,
        "port": args.port if args.port is not None else _config_port(),
        "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("--port", type=int, default=None,
                        help="Server port to record in app/converter/config.py "
                             "(default: the port in FASTER_API_URL)")
    parser.add_argument("--skip-install", action="store_true",
                        help="Do not pip install faster-qwen3-tts[demo]")
    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()),
                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 uninstall(*, emit=None, cancel=None) -> int:
    """Remove the faster-qwen3-tts backend entirely.

    Uninstalls the pip package (``faster-qwen3-tts``) from the managed venv
    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)
    if rc != 0:
        print("[WARNING] pip uninstall failed (exit "
              f"{rc}); remove faster-qwen3-tts from the managed venv 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 _interactive():
        return run_tui(args)

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


def _interactive() -> bool:
    try:
        import curses  # noqa: F401
    except ImportError:
        return False
    try:
        return sys.stdin.isatty() and sys.stdout.isatty()
    except (AttributeError, ValueError):
        return False


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