path: root/app/backends/envs.py

"""The managed Python environments for the audiobook generator and its backends.

audiobook.py is meant to be launched from any Python (a bare system interpreter
is fine): on startup it bootstraps a single tool-managed app venv at
``app/envs/tts`` and re-execs itself inside it. That venv holds the audiobook
app's own ``requirements.txt`` dependencies only. The backend TTS packages the
setup wizards pip install get dedicated venvs of their own —
``app/envs/qwen`` for ``qwen-tts``, ``app/envs/faster`` for
``faster-qwen3-tts[demo]`` — so their heavyweight torch/transformers stacks
never share an environment with each other or with the app (qwen-tts and
faster-qwen3-tts both ship a ``qwen_tts`` module with conflicting
transformers pins, which makes a shared install break whichever distribution
lands last). Nothing is ever installed into the launching interpreter's
environment.

A parent process never needs to "activate" an environment — activation is
just a shell convenience that puts an env's ``bin`` on PATH. Instead every
helper here resolves an env's binaries by absolute path
(``app/envs/tts/bin/python``, ``app/envs/qwen/bin/qwen-tts-demo``), so the hub can
spawn servers in these envs from any parent environment.

This module is imported before audiobook.py's third-party dependencies, so
it must stay stdlib-only (it may import ``backends.common``, which is also
stdlib-only, but never ``converter`` or the backend modules).
"""

import hashlib
import json
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from backends import common

# The tts-audiobook-generator checkout root (where audiobook.py lives).
TTS_ROOT = Path(__file__).resolve().parent.parent.parent

# One app venv for the app requirements; one venv per pip-installed TTS
# backend. The per-backend split keeps qwen-tts's transformers 4 pin away
# from faster-qwen3-tts's transformers 5 requirement (and away from the app).
ENV_DIR = TTS_ROOT / "app" / "envs" / "tts"
QWEN_ENV_DIR = TTS_ROOT / "app" / "envs" / "qwen"
FASTER_ENV_DIR = TTS_ROOT / "app" / "envs" / "faster"
REQUIREMENTS_PATH = TTS_ROOT / "requirements.txt"

# requirements.txt lines whose comment starts with this tag are installed
# best-effort: they gate features that degrade gracefully at runtime (e.g.
# faster-whisper falls back to x-vector-only cloning), so on platforms with
# no compatible wheels (ctranslate2 has no musllinux builds) the install
# retries without them instead of failing the whole bootstrap.
OPTIONAL_TAG = "# optional:"

# Marker file recording what was last installed into the env, so
# ensure_app_env() re-installs when requirements.txt changes. Its content is
# "<requirements sha256>:<MARKER_VERSION>"; bump MARKER_VERSION whenever
# ensure_app_env gains a new post-install obligation, so envs installed by
# older tool versions are re-installed (and re-verified) once on next launch.
MARKER_PATH = ENV_DIR / ".audiobook_env_ready"
MARKER_VERSION = "2"


def _is_windows() -> bool:
    return sys.platform == "win32"


def env_python(env_dir: Optional[Path] = None) -> Path:
    """Absolute path to an env's python interpreter (the app env by default)."""
    base = env_dir if env_dir is not None else ENV_DIR
    return base / ("Scripts/python.exe" if _is_windows() else "bin/python")


def env_script(name: str, env_dir: Optional[Path] = None) -> Path:
    """Absolute path to a console script installed in an env (e.g. qwen-tts-demo).

    ENV_DIR defaults to the app env; pass a backend's env dir (QWEN_ENV_DIR,
    FASTER_ENV_DIR) for its scripts.
    """
    subdir = "Scripts" if _is_windows() else "bin"
    suffix = ".exe" if _is_windows() else ""
    return (env_dir if env_dir is not None else ENV_DIR) / subdir / f"{name}{suffix}"


def env_exists(env_dir: Optional[Path] = None) -> bool:
    """True when the given env's python interpreter is present on disk."""
    return env_python(env_dir).is_file()


def is_managed_env() -> bool:
    """True when the current process is already running inside the app venv."""
    try:
        return Path(sys.executable).resolve() == env_python().resolve()
    except OSError:
        return False


def create_env(env_dir: Optional[Path] = None) -> int:
    """Create ENV (an env dir, defaulting to the app one) with the launching
    interpreter (inherits its version).

    pip is bootstrapped inside the venv by ensurepip. Returns the ``python -m
    venv`` exit code; a non-zero result is reported with platform remediation.
    """
    target = env_dir if env_dir is not None else ENV_DIR
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("  Setting up your environment for the first time...")
    print("  This may take a minute.")
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print(f"[INFO] creating managed environment at {target}...")
    rc = common.run_console_subprocess(
        [sys.executable, "-m", "venv", str(target)])
    if rc != 0:
        print(f"[ERROR] python -m venv failed (exit {rc}).")
        if _is_windows():
            print("        On Windows make sure the launcher has the venv module.")
        else:
            print("        On Debian/Ubuntu install the venv package, e.g.:")
            print("          sudo apt install python3-venv")
    return rc


def _marker_applies(marker: str) -> bool:
    """Best-effort evaluation of a requirements.txt environment marker.

    Only the ``sys_platform == "win32"`` gate is interpreted (the one form
    this project uses): it applies everywhere except non-Windows hosts,
    where the entry must not be installed *or* probed. Any other marker is
    assumed to apply.
    """
    if not marker.strip():
        return True
    return not ("win32" in marker and not _is_windows())


def requirement_specs() -> List[Tuple[str, bool]]:
    """Parse requirements.txt into ``(spec, is_optional)`` pairs.

    SPEC is the pip requirement (e.g. ``faster-whisper>=1.0.0``), with
    comments stripped and environment markers evaluated best-effort by
    _marker_applies (entries excluded by their marker are left out here so
    neither the install nor the import probes see them). A line is
    optional when its comment starts with OPTIONAL_TAG. Option/flag lines
    (``-r``, ``--index-url``, ...) are ignored — this file holds plain
    requirement lines only.
    """
    try:
        lines = REQUIREMENTS_PATH.read_text(encoding="utf-8").splitlines()
    except OSError:
        return []
    specs: List[Tuple[str, bool]] = []
    for line in lines:
        req, _, comment = line.partition("#")
        optional = comment.strip().lower().startswith(OPTIONAL_TAG[2:])
        req, _, marker = req.partition(";")
        if not _marker_applies(marker):
            continue
        req = req.strip().rstrip("\\").strip()
        if not req or req.startswith(("-", "--")):
            continue
        specs.append((req, optional))
    return specs


def _base_name(spec: str) -> str:
    """The distribution name portion of a pip requirement spec."""
    return re.split(r"[<>=!~;\[ ]", spec, maxsplit=1)[0].strip()


def install_requirements(skip_optional: bool = False) -> int:
    """Install requirements.txt into the venv. Returns pip's exit code.

    With SKIP_OPTIONAL the lines tagged OPTIONAL_TAG are left out — the
    fallback for platforms where an optional dependency cannot resolve
    (see ensure_app_env, which retries core-only before giving up).
    """
    specs = requirement_specs()
    installable = [spec for spec, optional in specs
                   if not (skip_optional and optional)]
    if not installable:
        print(f"[ERROR] no installable requirements found in {REQUIREMENTS_PATH}")
        return 1
    skipped = [spec for spec, optional in specs if optional]
    label = str(REQUIREMENTS_PATH) if not skip_optional else \
        f"{REQUIREMENTS_PATH} (without optional: {', '.join(_base_name(s) for s in skipped)})"
    print(f"[INFO] pip install {label} into {ENV_DIR}...")
    return common.run_console_subprocess(
        [str(env_python()), "-m", "pip", "install", *installable])


def pip_install(packages: List[str], *, emit=None, cancel=None,
                env_dir: Optional[Path] = None,
                upgrade: bool = False) -> int:
    """pip install PACKAGES into ENV (an env dir, default the app env),
    creating it first if needed.

    Used by the qwen/faster setup wizards to install their TTS packages into
    their dedicated backend venvs (QWEN_ENV_DIR / FASTER_ENV_DIR), never
    alongside each other or the app requirements. With UPGRADE the install
    runs with ``-U``: pip then resolves the latest version itself and
    reports "Requirement already satisfied" when the env already holds it —
    the backend update action's cheap freshness check. Returns pip's exit
    code. With EMIT given (the in-TUI task view) pip runs with
    ``--progress-bar off`` so its output is clean status lines rather than
    carriage-return progress spam.
    """
    if not env_exists(env_dir) and create_env(env_dir) != 0:
        return 1
    target = env_dir if env_dir is not None else ENV_DIR
    print(f"[INFO] pip install {' '.join(packages)} into {target}...")
    argv = [str(env_python(env_dir)), "-m", "pip", "install"]
    if upgrade:
        argv.append("-U")
    if emit is not None:
        argv.append("--progress-bar")
        argv.append("off")
    argv.extend(packages)
    return common.run_console_subprocess(argv, emit=emit, cancel=cancel)


def pip_uninstall(packages: List[str], *, emit=None,
                  env_dir: Optional[Path] = None) -> int:
    """pip uninstall PACKAGES from ENV (an env dir, default the app env).
    Returns pip's exit code.

    Used by the backends' ``uninstall`` action to remove pip-installed TTS
    packages from their dedicated environments. A missing env is a no-op
    (there is nothing to uninstall from), reported as success. With EMIT
    given (the in-TUI task view) pip runs with its output piped and streamed
    to EMIT, so nothing writes to the terminal behind curses.
    """
    if not env_exists(env_dir):
        return 0
    target = env_dir if env_dir is not None else ENV_DIR
    print(f"[INFO] pip uninstall {' '.join(packages)} from {target}...")
    return common.run_console_subprocess(
        [str(env_python(env_dir)), "-m", "pip", "uninstall", "-y", *packages],
        emit=emit)


def module_available(module: str, env_dir: Optional[Path] = None) -> bool:
    """True when MODULE imports inside ENV (e.g. qwen_tts in QWEN_ENV_DIR).

    A short subprocess probe against the env's interpreter — the equivalent
    of importlib.util.find_spec, but for a managed env rather than the
    current one. Used by each backend's ``_is_installed``.
    """
    if not env_exists(env_dir):
        return False
    try:
        result = subprocess.run(
            [str(env_python(env_dir)), "-c", f"import {module}"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            timeout=30, check=False)
    except (OSError, subprocess.TimeoutExpired):
        return False
    return result.returncode == 0


# Import names that differ from their requirements.txt distribution name.
_IMPORT_NAMES = {
    "beautifulsoup4": "bs4",
    "faster-whisper": "faster_whisper",
}


def _imports_ok(import_names: List[str], *, python: Optional[Path] = None) -> bool:
    """True when every name in IMPORT_NAMES imports inside the target env.

    One combined probe subprocess: a healthy env costs a single interpreter
    start-up; only failures are isolated per-name afterwards.
    """
    python = python or env_python()
    try:
        result = subprocess.run(
            [str(python), "-c", f"import {', '.join(import_names)}"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            timeout=300, check=False)
    except (OSError, subprocess.TimeoutExpired):
        return False
    return result.returncode == 0


# Runs INSIDE the target env: collect one candidate import name per installed
# distribution (its top_level.txt, falling back to the normalized project
# name) and actually import each, reporting the failures. Two subtleties the
# wrong-platform-wheel case demands: find_spec gates the heuristic fallback
# names (a distribution like protobuf exposes no top-level "protobuf" module,
# and guessing must not turn that into a false positive), and distributions
# shipping compiled extensions get their extension submodules probed too —
# lxml's pure-Python __init__ imports fine while every binary submodule is
# missing.
_ENV_SCAN_CODE = """
import importlib.metadata as md
import importlib.util as iu
import json
import os

names = set()
for dist in md.distributions():
    top = dist.read_text("top_level.txt")
    if top:
        names.update(part.strip() for part in top.split())
    else:
        name = (dist.metadata.get("Name") or "").strip()
        if name:
            names.add(name.lower().replace("-", "_"))


def compiled_submodules(top):
    \"\"\"top.* extension modules worth probing (only for .so-shipping tops).\"\"\"
    try:
        spec = iu.find_spec(top)
        if spec is None or not spec.submodule_search_locations:
            return []
        mods = []
        for location in spec.submodule_search_locations:
            mods.extend(
                top + "." + entry.split(".")[0]
                for entry in os.listdir(location)
                if entry.endswith(".so")
                and entry.split(".")[0].isidentifier()
                and not entry.split(".")[0].startswith("_"))
        return mods
    except Exception:
        return []


failed = []
for name in sorted(names):
    if not name.isidentifier() or name.startswith("_"):
        continue
    try:
        if iu.find_spec(name) is None:
            continue
        __import__(name)
        probes = compiled_submodules(name)
    except Exception:
        failed.append(name)
        continue
    for module in probes:
        try:
            __import__(module)
        except Exception:
            failed.append(name)
            break
print(json.dumps(failed))
"""


def scanned_broken_imports(*, python: Optional[Path] = None) -> Optional[List[str]]:
    """Every installed distribution whose top-level import fails in the env.

    Unlike broken_imports this sees transitive dependencies too (ebooklib's
    lxml, faster-whisper's ctranslate2), where wrong-platform wheels do
    their silent damage. Returns None when the scan itself could not run.
    """
    python = python or env_python()
    try:
        proc = subprocess.run([str(python), "-c", _ENV_SCAN_CODE],
                              capture_output=True, text=True,
                              timeout=600, check=False)
        return json.loads(proc.stdout.strip().splitlines()[-1])
    except (OSError, subprocess.TimeoutExpired, ValueError, IndexError):
        return None


def broken_imports(import_names: List[str],
                   *, python: Optional[Path] = None) -> List[str]:
    """The subset of IMPORT_NAMES that fails to import inside the env.

    Targeted form of scanned_broken_imports: names are probed individually,
    so callers get exactly which of the given names are broken.
    """
    python = python or env_python()
    if not import_names:
        return []
    if _imports_ok(import_names, python=python):
        return []
    return [name for name in import_names
            if not _imports_ok([name], python=python)]


def _venv_tags(python: Optional[Path] = None) -> Optional[dict]:
    """Platform facts of the target env's interpreter, or None on any failure.

    Returns ``{"musl": bool, "pyver": "3.14", "impl": "cp", "abi": "cp314",
    "arch": "x86_64"}`` — everything pip's ``--platform`` repair needs —
    derived from EXT_SUFFIX (``.cpython-314-x86_64-linux-musl.so``) rather
    than sysconfig.get_platform(), which misreports ``linux-x86_64`` for
    portable musl builds and is what lures pip into glibc wheels.
    """
    python = python or env_python()
    code = ("import json, sys, sysconfig; suffix = "
            "sysconfig.get_config_var('EXT_SUFFIX') or ''; "
            "print(json.dumps({'suffix': suffix, "
            "'vi': list(sys.version_info[:2])}))")
    try:
        proc = subprocess.run([str(python), "-c", code],
                              capture_output=True, text=True,
                              timeout=60, check=False)
        data = json.loads(proc.stdout.strip().splitlines()[-1])
    except (OSError, subprocess.TimeoutExpired, ValueError, IndexError):
        return None
    match = re.match(r"\.([a-z]+)-(\d+)-([^-]+)-", data["suffix"])
    if match is None:
        return None
    impl, version, arch = match.groups()
    major, minor = data["vi"]
    return {"musl": "-musl" in data["suffix"],
            "pyver": f"{major}.{minor}",
            "impl": impl[:2],
            "abi": f"{impl[:2]}{version}",
            "arch": arch}


def _installed_specs(python: Optional[Path] = None) -> Dict[str, str]:
    """Map canonical distribution name -> pinned spec (``name==version``).

    Reads ``pip list --format=json`` from the target env so repairs pin the
    exact installed version instead of re-resolving (and never confuse an
    import name like ``bs4`` with a same-named-but-different PyPI project).
    """
    python = python or env_python()
    try:
        proc = subprocess.run(
            [str(python), "-m", "pip", "list", "--format=json",
             "--disable-pip-version-check"],
            capture_output=True, text=True, timeout=120, check=False)
        entries = json.loads(proc.stdout)
    except (OSError, subprocess.TimeoutExpired, ValueError):
        return {}
    specs: Dict[str, str] = {}
    for entry in entries:
        name = entry.get("name")
        version = entry.get("version")
        if name and version:
            canonical = re.sub(r"[-_.]+", "-", name).lower()
            specs.setdefault(canonical, f"{name}=={version}")
    return specs


def _spec_for_import(name: str, installed: Dict[str, str]) -> Optional[str]:
    """Pinned spec for the dist behind import NAME, or None when unknown."""
    for candidate in (name,
                      next((k for k, v in _IMPORT_NAMES.items()
                            if v == name), None)):
        if candidate and re.sub(r"[-_.]+", "-", candidate).lower() in installed:
            return installed[re.sub(r"[-_.]+", "-", candidate).lower()]
    return None


def _site_packages(python: Path) -> Optional[Path]:
    """The env interpreter's pure-Python site-packages directory."""
    try:
        proc = subprocess.run(
            [str(python), "-c",
             "import sysconfig; print(sysconfig.get_paths()['purelib'])"],
            capture_output=True, text=True, timeout=60, check=False)
        return Path(proc.stdout.strip().splitlines()[-1])
    except (OSError, subprocess.TimeoutExpired, IndexError):
        return None


_MUSLLINUX_PLATFORMS = ("musllinux_1_2", "musllinux_1_1")


def repair_imports(broken: List[str], *, python: Optional[Path] = None,
                   emit=None) -> List[str]:
    """Reinstall BROKEN packages from musllinux wheels where possible.

    One targeted reinstall per package: pip uninstalls the mismatched
    distribution and reinstalls the exact same version with the platform
    override flags plus ``--target`` into the env's site-packages, so the
    resolver can only pick musllinux wheels (matching this env's
    interpreter) instead of the glibc wheels it guessed before. Returns
    the names still broken after the attempt — typically packages with no
    musllinux builds at all, which the caller should report as unavailable.
    """
    python = python or env_python()
    tags = _venv_tags(python)
    if tags is None or not tags["musl"] or not broken:
        return list(broken)
    site_packages = _site_packages(python)
    if site_packages is None:
        return list(broken)
    installed = _installed_specs(python)
    still_broken = []
    for name in broken:
        spec = _spec_for_import(name, installed)
        if spec is None:
            print(f"[WARN] {name}: cannot determine the installed package; "
                  "not repaired")
            still_broken.append(name)
            continue
        common.run_console_subprocess(
            [str(python), "-m", "pip", "uninstall", "-y",
             _base_name(spec)], emit=emit)
        argv = [str(python), "-m", "pip", "install", "--no-deps",
                "--upgrade", "--only-binary=:all:",
                "--target", str(site_packages),
                "--python-version", tags["pyver"],
                "--implementation", tags["impl"],
                "--abi", tags["abi"]]
        for platform_base in _MUSLLINUX_PLATFORMS:
            argv += ["--platform", f"{platform_base}_{tags['arch']}"]
        argv.append(spec)
        print(f"[INFO] reinstalling from musllinux wheels: {spec}")
        rc = common.run_console_subprocess(argv, emit=emit)
        if rc != 0:
            print(f"[WARN] {name}: reinstalling {spec} from musllinux "
                  f"wheels failed (pip exit {rc})")
            still_broken.append(name)
    # Verify with the deep scan: a top-level import can succeed while the
    # compiled submodules underneath it still cannot load (lxml's
    # pure-Python __init__ hides exactly this).
    attempted = [name for name in broken if name not in still_broken]
    deep = scanned_broken_imports(python=python)
    if deep is None:
        return list(broken)
    return sorted(set(still_broken)
                  | {name for name in attempted if name in deep})


def _requirements_sha() -> str:
    try:
        data = REQUIREMENTS_PATH.read_bytes()
    except OSError:
        return ""
    return hashlib.sha256(data).hexdigest()


def _marker_valid() -> bool:
    """True when the marker matches both the requirements hash and MARKER_VERSION.

    A bare-hash marker (written by tool versions before MARKER_VERSION
    existed) is treated as invalid, so envs installed before a new
    post-install obligation was added get one re-install + verification.
    """
    try:
        expected = f"{_requirements_sha()}:{MARKER_VERSION}"
        return MARKER_PATH.read_text(encoding="utf-8").strip() == expected
    except OSError:
        return False


def _write_marker() -> None:
    try:
        content = f"{_requirements_sha()}:{MARKER_VERSION}\n"
        MARKER_PATH.write_text(content, encoding="utf-8")
    except OSError:
        pass


def ensure_importable(*, emit=None) -> List[str]:
    """Verify every installed distribution actually imports inside the venv.

    Wheels built for the wrong platform can install "successfully" (pip
    exits 0) while their compiled modules fail to import — glibc wheels
    under a musl interpreter, or a pip wheel cache primed on another
    machine. Broken packages are reinstalled from musllinux wheels when
    possible; whatever cannot be repaired (no compatible build exists) is
    returned so callers can warn about the features it takes down. EMIT
    streams pip's output to an in-TUI task view when given.
    """
    failed = scanned_broken_imports()
    if failed is None:
        print("[WARN] could not scan the managed environment's imports")
        return []
    if not failed:
        return []
    print(f"[WARN] these installed packages fail to import inside "
          f"{ENV_DIR}: {', '.join(failed)}")
    remaining = repair_imports(failed, emit=emit)
    if remaining:
        print(f"[WARN] could not repair: {', '.join(remaining)}. The "
              "related features will be unavailable until compatible "
              "builds exist for this platform.")
    else:
        print("[OK] repaired all previously broken imports")
    return remaining


def ensure_app_env() -> None:
    """Make sure the venv exists and has the current requirements.txt installed.

    Creates the venv when missing, and (re)installs requirements.txt when it
    is missing or has changed since the last install (tracked by a hash +
    version marker). If the full install cannot resolve — platforms whose
    wheels don't cover some OPTIONAL_TAG dependency — it retries without the
    optional lines rather than failing the bootstrap. Afterwards the imports
    are verified and wrong-platform wheels repaired (ensure_importable).
    Raises RuntimeError on any failure so the caller can abort before re-exec.
    """
    if not env_exists() and create_env() != 0:
        raise RuntimeError("could not create the managed environment")
    if not _marker_valid():
        if install_requirements() != 0:
            skipped = [spec for spec, optional in requirement_specs()
                       if optional]
            if not skipped:
                raise RuntimeError("pip install -r requirements.txt failed")
            print(f"[WARN] retrying without optional requirements: "
                  f"{', '.join(_base_name(spec) for spec in skipped)}")
            if install_requirements(skip_optional=True) != 0:
                raise RuntimeError("pip install -r requirements.txt failed")
        ensure_importable()
        _write_marker()


def bootstrap(script_path: str) -> None:
    """Run audiobook.py inside the managed venv, creating it first if needed.

    A no-op when the current process is already the venv's interpreter. Otherwise
    ensures the env (and requirements) are ready, then replaces the process with
    the venv's python running the same script and CLI args. Called at the top of
    audiobook.py before any third-party import.
    """
    if is_managed_env():
        return
    try:
        ensure_app_env()
    except RuntimeError as exc:
        print(f"[FATAL] {exc}", file=sys.stderr)
        sys.exit(1)
    py = str(env_python())
    target = str(Path(script_path).resolve())
    print(f"[INFO] re-launching inside managed environment: {py}")
    os.execv(py, [py, target, *sys.argv[1:]])