ModuleUpdate.py

import sys
import os
import subprocess
import multiprocessing
import json
import shutil
import time
import datetime
import zipfile
import re
import logging
import tempfile
import contextlib
import errno
import importlib.metadata
import importlib.util

logger = logging.getLogger("Update")

if not logging.getLogger().hasHandlers():
    logging.basicConfig(level=logging.DEBUG, format='%(message)s', stream=sys.stdout)

from pathlib import Path
from collections.abc import Iterable
from typing import Any, List, Optional, TypeVar, override

from importlib import invalidate_caches
from BaseUtils import tuplize_version, Version, local_path, mwgg_venv_site_packages, use_worlds_venv, is_frozen
from APContainer import APWorldContainer

# mwgg_igdb package source — orphan branch on the Index repo
# See MultiworldGG-Index/scripts/build_variants.py for variant definitions.
DEFAULT_MWGG_IGDB_VARIANT = "sixteen"  # ultimate fallback when nothing is installed
_VARIANTS = ("nr", "ao", "twelve", "sixteen")
_EXPLICIT_VARIANT: Optional[str] = None  # set via set_variant(); wins over detection
MWGG_INDEX_REPO = "MultiworldGG/MultiworldGG-Index"
# The three globals below mirror the currently *resolved* variant. _resolve_variant()
# keeps them consistent; callers should treat them as read-only.
MWGG_IGDB_VARIANT = DEFAULT_MWGG_IGDB_VARIANT
MWGG_IGDB_BRANCH = f"game_index_{MWGG_IGDB_VARIANT}"
MWGG_IGDB_GIT_URL = f"git+https://github.com/{MWGG_INDEX_REPO}@{MWGG_IGDB_BRANCH}"

def is_frozen() -> bool:
    return getattr(sys, 'frozen', False)

def is_windows() -> bool:
    return sys.platform in ("win32", "cygwin", "msys")

def is_macos() -> bool:
    return sys.platform == "darwin"

def is_linux() -> bool:
    return sys.platform.startswith("linux")

def install_path() -> Path:
    # Returns the path to the install directory for the python modules
    # Frozen builds only
    if is_windows():
        return Path.home() / "AppData" / "Local" / "MultiworldGG" / "mwgg_venv"
    elif is_macos():
        return Path.home() / "Library" / "Application Support" / "MultiworldGG" / "mwgg_venv"
    elif is_linux():
        return Path.home() / ".local" / "share" / "MultiworldGG" / "mwgg_venv"
    else:
        raise RuntimeError("Unsupported platform")

# Version compatibility checks
if (is_windows() or is_macos()) and sys.version_info < (3, 13, 0):
    raise RuntimeError(f"Incompatible Python Version found: {sys.version_info}. Official 3.13.+ is supported.")
elif sys.version_info < (3, 13, 0):
    raise RuntimeError(f"Incompatible Python Version found: {sys.version_info}. 3.13.+ is supported.")

def _worlds_venv_is_readonly() -> bool:
    """True when the worlds venv lives on a read-only mount (e.g. a Docker `:ro`
    bind mount). Such consumers must never attempt installs — the mwgg_upgrader
    service is the sole writer of the venv. Only meaningful when use_worlds_venv()
    is set; always False on a normal dev/user install (writable venv)."""
    if not use_worlds_venv():
        return False
    venv_dir = install_path()
    try:
        venv_dir.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(dir=venv_dir):
            pass
    except OSError:
        return True
    return False


# Detected once at import: a read-only worlds venv disables every install path
# below, exactly like SKIP_ALL_INSTALLS, so a consumer that reaches update()
# without the env opt-out still can't crash writing the upgrader-owned venv.
_VENV_READONLY = _worlds_venv_is_readonly()
if _VENV_READONLY:
    logger.info("Worlds venv is read-only; installs disabled (mwgg_upgrader owns writes).")


def _skip_all_installs() -> bool:
    """Installs are off via explicit env opt-out or a read-only worlds venv."""
    return bool(os.environ.get("SKIP_ALL_INSTALLS")) or _VENV_READONLY


# Skip update if running in splash screen process
# Allow updates in main process and main client process
_skip_update = bool(
    (multiprocessing.parent_process() and multiprocessing.current_process().name != "MultiWorldGG")
    or os.environ.get("SKIP_REQUIREMENTS_UPDATE", "")
    or _skip_all_installs()
)

update_ran = _skip_update
need_update: List[str] = []

_T = TypeVar("_T")


class RequirementsSet(set[_T]):
    """Custom set that tracks whether updates have been run."""

    @override
    def add(self, e: _T) -> None:
        global update_ran
        update_ran &= _skip_update
        super().add(e)

    @override
    def update(self, *s: Iterable[_T]) -> None:
        global update_ran
        update_ran &= _skip_update
        super().update(*s)


# Initialize file sets

requirements_files: RequirementsSet[Path] = RequirementsSet({Path(local_path("requirements.txt"))})
worlds_files: dict[str, RequirementsSet[str]] = {"wheels": RequirementsSet(), "apworlds": RequirementsSet()}

# custom_worlds always lives next to the executable / source checkout -- the
# upstream location, and where users actually drop their apworlds. This is the
# single source of truth for the launch scan (register_custom_worlds /
# get_available_worlds), Utils.set_game_names, and the launch path. Do NOT
# special-case frozen builds to write_path(): that splits the scan from the launch
# path and custom worlds silently stop being selectable.
def _resolve_custom_worlds_dir() -> Path:
    return Path(local_path("custom_worlds"))


custom_worlds_dir = _resolve_custom_worlds_dir()

# Best-effort mkdir. Skipped silently on the rare read-only filesystem; the
# downstream readers already handle missing directories.
try:
    custom_worlds_dir.mkdir(parents=True, exist_ok=True)
except OSError:
    pass

def _scan_custom_worlds() -> None:
    """Register any .whl/.apworld files in custom_worlds_dir into worlds_files.

    Skipped once a full update has run (those files are handled by the updater).
    Reads the module-level update_ran / custom_worlds_dir / worlds_files, so tests
    can monkeypatch them and call this directly instead of re-implementing it.
    """
    if update_ran or not custom_worlds_dir.exists():
        return
    for world_file in custom_worlds_dir.glob("*.whl"):
        worlds_files["wheels"].add(str(world_file))
    for world_file in custom_worlds_dir.glob("*.apworld"):
        worlds_files["apworlds"].add(str(world_file))


# Add wheel files if update hasn't run
_scan_custom_worlds()

# Default for dev mode (not frozen): use the running interpreter and let uv install into its venv.
python_cmd = sys.executable

_uv_resolved_path: Optional[Path] = None
_uv_unavailable: bool = False


def _uv_candidate_paths() -> list[Path]:
    """uv lookup order; don't hunt for it
    """
    candidates: list[Path] = []
    # Frozen builds on Linux/macOS ship uv next to the executable 
    if is_frozen():
        exe_dir = Path(sys.executable).parent
        if is_macos():
            import platform
            arch = platform.machine()  # "arm64" on Apple Silicon, "x86_64" on Intel
            candidates.append(exe_dir / f"uv-{arch}")
        elif not is_windows():
            candidates.append(exe_dir / "uv")
    candidates.append(Path("uv"))  # PATH lookup
    if is_windows():
        candidates += [
            Path.home() / "AppData" / "Local" / "Microsoft" / "WinGet" / "Links" / "uv.exe",  # winget shim
            Path.home() / ".local" / "bin" / "uv.exe",                          # astral PS installer
        ]
    else:
        candidates += [
            Path.home() / ".local" / "bin" / "uv",     # astral installer / pipx
            Path("/opt/homebrew/bin/uv"),              # Homebrew (Apple Silicon)
            Path("/usr/local/bin/uv"),                 # Homebrew (Intel) / generic
        ]
    return candidates


def _uv_pip(*args: str) -> list[str]:
    return ["pip", *args, "--python", str(python_cmd)]


def _uv_run(args: list[str], timeout: float = 120, check: bool = False) -> subprocess.CompletedProcess[str]:
    """Run `uv <args>` against the first reachable uv binary."""
    global _uv_resolved_path, _uv_unavailable

    # Windows-only: detach into a new process group with no console window so a uv
    # subprocess can't flash a window or steal the parent's Ctrl-C. 0 is the
    # cross-platform no-op default elsewhere.
    creationflags = (
        subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.CREATE_NO_WINDOW
        if is_windows() else 0
    )

    if _uv_unavailable:
        return subprocess.CompletedProcess(args, 127, "", "uv not found at any known path")

    candidates = [_uv_resolved_path] if _uv_resolved_path else _uv_candidate_paths()

    for cand in candidates:
        cmd = [cand] + args
        try:
            result = subprocess.run(
                cmd,
                check=check,
                capture_output=True,
                text=True,
                stdin=subprocess.DEVNULL,
                timeout=timeout,
                creationflags=creationflags,
            )
        except OSError as e:
            # The candidate is unusable — try next.
            logger.debug(f"uv not usable at {cand} ({e!r}); trying next candidate")
            continue
        if _uv_resolved_path is None:
            _uv_resolved_path = cand
            logger.debug(f"Using uv at {cand}")
        return result

    _uv_unavailable = True
    if is_windows():
        install_hint = (
            "install uv via `winget install astral-sh.uv`, "
            "`irm https://astral.sh/uv/install.ps1 | iex`, or `choco install uv`"
        )
    else:
        install_hint = "install uv via `curl -LsSf https://astral.sh/uv/install.sh | sh`"
    logger.warning(
        "uv not found at any known install path. Worlds cannot be pre-installed; "
        f"they will be installed on demand when needed. To pre-install, {install_hint}."
    )
    return subprocess.CompletedProcess(args, 127, "", "uv not found at any known path")


def venv_is_healthy(venv_path: Path) -> bool:
    """True if the venv's interpreter actually runs.

    No `.exists()` pre-probes: stat()-ing the venv's `home =` (a uv-managed
    python) raises OSError on untraversable mount points (WinError 448) and
    tells us nothing that running the interpreter doesn't. We just invoke it —
    a missing exe, a dead base python, an untraversable path, or a timeout all
    surface as a failure, which means "unhealthy" and the caller recreates.
    """
    venv_python = venv_path / ("Scripts" if is_windows() else "bin") / ("python.exe" if is_windows() else "python")
    try:
        return subprocess.run([str(venv_python), "--version"], capture_output=True, timeout=10).returncode == 0
    except (subprocess.TimeoutExpired, OSError):
        return False


if use_worlds_venv():
    # Route worlds + mwgg_igdb into a dedicated venv under user data
    if is_frozen():
        exe_dir = Path(sys.exec_prefix)
        default_libs_dir = Path(exe_dir, "lib")
        if str(default_libs_dir) not in sys.path:
            sys.path.append(str(default_libs_dir))

    venv_path = install_path()
    venv_ready = True
    if not venv_is_healthy(venv_path):
        # Any failure here (uv missing/timeout, unwritable dir, untraversable
        # mount point) must degrade to "install on demand later", never crash
        # the import — this runs at module load for every consumer.
        try:
            venv_path.mkdir(parents=True, exist_ok=True)
            if any(venv_path.iterdir()):
                logger.info(f"Repairing stale venv at {venv_path} via uv (site-packages preserved).")
            else:
                logger.info(f"Creating venv at {venv_path} via uv.")
            # uv reuses an existing system Python 3.13 if one is present; otherwise it
            # downloads python-build-standalone.
            venv_result = _uv_run(
                ["venv", str(venv_path), "--allow-existing", "--python", "3.13"],
                timeout=600,
            )
            venv_ready = venv_result.returncode == 0
        except Exception as e:
            logger.debug(f"Worlds venv setup failed: {e!r}")
            venv_ready = False
        if not venv_ready:
            logger.warning(
                "Could not create the worlds venv. Worlds will be installed on demand "
                "the next time uv is available."
            )

    if venv_ready:
        python_cmd = venv_path / ("Scripts" if is_windows() else "bin") / ("python.exe" if is_windows() else "python")

        # Make packages installed into the worlds venv (mwgg_igdb, plus any
        # top-level helpers shipped alongside worlds) importable from the
        # running process.
        site_packages = mwgg_venv_site_packages()
        if site_packages not in sys.path:
            sys.path.insert(0, site_packages)


def confirm(msg: str) -> None:
    """Get user confirmation for an action."""
    try:
        input(f"\n{msg}")
    except KeyboardInterrupt:
        logger.info("\nAborting")
        sys.exit(1)


def _format_manual_install_hint(module_location: str) -> str:
    """Render a copy-pasteable `uv pip install` command for the host shell.
    """
    return (
        "To install manually on the host (where build tools are available), run:\n"
        f"    bash:      source {mwgg_venv_site_packages()}/bin/activate\n"
        f"    win PS:    & {mwgg_venv_site_packages()}/Scripts/Activate.ps1\n"
        f"    docker:    source ~/<your-local-venv>/bin/activate\n"
        f"    then run:\n"
        f"               uv pip install '{module_location}' --upgrade --no-cache\n"
        f"    docker admins will need to copy the site packages to /var/lib/mwgg/mwgg_venv"
    )


def parse_requirements_file(file_path: Path) -> List[str]:
    """
    Parse a requirements.txt file and return a list of requirement strings.
    Handles line continuations, comments, and various requirement formats.
    """
    requirements: list[str] = []
    
    with open(file_path, 'r') as f:
        lines = f.readlines()
    
    prev_line = ""
    
    for line in lines:
        line = line.rstrip('\r\n')
        
        # Handle line continuations
        if line.endswith('\\'):
            prev_line += line[:-1] + " "
            continue
        
        line = prev_line + line
        prev_line = ""
        
        # Skip empty lines and comments
        if not line.strip() or line.strip().startswith('#'):
            continue
        
        # Remove hash specifications for version checking
        line = line.split("--hash=")[0].strip()
        
        # Handle URL-based requirements
        if line.startswith(("https://", "git+https://")):
            line = _parse_url_requirement(line)
        
        # Handle custom PEP 508 syntax
        elif "@" in line and "#" in line:
            line = _parse_custom_pep508_requirement(line)
        
        if line.strip():
            requirements.append(line.strip())
    
    return requirements


def _parse_url_requirement(line: str) -> str:
    """Parse URL-based requirements and extract package name and version."""
    rest = line.split('/')[-1]
    
    # Extract from filename
    if "@" in rest:
        raise ValueError("Can't deduce version from requirement")
    
    rest = rest.replace(".zip", "-").replace(".tar.gz", "-")
    try:
        name, version, _ = rest.split("-", 2)
        return f'{name}=={version}'
    except ValueError:
        return ""


def _parse_custom_pep508_requirement(line: str) -> str:
    """Parse custom PEP 508 syntax: name @ url#version ; marker."""
    name, rest = line.split("@", 1)
    version = rest.split("#", 1)[1].split(";", 1)[0].rstrip()
    result = f"{name.rstrip()}=={version}"
    
    if ";" in rest:  # keep marker
        result += rest[rest.find(";"):]
    
    return result


def _detect_installed_variant() -> Optional[str]:
    """Return the variant currently installed locally, or None if undetectable.

    Reads the `__variant__` constant the Index build bakes into the generated
    `mwgg_igdb` module. Absent (pre-`__variant__` build) or unimportable → None,
    and callers fall back to DEFAULT_MWGG_IGDB_VARIANT.
    """
    if importlib.util.find_spec("mwgg_igdb") is None:
        return None
    try:
        import mwgg_igdb
    except ImportError:
        return None
    variant = getattr(mwgg_igdb, "__variant__", None)
    if isinstance(variant, str) and variant in _VARIANTS:
        return variant
    return None


def _resolve_variant() -> str:
    """Pick the variant to act on, refresh the derived globals, and return it.

    Precedence: explicit `set_variant()` > detected install > default fallback.
    """
    global MWGG_IGDB_VARIANT, MWGG_IGDB_BRANCH, MWGG_IGDB_GIT_URL
    if _EXPLICIT_VARIANT is not None:
        variant = _EXPLICIT_VARIANT
    else:
        variant = _detect_installed_variant() or DEFAULT_MWGG_IGDB_VARIANT
    # These UPPERCASE names are public, documented module globals that are
    # intentionally reassigned here (callers and tests read ModuleUpdate.MWGG_IGDB_*);
    # the "constant" rule doesn't apply and renaming would break the public API.
    MWGG_IGDB_VARIANT = variant  # pyright: ignore[reportConstantRedefinition]
    MWGG_IGDB_BRANCH = f"game_index_{variant}"  # pyright: ignore[reportConstantRedefinition]
    MWGG_IGDB_GIT_URL = f"git+https://github.com/{MWGG_INDEX_REPO}@{MWGG_IGDB_BRANCH}"  # pyright: ignore[reportConstantRedefinition]
    return variant


def _igdb_install_date() -> Optional[datetime.date]:
    """Local date `mwgg_igdb` was last written to disk, or None if not installed.

    The module file's mtime is the package's own install datestamp — set fresh
    every time uv (re)installs it — so no separate stamp file is needed.
    """
    spec = importlib.util.find_spec("mwgg_igdb")
    if spec is None or not spec.origin or not os.path.exists(spec.origin):
        return None
    return datetime.date.fromtimestamp(os.path.getmtime(spec.origin))


def _igdb_upgraded_recently() -> bool:
    """True when an upgrade pull would be a no-op: `mwgg_igdb` was installed
    today. Variant switches do NOT rely on this throttle — they go through the
    callers that pass force=True (the `mwgg_igdb_<variant>` token path in
    install_worlds and the mwgg_upgrader), which bypass it entirely.
    """
    install_date = _igdb_install_date()
    return install_date is not None and install_date == datetime.date.today()


# Consumed by the upgrader tools (tools/mwgg_upgrade.py, tools/mcp_mwgg_upgrader.py),
# so it is unused *within* this module — hence the targeted ignore.
def _venv_has_worlds() -> bool:  # pyright: ignore[reportUnusedFunction]
    try:
        worlds_dir = _venv_worlds_dir()
        return worlds_dir.exists() and any(worlds_dir.iterdir())
    except OSError:
        return False


def install_mwgg_igdb(upgrade: bool = False, force: bool = False) -> bool:
    """Install or refresh the mwgg_igdb package from the Index repo orphan branch.

    Called before any code path that imports `mwgg_igdb` — the package is the
    runtime source-of-truth for which worlds exist and where to fetch them.

    Args:
        upgrade: Run pip with --upgrade.
        force: With upgrade=True, bypass the once-daily throttle. Use for variant
               switches and standalone CLI invocation where a fresh pull is required.

    Concurrency: two processes that race on a stale stamp can both run pip into the
    same venv. uv pip writes to a temp location before rename, so the worst outcome
    is two pulls instead of one — not corruption.

    Returns True if the install succeeded (or was throttled).
    """
    if _skip_all_installs():
        return True
    _resolve_variant()
    if upgrade and not force and _igdb_upgraded_recently():
        logger.debug("mwgg_igdb already installed today; skipping upgrade pull")
        return True
    args = _uv_pip("install", MWGG_IGDB_GIT_URL, "--no-cache")
    if upgrade:
        # --reinstall rewrites the tiny package even when the branch HEAD is
        # unchanged, so the module's install date (its mtime) advances to today
        # and the once-daily throttle above stays satisfied until tomorrow.
        args.append("--reinstall")
    logger.info(f"Installing mwgg_igdb ({MWGG_IGDB_VARIANT}) from {MWGG_IGDB_BRANCH}")
    try:
        result = _uv_run(args, timeout=300)
    except subprocess.TimeoutExpired:
        logger.warning("uv install of mwgg_igdb timed out.")
        return False
    if result.returncode != 0:
        logger.warning(f"Failed to install mwgg_igdb: {result.stderr}")
        return False
    return True


def _get_game_index():
    """Lazy-import GameIndex; install mwgg_igdb if missing. Returns None on failure."""
    try:
        from mwgg_igdb import GameIndex
        return GameIndex
    except ImportError:
        if install_mwgg_igdb():
            invalidate_caches()
            try:
                from mwgg_igdb import GameIndex
                return GameIndex
            except ImportError as e:
                logger.warning(f"mwgg_igdb still unimportable after install: {e}")
        return None


def _module_location_tag(url: str) -> Optional[str]:
    """Extract the version from a release-asset wheel URL.

    Expects URLs like
    ``https://github.com/<owner>/<repo>/releases/download/<release_tag>/<dist>-<ver>-py3-none-any.whl``,
    optionally with a ``#sha256=<hex>`` fragment. Returns None for anything
    that isn't a recognizable wheel URL (legacy ``git+...@<ref>`` URLs from
    the v2 publish flow degrade to None — the caller then skips the
    comparison rather than crashing).
    """
    if not url:
        return None
    name = url.rsplit("/", 1)[-1]
    name = name.split("#", 1)[0].split("?", 1)[0]
    if not name.endswith(".whl"):
        return None
    parts = name[:-len(".whl")].split("-")
    # PEP 427: dist, version, [build,] python, abi, platform — version is index 1.
    if len(parts) < 5:
        return None
    return parts[1]


def _parse_variant_token(token: str) -> Optional[str]:
    """Return the variant name if `token` is `mwgg_igdb` or `mwgg_igdb_<variant>`, else None.

    Bare `mwgg_igdb` (no suffix) maps to the canonical default `sixteen`. Inno Setup
    passes one of these tokens in the `--worlds` list to select the parental-rating gate.
    """
    if token == "mwgg_igdb":
        return "sixteen"
    prefix = "mwgg_igdb_"
    if token.startswith(prefix):
        variant = token[len(prefix):]
        if variant in _VARIANTS:
            return variant
    return None


def set_variant(variant: str) -> None:
    """Switch the runtime mwgg_igdb variant; takes effect on next install_mwgg_igdb call.

    Sets the explicit-override sentinel so this choice wins over any detected
    installed variant on subsequent _resolve_variant() calls.
    """
    global _EXPLICIT_VARIANT
    _EXPLICIT_VARIANT = variant  # pyright: ignore[reportConstantRedefinition]  # intentional reassignment of the override sentinel
    _resolve_variant()


def _world_slug(world: str) -> str:
    return world.removeprefix("worlds.")


def _world_requires_install(slug: str, games: dict[str, dict[str, object]]) -> bool:
    try:
        dist = importlib.metadata.distribution(f"worlds.{slug}")
    except importlib.metadata.PackageNotFoundError:
        return True

    module_location = games.get(slug, {}).get("module_location")
    if not isinstance(module_location, str):
        return False

    tag = _module_location_tag(module_location)
    return bool(tag and dist.version != tag)


def _worlds_requiring_install(worlds: list[str], games: dict[str, dict[str, object]]) -> list[str]:
    return [world for world in worlds if _world_requires_install(_world_slug(world), games)]


def check_for_updates(worlds_only: bool = False) -> List[str]:
    """
    Return packages with newer versions available.

    For worlds: re-pull mwgg_igdb (always latest), then return slugs whose
    installed dist version doesn't match the tag in `module_location`.
    For non-world packages: query PyPI against requirements.txt entries.
    """
    if is_frozen() and not worlds_only:
        return []

    if worlds_only:
        install_mwgg_igdb(upgrade=True)
        index = _get_game_index()
        if index is None:
            return []
        outdated: List[str] = []
        for slug, entry in index.get_all_games().items():
            loc = entry.get("module_location")
            if not loc:
                continue
            tag = _module_location_tag(loc)
            if not tag:
                continue
            try:
                dist = importlib.metadata.distribution(f"worlds.{slug}")
            except importlib.metadata.PackageNotFoundError:
                continue
            if dist.version != tag:
                outdated.append(f"worlds.{slug}")
        logger.info(f"Worlds with available updates: {outdated}")
        return outdated

    # Dev-only path: ask uv for outdated dists. uv's resolver enforces requirements.txt
    # specifiers at install time, so we don't need to pre-filter here.
    try:
        executable_args = _uv_pip("list", "--outdated", "--format", "json")
        logger.info(f"Executing subprocess command: {executable_args}")
        response = _uv_run(executable_args, timeout=45)
        if response.returncode != 0:
            logger.warning(f"Could not check for updates: {response.stderr}")
            return []

        outdated_packages = json.loads(response.stdout)
        logger.info(f"Newer versions of the following packages are available: {outdated_packages}")
        return [pkg["name"] for pkg in outdated_packages]

    except (subprocess.TimeoutExpired, json.JSONDecodeError, FileNotFoundError) as e:
        logger.warning(f"Could not check for updates: {e}")
        return []


def uninstall_worlds(worlds: List[str]) -> None:
    """Uninstall a list of `worlds.<slug>` packages from the venv."""
    for world in worlds:
        # uv pip uninstall is non-interactive by default; no --yes equivalent needed.
        try:
            _uv_run(_uv_pip("uninstall", world), timeout=60)
        except subprocess.TimeoutExpired:
            logger.warning(f"uv uninstall of {world} timed out.")


def find_world_modules() -> set[str]:
    """Return all known world slugs: union of mwgg_igdb entries and currently installed `worlds.<slug>` dists."""
    world_modules_set: set[str] = set()

    index = _get_game_index()
    if index is not None:
        world_modules_set.update(index.get_all_games().keys())

    try:
        executable_args = _uv_pip("list", "--format", "json")
        logger.debug(f"Executing subprocess command to find installed worlds: {executable_args}")
        response = _uv_run(executable_args, timeout=45)
        if response.returncode == 0:
            for package in json.loads(response.stdout):
                package_name = package.get("name", "")
                if package_name.startswith("worlds") and len(package_name) > 7:
                    # uv hyphenates dist names (worlds.dark_souls_3 -> worlds-dark-souls-3); restore the slug.
                    world_name = package_name[7:].replace("-", "_")
                    if not world_name.startswith("_"):
                        world_modules_set.add(world_name)
        else:
            logger.warning(f"Could not list installed packages: {response.stderr}")
    except (subprocess.TimeoutExpired, json.JSONDecodeError, FileNotFoundError) as e:
        logger.warning(f"Could not check installed world modules: {e}")
    except Exception as e:
        logger.warning(f"Unexpected error while checking installed world modules: {e}")

    return world_modules_set


def _venv_worlds_dir() -> Path:
    """Return the venv worlds dir from which worlds/__init__.py extends __path__.
    Apworlds get extracted here so they're importable via the normal file loader
    (multiprocessing.spawn in child processes needs disk-based modules; zipimport-
    only modules can't be re-imported in the spawned child).
    """
    if use_worlds_venv():
        return Path(mwgg_venv_site_packages("worlds"))
    # Dev: matches the hardcoded path in src/worlds/__init__.py
    from sysconfig import get_path
    return Path(get_path("purelib")) / "worlds"


def _install_apworld_to_venv(apworld_file: Path, slug: str) -> bool:
    """Extract the `<slug>/` directory from apworld_file into the venv worlds dir.
    Returns True on success. Overwrites existing files in place rather than
    rmtree'ing (rmtree fails on Windows if the module is currently loaded).
    """
    venv_worlds = _venv_worlds_dir()
    try:
        venv_worlds.mkdir(parents=True, exist_ok=True)
        prefix = f"{slug}/"
        with zipfile.ZipFile(apworld_file, "r") as zf:
            members = [m for m in zf.namelist() if m == prefix or m.startswith(prefix)]
            if not members:
                logger.warning(f"Apworld {apworld_file} contains no '{slug}/' directory")
                return False
            for member in members:
                zf.extract(member, str(venv_worlds))
        # Refresh the target dir's mtime so the stale-extraction pruner uses
        # the most recent extraction as the "last used" signal, even if zipfile
        # restored historical timestamps from the archive members.
        target_dir = venv_worlds / slug
        try:
            os.utime(target_dir, None)
        except OSError:
            pass
        logger.info(f"Extracted apworld {apworld_file} to {target_dir}")
        # If the worlds package is already imported, extend its __path__ so the
        # newly-extracted module is discoverable without a restart. (At startup,
        # worlds/__init__.py does this itself when it first runs.)
        worlds_pkg = sys.modules.get("worlds")
        if worlds_pkg is not None and hasattr(worlds_pkg, "__path__"):
            venv_str = str(venv_worlds)
            if venv_str not in worlds_pkg.__path__:
                worlds_pkg.__path__.append(venv_str)
        return True
    except Exception as e:
        logger.error(f"Failed to extract apworld {apworld_file} to {venv_worlds}: {e}")
        return False


def _prune_stale_apworld_extractions(max_age_days: int = 30) -> None:
    """Remove extracted-apworld dirs in the venv worlds dir whose mtime is older
    than max_age_days. Skips any dir backed by a real importlib.metadata
    Distribution (pip-installed wheels), so this only ever touches our own
    extraction output.
    """
    import importlib.metadata
    venv_worlds = _venv_worlds_dir()
    if not venv_worlds.exists():
        return
    cutoff = time.time() - max_age_days * 86400
    for entry in venv_worlds.iterdir():
        if not entry.is_dir() or entry.name.startswith("_"):
            continue
        try:
            importlib.metadata.distribution(f"worlds.{entry.name}")
            continue  # pip-installed; leave alone
        except importlib.metadata.PackageNotFoundError:
            pass
        try:
            if entry.stat().st_mtime >= cutoff:
                continue
        except OSError:
            continue
        try:
            shutil.rmtree(entry)
            logger.info(f"Pruned stale extracted apworld {entry} (mtime > {max_age_days}d)")
        except OSError as e:
            logger.warning(f"Could not prune stale apworld {entry}: {e}")


def install_worlds(worlds: List[str], update: bool = False, with_deps: bool = False) -> list[str]:
    """
    Install worlds by resolving each apworld's `module_location` from mwgg_igdb and pip-installing the URL.

    `module_location` is a `https://.../<dist>-<world_version>-py3-none-any.whl#sha256=<hex>`
    release-asset URL set by the Index repo.

    Falls back to a custom_worlds/<slug>.apworld lookup if the apworld isn't in the index or its
    `module_location` install fails.

    Args:
        worlds: List of apworlds to install.
        update: If True, uninstall old versions first.
        with_deps: If True, install the wheel *with* its transitive dependencies.
            Otherwise, dependencies are still installed for new worlds and skipped
            for worlds that were already present.

    Returns:
        List of apworlds that fell back to a custom apworld.
    """
    if _skip_all_installs():
        return []
    apworlds: list[str] = []

    world_slugs: list[str] = []
    selected_variant: Optional[str] = None
    for entry in worlds:
        variant = _parse_variant_token(entry)
        if variant is not None:
            selected_variant = variant
        else:
            world_slugs.append(entry)
    if selected_variant is not None:
        set_variant(selected_variant)
        install_mwgg_igdb(upgrade=True, force=True)

    index = _get_game_index()
    games: dict[str, dict[str, Any]] = index.get_all_games() if index is not None else {}
    # Snapshot BEFORE uninstall_worlds: a world properly installed at the current
    # mwgg_igdb tag stays in the set so update=True reinstalls can skip deps.
    installed_world_slugs = {
        _world_slug(world) for world in world_slugs
        if not _world_requires_install(_world_slug(world), games)
    }

    if update:
        logger.info(f"Uninstalling old versions of: {world_slugs}")
        uninstall_worlds(world_slugs)

    if not update and not with_deps and world_slugs:
        worlds_to_install = _worlds_requiring_install(world_slugs, games)
        skipped_worlds = sorted(set(world_slugs) - set(worlds_to_install))
        if skipped_worlds:
            logger.debug(f"Skipping already-installed worlds: {skipped_worlds}")
        if not worlds_to_install:
            _prune_stale_apworld_extractions()
            invalidate_caches()
            return apworlds
        world_slugs = worlds_to_install

    for world in world_slugs:
        slug = world.removeprefix("worlds.")
        target = f"worlds.{slug}"

        if update:
            logger.info(f"Updating world: {target}")
        else:
            logger.info(f"Installing world: {target}")

        entry = games.get(slug, {})
        module_location = entry.get("module_location")

        if not module_location:
            logger.warning(f"No module_location for {slug} in mwgg_igdb; checking custom_worlds")
            apworld_file = custom_worlds_dir / f"{slug}.apworld"
            if apworld_file.exists():
                logger.info(f"Found apworld file: {apworld_file}")
                if _install_apworld_to_venv(apworld_file, slug):
                    apworlds.append(target)
            else:
                logger.warning(f"Custom apworld file not found at {apworld_file}, {slug} cannot be installed")
            continue

        install_args = ["install"]
        if not with_deps and slug in installed_world_slugs:
            install_args.append("--no-deps")
        install_args += [module_location, "--upgrade", "--no-cache"]
        executable_args = _uv_pip(*install_args)
        logger.info(f"Executing subprocess command: {executable_args}")
        try:
            result = _uv_run(executable_args, timeout=300)
        except subprocess.TimeoutExpired:
            logger.warning(f"uv install of {target} timed out; treating as failure.")
            logger.warning(_format_manual_install_hint(module_location))
            apworld_file = custom_worlds_dir / f"{slug}.apworld"
            if apworld_file.exists():
                logger.info(f"Found apworld file: {apworld_file}")
                if _install_apworld_to_venv(apworld_file, slug):
                    apworlds.append(target)
            continue
        logger.info(result.stdout)

        if result.returncode != 0:
            stderr_text = (result.stderr or "").strip() or "uv returned non-zero with no stderr"
            logger.warning(f"World {target} failed to install from {module_location}:\n{stderr_text}")
            logger.warning(_format_manual_install_hint(module_location))
            apworld_file = custom_worlds_dir / f"{slug}.apworld"
            if apworld_file.exists():
                logger.info(f"Found apworld file: {apworld_file}")
                if _install_apworld_to_venv(apworld_file, slug):
                    apworlds.append(target)
            else:
                logger.warning(f"Custom apworld file not found at {apworld_file}")
        else:
            logger.info(f"Successfully installed {target}")

    _prune_stale_apworld_extractions()
    invalidate_caches()
    return apworlds

def update_world_from_package() -> None:
    """Install/update wheel files from custom_worlds directory."""
    # Use threading version if frozen, otherwise use subprocess
    if is_frozen():
        for world in worlds_files["wheels"]:
            logger.info(f"Installing wheel: {world}")
            # uv prefers wheels by default, no --prefer-binary equivalent needed.
            executable_args = _uv_pip("install", world, "--upgrade", "--no-cache")

            # Use threading instead of multiprocessing to avoid argument contamination
            import threading
            import queue

            result_queue: queue.Queue[tuple[int, str, str]] = queue.Queue()

            def _pip_install_thread():
                try:
                    result = _uv_run(executable_args, timeout=30)
                    result_queue.put((result.returncode, result.stdout, result.stderr))
                except Exception as e:
                    result_queue.put((1, "", str(e)))
            
            install_thread = threading.Thread(target=_pip_install_thread, daemon=True)
            install_thread.start()
            install_thread.join()
            
            # Get the return values from the worker thread
            try:
                returncode, stdout, stderr = result_queue.get_nowait()
                logger.info(stdout)
            except:
                returncode = 1  # Assume failure if we can't get the result
                stdout = ""
                stderr = "Failed to get process result"
            
            if returncode != 0:
                logger.warning(f"Failed to install wheel {world}")
                if stderr:
                    logger.error(f"{stderr}")
            else:
                logger.info(f"Successfully installed wheel {world}")

        for world in worlds_files["apworlds"]:
            logger.info(f"APWorld found, checking versions: {world}")
            try:
                # Extract module name from apworld filename (e.g., "world_name.apworld" -> "world_name")
                world_path = Path(world)
                module_name = world_path.stem  # Gets filename without extension
                
                # Read version from the apworld zip file using APWorldContainer
                new_version: Optional[Version] = None
                manifest: dict[str, object] = {}
                try:
                    apworld_container = APWorldContainer(world_path)
                    # Set manifest path to expected location
                    with zipfile.ZipFile(world, 'r') as apworld_zip:
                        manifest = apworld_container.read_contents(apworld_zip)
                    if "world_version" in manifest:
                        # manifest is untrusted external data (dict[str, object]); coerce the
                        # value to str so a non-string world_version can't crash tuplize_version.
                        new_version = tuplize_version(str(manifest["world_version"]))
                        logger.info(f"APworld {world} has version {new_version}")
                    else:
                        logger.info(f"APworld {world} has no world_version specified")
                except Exception as e:
                    logger.warning(f"Failed to read version from APworld {world}: {e}")
                
                # Check if world is already installed using pip show
                package_name = f"worlds.{module_name}"
                installed_version: Optional[Version] = None
                
                try:
                    executable_args = _uv_pip("show", package_name)
                    result = _uv_run(executable_args, timeout=10)

                    if result.returncode == 0:
                        # Package is installed, parse version from output
                        for line in result.stdout.splitlines():
                            if line.startswith("Version:"):
                                version_str = line.split(":", 1)[1].strip()
                                installed_version = tuplize_version(version_str)
                                logger.info(f"Installed world {module_name} has version {version_str}")
                                break
                        else:
                            logger.info(f"Installed world {module_name} found but no version in pip show output")
                    else:
                        # Package not installed
                        logger.info(f"World {module_name} is not installed")
                except (subprocess.TimeoutExpired, Exception) as e:
                    logger.warning(f"Failed to check installed version for {module_name} using pip show: {e}")
                
                # Compare versions: install if new_version > installed_version
                # According to spec: "An APWorld without a world_version is always treated as older than one with a version"
                if new_version is None and installed_version is not None:
                    logger.info(f"There is a custom apworld file with no world version specified, please remove it from your custom_worlds directory.")
                elif installed_version is None or (new_version is not None and new_version > installed_version):
                    if installed_version is not None and new_version is not None:
                        uninstall_worlds([package_name])
                        logger.info(f"New version {new_version.as_simple_string()} > installed {installed_version.as_simple_string()}, uninstalling old version so new version will be picked up.")
                else:
                    logger.info(f"There is a custom apworld file with an older version than what is installed. Please remove it from your custom_worlds directory.")

            except Exception as e:
                logger.warning(f"Failed to check versions for APworld {world}: {e}")
    else:
        for wheel in worlds_files["wheels"]:
            logger.info(f"Installing wheel: {wheel}")
            executable_args = _uv_pip("install", wheel, "--upgrade")
            result = _uv_run(executable_args, timeout=30)
            if result.returncode != 0:
                logger.warning(f"Failed to install wheel {wheel}")
            else:
                logger.info(f"Successfully installed wheel {wheel}")


def update_requirements(needed_packages: List[str]) -> None:
    """
    Update packages from requirements.txt files and install worlds.

    uv's resolver respects each requirement's version specifier on its own; we don't
    pre-parse with `packaging`. When `needed_packages` is empty, upgrade everything in
    each requirements file. Otherwise, upgrade only the named entries.
    """
    if is_frozen():
        # Base requirements are baked into the frozen bundle at build time; the runtime
        # venv only carries additional packages (worlds, mwgg_igdb). Worlds in
        # `needed_packages` are still handled below.
        worlds_to_install = [pkg for pkg in needed_packages if pkg.startswith("worlds") or pkg.startswith("mwgg")]
        if worlds_to_install:
            logger.info(f"Installing/updating worlds: {worlds_to_install}")
            install_worlds(worlds_to_install)
        return

    update_all = len(needed_packages) == 0

    for req_file in requirements_files:
        if not req_file.exists():
            logger.warning(f"Requirements file not found: {req_file}")
            continue

        logger.debug(f"Processing requirements from: {req_file}")
        if update_all:
            executable_args = _uv_pip("install", "--upgrade", "-r", str(req_file))
        else:
            # Resolve the whole file so transitive caps and sibling constraints are
            # respected; --upgrade-package targets only the dists check_for_updates
            # flagged. The previous per-spec-line `pip install --upgrade <req_line>`
            # ignored sibling constraints and oscillated against check_requirements_satisfied.
            req_pkg_names = {
                re.split(r'[<>=!~;@\s]', line, maxsplit=1)[0].strip()
                for line in parse_requirements_file(req_file)
            }
            relevant = [p for p in needed_packages if p in req_pkg_names]
            if not relevant:
                continue
            executable_args = _uv_pip("install", "-r", str(req_file))
            for pkg in relevant:
                executable_args.extend(["--upgrade-package", pkg])

        result = _uv_run(executable_args, timeout=30)
        if result.returncode != 0:
            logger.warning(f"Failed to install/update from {req_file.name}")

    # Handle worlds (these are not in requirements.txt files)
    worlds_to_install = [pkg for pkg in needed_packages if pkg.startswith("worlds") or pkg.startswith("mwgg")]
    if worlds_to_install:
        logger.info(f"Installing/updating worlds: {worlds_to_install}")
        install_worlds(worlds_to_install)


def check_requirements_satisfied(yes: bool = False) -> bool:
    """
    Ensure all requirements files are satisfied. Returns True on success.

    With uv this is fast and idempotent — install runs unconditionally; if everything
    is already present, uv reports "Audited N packages" and exits in milliseconds.
    """
    if is_frozen():
        # Base requirements are baked into the frozen bundle at build time.
        return True
    for req_file in requirements_files:
        if not req_file.exists():
            logger.warning(f"Requirements file not found: {req_file}")
            continue
        logger.info(f"Ensuring requirements from {req_file.name} are satisfied")
        result = _uv_run(_uv_pip("install", "-r", str(req_file)), timeout=30)
        if result.returncode != 0:
            logger.warning(f"Failed to install requirements from {req_file.name}: {result.stderr}")
            return False
    return True


@contextlib.contextmanager
def _install_lock():
    """Serialize the top-level update pipeline across processes.

    Frozen clients and webhost workers share the runtime worlds venv. Direct
    install_worlds() callers remain unlocked unless contention shows up there.
    """
    lock_dir = install_path().parent if use_worlds_venv() else Path(tempfile.gettempdir())
    lock_dir.mkdir(parents=True, exist_ok=True)
    lock_path = lock_dir / ".mwgg-install.lock"
    with open(lock_path, "a+b") as lock_file:
        lock_file.seek(0, os.SEEK_END)
        if lock_file.tell() == 0:
            lock_file.write(b"\0")
            lock_file.flush()

        if is_windows():
            import msvcrt
            lock_file.seek(0)
            while True:
                try:
                    msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
                    break
                except OSError as e:
                    if (
                        e.errno not in (errno.EACCES, errno.EDEADLK)
                        and getattr(e, "winerror", None) not in (32, 33)
                    ):
                        raise
                    time.sleep(0.25)
            try:
                yield
            finally:
                lock_file.seek(0)
                msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
        else:
            # fcntl is Unix-only and this branch only runs off-Windows, but the gate
            # type-checks with pythonPlatform=Windows (where typeshed hides flock/LOCK_*),
            # so these are platform false positives, not real attribute errors.
            import fcntl
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
            try:
                yield
            finally:
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]


def update(yes: bool = True, force: bool = False, worlds: Optional[List[str]] = None) -> None:
    """
    Main update function.

    Args:
        yes: Answer yes to all prompts
        force: Force update without checking
        worlds: List of specific worlds to update

    Returns:
        None
    """
    if _skip_all_installs():
        return
    with _install_lock():
        _update_locked(yes=yes, force=force, worlds=worlds)


def _update_locked(yes: bool, force: bool, worlds: Optional[List[str]]) -> None:
    if _skip_all_installs():
        return
    # Install/refresh mwgg_igdb upfront
    install_mwgg_igdb(upgrade=True)

    if worlds:
        install_worlds(worlds, update=force)
        return

    if is_frozen():
        if (exe_dir / "custom_wheels").exists():
            logger.debug("Custom Worlds found, checking...")
            update_world_from_package()
        updates = check_for_updates(worlds_only=True)
        if updates:
            restart_needed = install_worlds(updates)
            if restart_needed:
                # Library updates were staged, need to restart
                from Utils import exit_restart_for_update
                exit_restart_for_update()
        else:
            logger.debug("No updates found.")
    global update_ran

    if update_ran:
        return

    update_ran = True

    if force:
        logger.debug("Force update requested - skipping update checks")
        # Force mode updates all requirements and worlds
        update_requirements([])  # Empty list means update all
    
    # Check for available updates
    logger.debug("Checking for available updates...")
    available_updates = check_for_updates()
    
    if available_updates:
        logger.debug(f"Found updates for: {available_updates}")
        if not yes:
            confirm("Updates available. Press enter to continue with updates.")
    else:
        logger.debug("No updates found.")
    
    # Check if requirements are satisfied
    logger.debug("Checking if all requirements are satisfied...")
    if not check_requirements_satisfied(yes=yes):
        logger.debug("Installing missing requirements...")
        update_requirements([])  # Empty list means update all missing requirements
    
    # Update packages that need updates (including worlds)
    if available_updates:
        logger.debug("Updating packages that need updates...")
        update_requirements(available_updates)
    
    logger.debug("Update process completed.")


class RestartException(Exception):
    """Exception raised when a restart is needed."""
    pass

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='Install archipelago requirements')
    parser.add_argument('-y', '--yes', dest='yes', action='store_true', 
                       help='answer "yes" to all questions')
    parser.add_argument('-f', '--force', dest='force', action='store_true', 
                       help='force update')
    parser.add_argument('-a', '--append', nargs="*", dest='additional_requirements',
                       help='List paths to additional requirement files.')
    parser.add_argument('-w', '--worlds', nargs="*", dest='worlds',
                       help='List of worlds to update.')
    
    args = parser.parse_args()

    # Standalone always pulls fresh; the stamp written here suppresses the
    # throttled install_mwgg_igdb(upgrade=True) calls inside the subsequent update().
    install_mwgg_igdb(upgrade=True, force=True)

    if args.additional_requirements:
        requirements_files.update([Path(req) for req in args.additional_requirements])

    if args.worlds:
        update(args.yes, args.force, args.worlds)
    else:
        update(args.yes, args.force)