path: root/backends/common.py

"""Shared helpers for the backend setup wizards.

Every TTS backend setup wizard (audio.cpp, qwen, faster) lives in its own
module under ``backends``; this module holds the pieces more than one of
them needs: .wav discovery, path normalization, and the regex edit that
keeps ``converter/config.py`` in sync with the choices made in a wizard.
It deliberately imports nothing from the other backend modules (or the
TUI) so it can be reused without pulling curses into a non-interactive
run.
"""

import os
import re
import urllib.parse
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

# The tts-audiobook-generator checkout root (where audiobook.py lives).
# Backend checkouts are cloned into subdirectories of this root
# (./audio.cpp, ./faster-qwen3-tts) so a single tree holds everything.
TTS_ROOT = Path(__file__).resolve().parent.parent

# converter/config.py — rewritten in place by update_config_value so the
# converter picks up the host/port/voice a wizard configured.
CONFIG_PATH = TTS_ROOT / "converter" / "config.py"

# Output directory of tts-audiobook-generator; never offered as a .wav
# source by detect_wav_dir.
TTS_OUTPUT_DIR = "output"

# The voice-transcript mapping file audio.cpp reads from its voice_dir.
# (The faster backend uses voices.json instead; see backends.faster.)
PROMPT_TEXT_FILENAME = "prompt_text"


def normalize_dir_arg(value: str) -> Path:
    """Normalize a user-supplied path 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 resolve_wav_dir_arg(value: str) -> Path:
    """Normalize a user-supplied wav directory argument."""
    return normalize_dir_arg(value)


def find_wav_files(input_dir: Path) -> List[Path]:
    """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 count_wavs(directory: Path) -> int:
    """Count the .wav files in DIRECTORY (0 when it cannot be read)."""
    try:
        return sum(1 for path in directory.iterdir()
                   if path.is_file() and path.suffix.lower() == ".wav")
    except OSError:
        return 0


def detect_wav_dir(audiocpp_dir: Path, tts_root: Path) -> Optional[Path]:
    """Find a unique directory that directly contains .wav files.

    Looks shallowly (the root itself and its immediate subdirectories) in
    both the audio.cpp checkout and the tts-audiobook-generator root (where
    audiobook.py lives), since clone reference .wavs commonly live in
    either. The tts-audiobook-generator ``output/`` directory is excluded.
    When exactly one candidate is found it is returned (as a starting
    directory for the .wav browser); when none or several are found None is
    returned so the caller falls back to its default start location.
    """
    candidates: List[Path] = []
    seen: Set[Path] = set()

    def consider(directory: Path) -> None:
        try:
            resolved = directory.resolve()
        except OSError:
            return
        if resolved in seen:
            return
        seen.add(resolved)
        if count_wavs(directory) > 0:
            candidates.append(directory)

    for root in (audiocpp_dir, tts_root):
        if not root.is_dir():
            continue
        consider(root)
        try:
            children = sorted(root.iterdir(), key=lambda p: p.name.lower())
        except OSError:
            continue
        for child in children:
            if not child.is_dir() or child.name.startswith("."):
                continue
            if root == tts_root and child.name == TTS_OUTPUT_DIR:
                continue
            consider(child)

    if len(candidates) == 1:
        return candidates[0]
    return None


def wav_dir_info(directory: Path) -> Tuple[str, str]:
    """TUI status describing the directory listed in the wav browser."""
    count = count_wavs(directory)
    if count:
        wavs = ".wav" if count == 1 else ".wavs"
        return (f"{count} {wavs} found in this directory. Press Enter.",
                "ok")
    return ("No .wav files found in this directory", "warn")


def wav_dir_preview(directory: Path) -> Tuple[str, str]:
    """TUI status describing a highlighted subdirectory in the wav browser."""
    count = count_wavs(directory)
    if count:
        wavs = ".wav" if count == 1 else ".wavs"
        return (f"{count} {wavs}", "ok")
    return ("no .wav files", "info")


def url_with_port(url: str, port: int) -> str:
    """Return URL with its port replaced/inserted as PORT."""
    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_value(key: str, value: str,
                        config_path: Optional[Path] = None) -> bool:
    """Rewrite a ``KEY = "value"`` line in converter/config.py.

    Only the quoted literal is replaced; surrounding lines and the trailing
    comment are preserved. Returns True when the file was changed. Used by
    the qwen and faster wizards to keep their API URL / voice / speaker
    settings in sync with the converter.
    """
    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*' + re.escape(key) + r'\s*=\s*")([^"]*)(")',
                      text)
    if not match or match.group(2) == value:
        return False
    text = text[:match.start(2)] + value + text[match.end(2):]
    try:
        path.write_text(text, encoding="utf-8")
    except OSError:
        return False
    return True


def read_prompt_text(prompt_path: Path) -> Dict[str, str]:
    """Parse a prompt_text file into a stem -> transcript mapping.

    Lines are ``<name>|<transcript>``; blank lines are skipped and a line
    without a ``|`` separator is treated as a name with an empty transcript.
    Returns an empty mapping when the file does not exist.
    """
    if not prompt_path.exists():
        return {}
    mapping: Dict[str, str] = {}
    for line in prompt_path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        if "|" in line:
            name, _, text = line.partition("|")
        else:
            name, text = line, ""
        mapping[name.strip()] = text
    return mapping


def write_prompt_text(wav_dir: Path,
                      transcripts: Dict[str, str]) -> Path:
    """Write the voice_dir prompt_text mapping into WAV_DIR.

    One ``<basename-without-extension>|<transcript>`` line per voice.
    Returns the path of the written file.
    """
    prompt_path = wav_dir / PROMPT_TEXT_FILENAME
    lines = [f"{name}|{text}" for name, text in transcripts.items()]
    prompt_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return prompt_path


def run_console_subprocess(argv: List[str], cwd: Optional[Path] = None) -> int:
    """Run a subprocess whose output streams to the plain console.

    Used inside ``tui.suspend`` for clone/build/pip steps: the caller has
    already left curses mode, so the child inherits the real terminal and
    its output appears normally. Returns the process exit code.
    """
    import subprocess
    try:
        result = subprocess.run(argv, cwd=str(cwd) if cwd is not None else None)
    except OSError as exc:
        print(f"[ERROR] Could not run {' '.join(argv)}: {exc}")
        return 1
    return result.returncode


def git_clone(url: str, target: Path) -> int:
    """Clone URL into TARGET, streaming to the console. Returns exit code."""
    print(f"[INFO] Cloning {url} into {target}...")
    return run_console_subprocess(["git", "clone", url, str(target)])


def pip_install(packages: List[str]) -> int:
    """pip install PACKAGES (into the current environment). Returns exit code."""
    print(f"[INFO] pip install {' '.join(packages)}...")
    import sys
    return run_console_subprocess([sys.executable, "-m", "pip", "install",
                                   *packages])