#!/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 (a directory argument or an
interactive prompt) 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, so running the tool with no arguments and pressing
Enter through produces a server.json hosting both Qwen3-TTS models on
127.0.0.1:8080 with the cuda backend.

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]
"""

import argparse
import json
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 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 ask_wav_dir() -> Optional[Path]:
    """Prompt for a directory of .wav clone references; Enter skips."""
    while True:
        try:
            answer = input("Directory with .wav files to clone "
                           "(Enter to skip): ").strip()
        except EOFError:
            return None
        if not answer:
            return None
        path = Path(answer)
        if path.is_dir():
            return path
        print(f"[WARNING] {answer} is not a directory; try again "
              "(or press Enter to skip).")


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]:
    """Resolve the clone-reference wav directory and transcribe it.

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

    wav_files = find_wav_files(wav_dir)
    if not wav_files:
        print(f"[WARNING] No .wav files found in {wav_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=Path, nargs="?", default=None,
                        help="Optional directory with .wav reference files "
                             "to add as voice cloning presets")
    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 args.input_dir is not None and not args.input_dir.is_dir():
        parser.error(f"WAV directory not found: {args.input_dir}")

    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())