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

Asks which Qwen3-TTS models 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

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 Base-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 models on
127.0.0.1:8080 with the cuda backend.

Usage:
    python tools/make_audiocpp_server_json.py [WAV_DIR] [--output 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, Optional

# 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

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


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_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 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 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 _print_next_steps(output_path: Path, include_custom: bool,
                      include_clone: bool, voice_presets: Dict[str, dict]) -> None:
    print("\nNext steps:")
    print("  1. Start the server (build path varies by platform, e.g.")
    print("     ./build/linux-cuda-release/bin/):")
    print(f"       audiocpp_server --config {output_path}")
    print("  2. Convert a book from this repository:")
    if include_custom:
        print("       python audiobook.py --backend audiocpp"
              "                # built-in speaker")
    if include_clone:
        names = ", ".join(voice_presets) or "none configured yet"
        print("       python audiobook.py --backend audiocpp --voice NAME"
              f"   # cloned voice ({names})")
    if include_clone and not include_custom:
        print("[INFO] Only the Base model is hosted: --voice is required, "
              "since speaker mode needs the CustomVoice model.")


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Generate a server.json for the audio.cpp audiocpp_server "
                    "hosting the Qwen3-TTS models 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("--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 models to host: both (default), custom "
                             "(CustomVoice speakers only), or clone "
                             "(Base voice cloning only)")
    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

    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}'")

    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)

    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)

    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)

    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 (Base) "
                  "model selected, so voice presets are not used")
    elif include_clone:
        wav_dir = ask_wav_dir()

    voice_presets: Dict[str, dict] = {}
    if wav_dir is not None:
        wav_files = find_wav_files(wav_dir)
        if wav_files:
            voice_presets = build_voice_presets(wav_files, args.whisper_model)
        else:
            print(f"[WARNING] No .wav files found in {wav_dir}; writing the "
                  "config without voice presets")

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

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

    print(f"\n[OK] Wrote {args.output} with {len(server_config['models'])} "
          f"model(s) and {len(voice_presets)} voice preset(s)")
    _print_next_steps(args.output, include_custom, include_clone, voice_presets)
    return 0


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