#!/usr/bin/env python3
"""Interactively generate a server.json for the audio.cpp audiocpp_server.

Asks which TTS model family to host, pulls the model ids expected by this
converter (AUDIOCPP_MODEL_ID / AUDIOCPP_CLONE_MODEL_ID) from
converter/config.py, and writes a server.json that can be passed to
audiocpp_server:

    audiocpp_server --config server.json

Hostable families: Qwen3-TTS (built-in CustomVoice speakers plus voice
cloning through the Base model) and the clone-only families Higgs Audio
v3 TTS 4B, VoxCPM2-2B, and IndexTTS-2 / 2.5 (see the "Option 4" section
of the README). The converter works with other audio.cpp TTS families
too; host them by writing server.json by hand.

Reference .wav files for voice cloning (the required WAV_DIR argument)
are transcribed with a local Whisper backend (faster_whisper or whisper)
and added as voice_presets on the cloning model entry.

Every value can also be supplied as a command-line flag; anything missing
is asked interactively with the default shown in brackets. Pressing Enter
accepts the default.

Usage:
    python tools/make_audiocpp_server_json.py WAV_DIR [--output PATH]
        [--family {qwen3_tts,higgs_audio_tts,voxcpm2,index_tts2,index_tts2_5}]
        [--model-id ID] [--model-path PATH]
        [--host HOST] [--port PORT] [--models {both,custom,clone}]
        [--backend {cuda,vulkan,hip,cpu}] [--lazy-load]
        [--whisper-model NAME] [--force]

WAV_DIR is required: a directory of .wav reference files used as voice
cloning presets. It is checked up front and reported with its resolved
absolute path if it does not exist.
"""

import argparse
import json
import os
import re
import sys
import urllib.parse
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# Allow running from any working directory.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from converter import config
from converter.tts import transcribe_reference_audio, whisper_backend_available

DEFAULT_HOST = "127.0.0.1"
FALLBACK_PORT = 8080
DEFAULT_CUSTOM_VOICE_PATH = "models/Qwen3-TTS-12Hz-1.7B-CustomVoice-GGUF"
DEFAULT_BASE_PATH = "models/Qwen3-TTS-12Hz-1.7B-Base-GGUF"
CONFIG_PATH = Path(__file__).resolve().parent.parent / "converter" / "config.py"

MODEL_SELECTIONS = ("both", "custom", "clone")
BACKENDS = ("cuda", "vulkan", "hip", "cpu")

FAMILY_QWEN3_TTS = "qwen3_tts"

# Families this tool can host, in menu order. "family" is the audio.cpp
# family name written to server.json (IndexTTS-2.5 uses the index_tts2
# family; its variant is selected by the downloaded model package);
# "install" is the model_manager_v2.py package that downloads the model;
# "default_id" is the suggested server entry id; "default_path" is where
# the package lands relative to the audio.cpp checkout.
FAMILY_ENTRIES = [
    {
        "key": FAMILY_QWEN3_TTS,
        "label": "Qwen3-TTS 1.7B - built-in speakers + voice cloning",
        "family": "qwen3_tts",
    },
    {
        "key": "higgs_audio_tts",
        "label": "Higgs Audio v3 TTS 4B - voice cloning, 100+ languages",
        "family": "higgs_audio_tts",
        "install": "higgs_audio_tts_4b_q8_0",
        "default_id": "higgs",
        "default_path": "models/Higgs-Audio-v3-TTS-4B-GGUF",
    },
    {
        "key": "voxcpm2",
        "label": "VoxCPM2-2B - voice cloning, multilingual, 48 kHz audio",
        "family": "voxcpm2",
        "install": "voxcpm2_q8_0",
        "default_id": "voxcpm2",
        "default_path": "models/VoxCPM2-GGUF",
    },
    {
        "key": "index_tts2",
        "label": "IndexTTS-2 - voice cloning, Chinese/English",
        "family": "index_tts2",
        "install": "index_tts2_q8_0",
        "default_id": "indextts2",
        "default_path": "models/IndexTTS2-GGUF",
    },
    {
        "key": "index_tts2_5",
        "label": "IndexTTS-2.5 - voice cloning, zh/en/ja/es/ar",
        "family": "index_tts2",
        "install": "index_tts2_5_q8_0",
        "default_id": "indextts25",
        "default_path": "models/IndexTTS2.5-GGUF",
    },
]
FAMILY_KEYS = tuple(entry["key"] for entry in FAMILY_ENTRIES)
FAMILY_BY_KEY = {entry["key"]: entry for entry in FAMILY_ENTRIES}


def resolve_wav_dir_arg(value: str) -> Path:
    """Normalize a user-supplied wav directory argument.

    Strips surrounding quotes (a common copy-paste artifact), expands a
    leading ``~``, and resolves the result to an absolute path so relative
    paths are always validated against the current working directory.
    """
    cleaned = value.strip()
    if len(cleaned) >= 2 and cleaned[0] == cleaned[-1] and cleaned[0] in "\"'":
        cleaned = cleaned[1:-1]
    return Path(os.path.expanduser(cleaned)).resolve()


def find_wav_files(input_dir: Path) -> list:
    """Return the .wav files in INPUT_DIR, sorted alphabetically by name."""
    return sorted(
        (path for path in input_dir.iterdir()
         if path.is_file() and path.suffix.lower() == ".wav"),
        key=lambda path: path.name.lower(),
    )


def prompt_overwrite(output_path: Path) -> bool:
    """Ask whether to overwrite an existing output file."""
    while True:
        try:
            answer = input(f"{output_path} already exists. Overwrite? (y/n): ").strip().lower()
        except EOFError:
            print("\n[WARNING] No interactive input available; keeping existing file")
            return False
        if answer in ("y", "yes"):
            return True
        if answer in ("n", "no"):
            return False
        print("Please answer 'y' or 'n'.")


def ask(prompt: str, default: Optional[str] = None) -> Optional[str]:
    """Prompt for a free-text value with a default; EOF returns the default."""
    suffix = f" [{default}]" if default is not None else ""
    try:
        answer = input(f"{prompt}{suffix}: ").strip()
    except EOFError:
        return default
    return answer or default


def ask_bool(prompt: str, default: bool = False) -> bool:
    """Prompt for a yes/no answer; Enter or EOF accepts the default."""
    suffix = " [Y/n]" if default else " [y/N]"
    while True:
        try:
            answer = input(f"{prompt}{suffix}: ").strip().lower()
        except EOFError:
            return default
        if not answer:
            return default
        if answer in ("y", "yes"):
            return True
        if answer in ("n", "no"):
            return False
        print("Please answer 'y' or 'n'.")


def ask_port(default: int) -> int:
    """Prompt for a port number; Enter or EOF accepts the default."""
    while True:
        try:
            answer = input(f"Port [{default}]: ").strip()
        except EOFError:
            return default
        if not answer:
            return default
        try:
            value = int(answer)
        except ValueError:
            value = None
        if value is not None and 1 <= value <= 65535:
            return value
        print("Please enter a port number between 1 and 65535.")


def ask_menu(title: str, options: list, default_index: int = 1) -> str:
    """Show a numbered menu and return the chosen option's value."""
    print(title)
    for number, (label, _) in enumerate(options, 1):
        print(f"  {number}) {label}")
    while True:
        try:
            answer = input(f"Choice [{default_index}]: ").strip()
        except EOFError:
            return options[default_index - 1][1]
        if not answer:
            return options[default_index - 1][1]
        if answer.isdigit() and 1 <= int(answer) <= len(options):
            return options[int(answer) - 1][1]
        print(f"Please enter a number between 1 and {len(options)}.")


def ask_family() -> str:
    """Ask which model family the server should host."""
    return ask_menu(
        "Which model family should the server host?",
        [(entry["label"], entry["key"]) for entry in FAMILY_ENTRIES])


def ask_models() -> str:
    return ask_menu(
        "Which models should the server host?",
        [
            ("Both (recommended) - built-in speakers + voice cloning", "both"),
            ("CustomVoice only - built-in speakers", "custom"),
            ("Base only - voice cloning (converting then requires --voice)", "clone"),
        ])


def ask_backend() -> str:
    return ask_menu(
        "Which inference backend was audiocpp_server built for?",
        [
            ("cuda - NVIDIA GPUs (fastest)", "cuda"),
            ("vulkan - cross-vendor GPU", "vulkan"),
            ("hip - AMD GPUs", "hip"),
            ("cpu - no GPU required", "cpu"),
        ])


def ask_distinct_clone_id(primary_id: str) -> str:
    """Prompt until a non-empty id different from PRIMARY_ID is entered."""
    prompt = (f"Enter a new id for the cloning (Base) model "
              f"(must differ from '{primary_id}'): ")
    while True:
        try:
            answer = input(prompt).strip()
        except EOFError:
            print()
            sys.exit("[FATAL] No interactive input available to resolve the "
                     "duplicate model id; give the two models distinct "
                     "AUDIOCPP_MODEL_ID / AUDIOCPP_CLONE_MODEL_ID values in "
                     "converter/config.py first")
        if answer and answer != primary_id:
            return answer
        print(f"[WARNING] The id must be unique; it cannot be empty or "
              f"equal to '{primary_id}'.")


def config_port() -> int:
    """Return the port of AUDIOCPP_API_URL in converter/config.py."""
    try:
        return urllib.parse.urlsplit(config.AUDIOCPP_API_URL).port or FALLBACK_PORT
    except ValueError:
        return FALLBACK_PORT


def _url_with_port(url: str, port: int) -> str:
    parts = urllib.parse.urlsplit(url)
    host = parts.hostname or "127.0.0.1"
    return urllib.parse.urlunsplit(
        (parts.scheme or "http", f"{host}:{port}", parts.path, "", ""))


def update_config_api_url_port(port: int, config_path: Optional[Path] = None) -> bool:
    """Rewrite the port inside AUDIOCPP_API_URL in converter/config.py.

    Only the quoted URL literal is replaced; surrounding lines and the
    trailing comment are preserved. Returns True when the file was changed.
    """
    path = Path(config_path) if config_path is not None else CONFIG_PATH
    try:
        text = path.read_text(encoding="utf-8")
    except OSError:
        return False
    match = re.search(r'(?m)^(\s*AUDIOCPP_API_URL\s*=\s*")([^"]*)(")', text)
    if not match:
        return False
    new_url = _url_with_port(match.group(2), port)
    if new_url == match.group(2):
        return False
    text = text[:match.start(2)] + new_url + text[match.end(2):]
    try:
        path.write_text(text, encoding="utf-8")
    except OSError:
        return False
    return True


def update_config_model_ids(model_id: str,
                            clone_model_id: Optional[str] = None,
                            config_path: Optional[Path] = None) -> bool:
    """Rewrite AUDIOCPP_MODEL_ID (and AUDIOCPP_CLONE_MODEL_ID when given).

    Only the quoted id literals are replaced; surrounding lines and
    comments are preserved. Returns True when the file was changed.
    """
    path = Path(config_path) if config_path is not None else CONFIG_PATH
    try:
        text = path.read_text(encoding="utf-8")
    except OSError:
        return False
    updates: List[Tuple[str, str]] = [("AUDIOCPP_MODEL_ID", model_id)]
    if clone_model_id is not None:
        updates.append(("AUDIOCPP_CLONE_MODEL_ID", clone_model_id))
    changed = False
    for name, value in updates:
        match = re.search(r'(?m)^(\s*' + name + r'\s*=\s*")([^"]*)(")', text)
        if match and match.group(2) != value:
            text = text[:match.start(2)] + value + text[match.end(2):]
            changed = True
    if not changed:
        return False
    try:
        path.write_text(text, encoding="utf-8")
    except OSError:
        return False
    return True


def build_voice_presets(wav_files: list, whisper_model: str) -> Dict[str, dict]:
    """Transcribe each wav file and build the voice_presets mapping."""
    presets: Dict[str, dict] = {}
    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}'; cloning works best "
                  "with an accurate transcript — consider editing server.json "
                  "by hand before starting the server")
        presets[name] = {
            "voice_ref": str(wav_file.resolve()),
            "reference_text": text or "",
        }
    return presets


def build_server_config(host: str, port: int, backend: str, lazy_load: bool,
                        include_custom: bool, include_clone: bool,
                        custom_voice_id: str, clone_model_id: str,
                        custom_voice_path: str, base_path: str,
                        voice_presets: Dict[str, dict]) -> dict:
    """Assemble the Qwen3-TTS server.json document."""
    models = []
    if include_custom:
        models.append({
            "id": custom_voice_id,
            "family": "qwen3_tts",
            "path": custom_voice_path,
            "task": "tts",
            "mode": "offline",
        })
    if include_clone:
        clone_entry = {
            "id": clone_model_id,
            "family": "qwen3_tts",
            "path": base_path,
            "task": "tts",
            "mode": "offline",
        }
        if voice_presets:
            clone_entry["voice_presets"] = voice_presets
        models.append(clone_entry)
    return {
        "host": host,
        "port": port,
        "backend": backend,
        "lazy_load": lazy_load,
        "models": models,
    }


def build_single_family_server_config(host: str, port: int, backend: str,
                                      lazy_load: bool, family: str,
                                      model_id: str, model_path: str,
                                      voice_presets: Dict[str, dict]) -> dict:
    """Assemble a server.json hosting one clone-only model family entry."""
    entry = {
        "id": model_id,
        "family": family,
        "path": model_path,
        "task": "tts",
        "mode": "offline",
    }
    if voice_presets:
        entry["voice_presets"] = voice_presets
    return {
        "host": host,
        "port": port,
        "backend": backend,
        "lazy_load": lazy_load,
        "models": [entry],
    }


def print_empty_transcript_warning(voice_presets: Dict[str, dict]) -> None:
    """Print a loud, final warning for voices whose transcript is empty."""
    empty = sorted(name for name, preset in voice_presets.items()
                   if not preset.get("reference_text"))
    if not empty:
        return
    bar = "=" * 70
    print()
    print(bar)
    print("[WARNING] MANUAL TRANSCRIPTION REQUIRED")
    print(bar)
    listing = "  - " + "\n  - ".join(empty) if len(empty) > 1 else f"  - {empty[0]}"
    print(f"The following voice preset(s) have an EMPTY reference_text in "
          f"server.json:\n{listing}")
    print("Those voices will NOT work until you add a manual transcription.")
    print('Edit server.json and fill in the "reference_text" field for each '
          "voice above with an accurate transcript of its reference .wav.")
    print(bar)


def _ask_host_port_backend_lazy(args: argparse.Namespace
                                ) -> Tuple[str, int, str, bool]:
    """Ask for (or take from flags) the shared server settings."""
    host = args.host if args.host else ask("Bind host", DEFAULT_HOST)
    port = args.port if args.port is not None else ask_port(config_port())
    if port != config_port():
        if ask_bool(f"Update AUDIOCPP_API_URL in converter/config.py to port "
                    f"{port} so audiobook.py talks to this server", True):
            if update_config_api_url_port(port):
                print(f"[OK] Updated AUDIOCPP_API_URL in {CONFIG_PATH}")
            else:
                print(f"[WARNING] Could not update {CONFIG_PATH}; edit "
                      "AUDIOCPP_API_URL by hand so audiobook.py uses the "
                      "new port")
        else:
            print("[WARNING] Left AUDIOCPP_API_URL unchanged; audiobook.py "
                  f"will still use port {config_port()}")
    backend = args.backend if args.backend else ask_backend()
    lazy_load = args.lazy_load or ask_bool(
        "Load models lazily (on first use instead of at startup)", False)
    return host, port, backend, lazy_load


def _collect_voice_presets(args: argparse.Namespace,
                           include_clone: bool) -> Dict[str, dict]:
    """Transcribe the wav directory into the voice_presets mapping.

    Returns the voice_presets mapping (empty when no wavs were found).
    Cloning entries only: a run without any cloning model ignores the wav
    directory entirely.
    """
    if not include_clone:
        print(f"[WARNING] Ignoring {args.input_dir}: no cloning model "
              "selected, so voice presets are not used")
        return {}

    wav_files = find_wav_files(args.input_dir)
    if not wav_files:
        print(f"[WARNING] No .wav files found in {args.input_dir}; writing the "
              "config without voice presets")
        return {}
    if whisper_backend_available() is None:
        print("[WARNING] Neither faster_whisper nor whisper was found, so "
              "reference .wav files cannot be transcribed automatically and "
              "every reference_text will be empty.")
        print('          Did you remember to "conda activate qwen3-tts"? '
              "Transcripts must be added by hand (see the warning at the end).")
    return build_voice_presets(wav_files, args.whisper_model)


def _offer_config_model_id_sync(model_id: str) -> None:
    """Offer to point converter/config.py at a non-Qwen model entry.

    The converter requests the model id configured in AUDIOCPP_MODEL_ID,
    and single-model servers use the same id for the clone entry, so both
    ids are rewritten together.
    """
    if config.AUDIOCPP_MODEL_ID == model_id \
            and config.AUDIOCPP_CLONE_MODEL_ID == model_id:
        return
    if ask_bool("Update AUDIOCPP_MODEL_ID and AUDIOCPP_CLONE_MODEL_ID in "
                f"converter/config.py to '{model_id}' so audiobook.py uses "
                "this model", True):
        if update_config_model_ids(model_id, model_id):
            print(f"[OK] Updated the model ids in {CONFIG_PATH}")
        else:
            print(f"[WARNING] Could not update {CONFIG_PATH}; edit "
                  "AUDIOCPP_MODEL_ID and AUDIOCPP_CLONE_MODEL_ID by hand so "
                  "audiobook.py uses this model")
    else:
        print("[WARNING] Left the model ids unchanged; audiobook.py will "
              f"still request model '{config.AUDIOCPP_MODEL_ID}'")


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Generate a server.json for the audio.cpp audiocpp_server "
                    "hosting a TTS model used by this converter.")
    parser.add_argument("input_dir", type=resolve_wav_dir_arg, metavar="WAV_DIR",
                        help="Directory with .wav reference files to add as "
                             "voice cloning presets (required)")
    parser.add_argument("--output", type=Path, default=Path("server.json"),
                        help="Output path for server.json (default: "
                             "server.json in the current directory)")
    parser.add_argument("--family", choices=FAMILY_KEYS, default=None,
                        help="Model family to host (default: Qwen3-TTS). "
                             "Non-Qwen families are clone-only and host a "
                             "single model entry")
    parser.add_argument("--model-id", type=str, default=None,
                        help="Server model id for a non-Qwen family entry "
                             "(default: a family-based name such as 'higgs')")
    parser.add_argument("--model-path", type=str, default=None,
                        help="Path to a non-Qwen family model package "
                             "(default: the model manager install location)")
    parser.add_argument("--host", type=str, default=None,
                        help="Bind host for the server (default: 127.0.0.1)")
    parser.add_argument("--port", type=int, default=None,
                        help="Port for the server (default: the port in "
                             "AUDIOCPP_API_URL from converter/config.py)")
    parser.add_argument("--models", choices=MODEL_SELECTIONS, default=None,
                        help="Which Qwen3-TTS models to host: both (default), "
                             "custom (CustomVoice speakers only), or clone "
                             "(Base voice cloning only). Only valid with "
                             "--family qwen3_tts")
    parser.add_argument("--backend", choices=BACKENDS, default=None,
                        help="Inference backend audiocpp_server was built "
                             "for (default: cuda)")
    parser.add_argument("--lazy-load", action="store_true",
                        help="Load models on first use instead of at startup "
                             "(default: load at startup)")
    parser.add_argument("--whisper-model", type=str, default="base",
                        help="Whisper model size for transcription "
                             "(default: base)")
    parser.add_argument("--force", action="store_true",
                        help="Overwrite the output file without prompting")
    args = parser.parse_args()

    if not args.input_dir.is_dir():
        parser.error(
            f"WAV directory not found: {args.input_dir}\n"
            f"      (resolved from the current working directory: "
            f"{Path.cwd()})\n"
            "      WAV_DIR must be a directory containing the .wav "
            "reference files to use as voice cloning presets")

    if args.output.exists() and not args.force \
            and not prompt_overwrite(args.output):
        print("[INFO] Aborted; existing server.json kept")
        return 1

    family_key = args.family if args.family is not None else ask_family()
    is_qwen = family_key == FAMILY_QWEN3_TTS
    if not is_qwen and args.models is not None:
        parser.error("--models only applies to --family qwen3_tts")

    print("[INFO] Model ids from converter/config.py:")
    print(f"       built-in speakers (CustomVoice): '{config.AUDIOCPP_MODEL_ID}'")
    print(f"       voice cloning (Base):            '{config.AUDIOCPP_CLONE_MODEL_ID}'")

    if is_qwen:
        selection = args.models if args.models is not None else ask_models()
        include_custom = selection in ("both", "custom")
        include_clone = selection in ("both", "clone")

        custom_voice_id = config.AUDIOCPP_MODEL_ID
        clone_model_id = config.AUDIOCPP_CLONE_MODEL_ID
        if include_custom and include_clone and custom_voice_id == clone_model_id:
            print(f"[WARNING] AUDIOCPP_MODEL_ID and AUDIOCPP_CLONE_MODEL_ID are "
                  f"both '{custom_voice_id}' in converter/config.py, but server "
                  "model ids must be unique.")
            clone_model_id = ask_distinct_clone_id(custom_voice_id)
    else:
        entry = FAMILY_BY_KEY[family_key]
        include_custom = False
        include_clone = True
        model_id = args.model_id if args.model_id else ask(
            f"Server model id for the {entry['label']} entry",
            entry["default_id"])
        _offer_config_model_id_sync(model_id)

    host, port, backend, lazy_load = _ask_host_port_backend_lazy(args)

    if is_qwen:
        custom_voice_path = base_path = None
        if include_custom:
            custom_voice_path = ask("Path to the Qwen3-TTS CustomVoice GGUF package",
                                    DEFAULT_CUSTOM_VOICE_PATH)
        if include_clone:
            base_path = ask("Path to the Qwen3-TTS Base GGUF package",
                            DEFAULT_BASE_PATH)
    else:
        model_path = args.model_path if args.model_path else ask(
            f"Path to the {entry['label']} package", entry["default_path"])

    voice_presets = _collect_voice_presets(args, include_clone)

    if is_qwen:
        server_config = build_server_config(
            host=host,
            port=port,
            backend=backend,
            lazy_load=lazy_load,
            include_custom=include_custom,
            include_clone=include_clone,
            custom_voice_id=custom_voice_id,
            clone_model_id=clone_model_id,
            custom_voice_path=custom_voice_path,
            base_path=base_path,
            voice_presets=voice_presets,
        )
    else:
        server_config = build_single_family_server_config(
            host=host,
            port=port,
            backend=backend,
            lazy_load=lazy_load,
            family=entry["family"],
            model_id=model_id,
            model_path=model_path,
            voice_presets=voice_presets,
        )

    print("\nGenerated server.json:")
    print(json.dumps(server_config, indent=2, ensure_ascii=False))
    if not ask_bool(f"\nWrite this to {args.output}", True):
        print("[INFO] Aborted; nothing written")
        return 1

    with args.output.open("w", encoding="utf-8") as handle:
        json.dump(server_config, handle, indent=2, ensure_ascii=False)
        handle.write("\n")

    if is_qwen:
        print(f"\n[OK] Wrote {args.output} with {len(server_config['models'])} "
              f"model(s) and {len(voice_presets)} voice preset(s)")
    else:
        print(f"\n[OK] Wrote {args.output} hosting {entry['label']} "
              f"(model id '{model_id}') with {len(voice_presets)} "
              f"voice preset(s)")
        print(f"[INFO] Install the model package from the audio.cpp checkout: "
              f"python3 tools/model_manager_v2.py install {entry['install']}")
        print("[INFO] Clone-only family: run audiobook.py with "
              f"--backend audiocpp --voice <preset name>")
        if not voice_presets:
            print("[WARNING] No voice presets were configured; clone-only "
                  "families have no built-in speakers, so add voice_presets "
                  "(or a voice_dir) to server.json before converting")
    print_empty_transcript_warning(voice_presets)
    return 0


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