#!/usr/bin/env python3
"""Stage A of the perf-seed workflow.

Fetches real module data from a configured S3 bucket via the hub hydration
pipeline and preserves the decompressed server-state trees to disk so a
later stage can replay them into a local MinIO backend.

Flow:
  1. Start a hub against real S3 (read-only SSO credentials), dehydration off,
     hibernated-watchdog timeout set huge so modules are not auto-deprovisioned
     mid-flow.
  2. Provision N modules (first N moduleIds from the bucket listing, filtered
     to UUID-shaped folders).
  3. Wait until every module reaches 'provisioned'.
  4. Hibernate every module (child processes exit, state dir intact and no
     remaining writer while we copy).
  5. Wait until every module reaches 'hibernated'.
  6. Copy <hub-data-dir>/servers/<moduleid>/ -> <snapshot-dir>/<moduleid>/
     with the hub still running - hibernate took the writers offline and the
     watchdog is muted.
  7. Shut the hub down via 'zen down --pid <hub pid>'.

Required environment variables (or pass via CLI flags):
  ZEN_PERF_S3_URI       e.g. s3://your-bucket/optional-prefix/
  ZEN_PERF_AWS_PROFILE  AWS SSO profile name configured with read access
  ZEN_PERF_AWS_REGION   defaults to us-east-1

Requirements:
  pip install boto3
  aws CLI v2 with the SSO profile configured
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional

_EXE_SUFFIX = ".exe" if sys.platform == "win32" else ""


def _rmtree_robust(path) -> None:
    """shutil.rmtree with a Windows-friendly retry for read-only files."""
    import os as _os
    import stat as _stat
    def _onerror(func, p, exc_info):
        try:
            _os.chmod(p, _stat.S_IWRITE)
            func(p)
        except Exception:
            pass
    # onexc was introduced in py3.12; fall back to onerror on older versions.
    if sys.version_info >= (3, 12):
        shutil.rmtree(path, onexc=lambda func, p, exc: _onerror(func, p, (type(exc), exc, exc.__traceback__)))
    else:
        shutil.rmtree(path, onerror=_onerror)


_DEFAULT_S3_URI = os.environ.get("ZEN_PERF_S3_URI", "")
_DEFAULT_AWS_PROFILE = os.environ.get("ZEN_PERF_AWS_PROFILE", "")
_DEFAULT_AWS_REGION = os.environ.get("ZEN_PERF_AWS_REGION", "us-east-1")

# Matches UUID-shaped module IDs (UUID-ish with mixed 4-char groups). Filters
# out stray top-level keys (readmes, test scratches) that aren't real modules.
_MODULEID_RE = re.compile(r"^[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}$")


# ---------------------------------------------------------------------------
# zenserver binary discovery - prefer release, fall back to debug.
# ---------------------------------------------------------------------------

def _find_zenserver(override: Optional[str]) -> Path:
    if override:
        p = Path(override) / f"zenserver{_EXE_SUFFIX}"
        if not p.exists():
            sys.exit(f"zenserver not found at {p}")
        return p

    script_dir = Path(__file__).resolve().parent
    repo_root = script_dir.parent.parent
    for mode in ("release", "debug"):
        for plat in (("windows", "x64"), ("linux", "x86_64"), ("macosx", "x86_64")):
            p = repo_root / "build" / plat[0] / plat[1] / mode / f"zenserver{_EXE_SUFFIX}"
            if p.exists():
                return p
    sys.exit("zenserver executable not found under build/. Build it or pass --zenserver-dir.")


def _find_zen(zenserver_exe: Path) -> Path:
    p = zenserver_exe.parent / f"zen{_EXE_SUFFIX}"
    if not p.exists():
        sys.exit(f"zen CLI not found at {p} (used for graceful hub shutdown)")
    return p


# ---------------------------------------------------------------------------
# AWS SSO credential resolution
# ---------------------------------------------------------------------------

def _require_boto3():
    try:
        import boto3  # type: ignore[import-not-found]
        import botocore.exceptions  # type: ignore[import-not-found]
    except ImportError:
        sys.exit("[aws] boto3 is required. Install it with: pip install boto3")
    return boto3, botocore.exceptions


def _sso_login(profile: str) -> None:
    print(f"[aws] running 'aws sso login --profile {profile}'...")
    rc = subprocess.call(["aws", "sso", "login", "--profile", profile])
    if rc != 0:
        sys.exit(f"[aws] 'aws sso login' failed with rc={rc}")


def _get_session(profile: str):
    boto3, exc_mod = _require_boto3()

    def _load_frozen():
        session = boto3.Session(profile_name=profile)
        creds = session.get_credentials()
        if creds is None:
            return session, None
        return session, creds.get_frozen_credentials()

    try:
        session, frozen = _load_frozen()
        if frozen is not None and frozen.access_key and frozen.secret_key:
            return session, frozen
    except exc_mod.ProfileNotFound:
        sys.exit(f"[aws] profile '{profile}' not found in ~/.aws/config")
    except Exception as e:
        print(f"[aws] initial credential load failed: {e}")

    _sso_login(profile)

    session, frozen = _load_frozen()
    if frozen is None or not frozen.access_key:
        sys.exit("[aws] could not resolve credentials after sso login")
    return session, frozen


def _parse_s3_uri(uri: str) -> tuple[str, str]:
    if not uri.startswith("s3://"):
        sys.exit(f"[aws] invalid S3 URI: {uri}")
    rest = uri[len("s3://"):]
    if "/" in rest:
        bucket, prefix = rest.split("/", 1)
    else:
        bucket, prefix = rest, ""
    return bucket, prefix


def _list_module_ids(session, bucket: str, prefix: str, region: str, limit: int) -> list[str]:
    """List UUID-shaped module folders under the bucket and return the `limit`
    most recently active ones, ranked by the LastModified of each module's
    `incremental-state.cbo` (newest first - these were last dehydrated most
    recently, which is the closest proxy for "most recently accessed").

    Falls back to listing order for any folder whose state file can't be
    HEADed (missing / 403 / transient error).
    """
    s3 = session.client("s3", region_name=region)
    prefix_norm = prefix if (not prefix or prefix.endswith("/")) else prefix + "/"

    # 1. Enumerate every UUID-shaped folder.
    paginator = s3.get_paginator("list_objects_v2")
    candidates: list[str] = []
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix_norm, Delimiter="/"):
        for cp in page.get("CommonPrefixes", []) or []:
            folder = cp.get("Prefix", "")[len(prefix_norm):].rstrip("/")
            if folder and _MODULEID_RE.match(folder):
                candidates.append(folder)
    print(f"[s3] {len(candidates)} module folders match UUID shape; ranking by state.cbo LastModified...")

    # 2. HEAD each module's state file in parallel. Missing/failed HEADs land
    #    at the tail via a sentinel epoch 0 timestamp.
    from datetime import datetime, timezone
    epoch = datetime(1970, 1, 1, tzinfo=timezone.utc)

    def _state_mtime(mid: str) -> datetime:
        key = f"{prefix_norm}{mid}/incremental-state.cbo"
        try:
            resp = s3.head_object(Bucket=bucket, Key=key)
            return resp.get("LastModified", epoch)
        except Exception:
            return epoch

    with ThreadPoolExecutor(max_workers=50) as pool:
        times = list(pool.map(_state_mtime, candidates))

    # 3. Sort descending (newest first). Folders without a state file sink.
    ranked = sorted(zip(candidates, times), key=lambda x: x[1], reverse=True)
    missing = sum(1 for _, t in ranked if t == epoch)
    if missing:
        print(f"[s3] {missing}/{len(ranked)} modules have no incremental-state.cbo (sorted to tail)")
    return [mid for mid, _ in ranked[:limit]]


# ---------------------------------------------------------------------------
# Hub lifecycle
# ---------------------------------------------------------------------------

def _start_hub(
    zenserver_exe: Path,
    data_dir: Path,
    port: int,
    log_file: Path,
    instance_limit: int,
    extra_args: list[str],
    extra_env: dict[str, str],
) -> tuple[subprocess.Popen, object]:
    data_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        str(zenserver_exe),
        "hub",
        "--enable-execution-history=false",
        f"--data-dir={data_dir}",
        f"--port={port}",
        "--hub-instance-corelimit=4",
        "--hub-provision-disk-limit-percent=99",
        "--hub-provision-memory-limit-percent=80",
        f"--hub-instance-limit={instance_limit}",
        # Seeding is not a perf-measurement path - we want it as fast as the
        # host can manage. Let the hub go wide on both provisioning and
        # hydration thread pools rather than matching prod limits.
        "--hub-instance-provision-threads=64",
        "--hub-hydration-threads=64",
        # With 1000 modules the seeding flow runs for 20+ minutes. Extend BOTH
        # watchdog inactivity timers so early-provisioned modules do not get
        # auto-deprovisioned while we're still hydrating the tail (hibernated
        # default 1800s, provisioned default 600s - the latter is the one that
        # bit us on the 1000-module run and dropped 335 modules).
        "--hub-watchdog-provisioned-inactivity-timeout-seconds=86400",
        "--hub-watchdog-hibernated-inactivity-timeout-seconds=86400",
    ] + extra_args

    env = os.environ.copy()
    env.update(extra_env)

    popen_kwargs: dict = {}
    if sys.platform == "win32":
        popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
    log_handle = log_file.open("wb")
    try:
        proc = subprocess.Popen(
            cmd, env=env, stdout=log_handle, stderr=subprocess.STDOUT,
            **popen_kwargs,
        )
    except Exception:
        log_handle.close()
        raise
    print(f"[hub] started (pid {proc.pid}), log: {log_file}")
    return proc, log_handle


def _wait_for_hub(proc: subprocess.Popen, port: int, timeout_s: float = 100.0) -> None:
    deadline = time.monotonic() + timeout_s
    req = urllib.request.Request(f"http://localhost:{port}/hub/status",
                                 headers={"Accept": "application/json"})
    while time.monotonic() < deadline:
        if proc.poll() is not None:
            sys.exit(f"[hub] process exited unexpectedly (rc={proc.returncode}) - "
                     f"is another zenserver already running on port {port}?")
        try:
            with urllib.request.urlopen(req, timeout=2):
                print("[hub] ready")
                return
        except Exception:
            time.sleep(0.2)
    sys.exit(f"[hub] timed out waiting for readiness after {timeout_s}s")


def _zen_down_hub(zen_exe: Path, hub_proc: subprocess.Popen, timeout_s: float = 300.0) -> None:
    """Shut the hub down via 'zen down --pid <pid>'. zen down performs a
    proper server-initiated shutdown (signals the shutdown event, waits for
    the server to drain) rather than a blunt kill."""
    if hub_proc.poll() is not None:
        return
    pid = hub_proc.pid
    print(f"[hub] zen down --pid {pid}")
    rc = subprocess.call([str(zen_exe), "down", "--pid", str(pid), "--force"])
    if rc != 0:
        print(f"[hub] zen down returned rc={rc}; waiting for exit anyway")
    try:
        hub_proc.wait(timeout=timeout_s)
    except subprocess.TimeoutExpired:
        print(f"[hub] did not exit after {timeout_s}s, killing")
        hub_proc.kill()
        hub_proc.wait()


# ---------------------------------------------------------------------------
# Hub HTTP API
# ---------------------------------------------------------------------------

def _hub_post(port: int, path: str, timeout_s: float = 60.0) -> tuple[int, dict]:
    url = f"http://localhost:{port}{path}"
    req = urllib.request.Request(url, data=b"{}", method="POST",
                                 headers={"Content-Type": "application/json",
                                          "Accept": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            try:
                body = json.loads(resp.read())
            except Exception:
                body = {}
            return resp.status, body
    except urllib.error.HTTPError as e:
        try:
            body = json.loads(e.read())
        except Exception:
            body = {}
        return e.code, body
    except Exception as e:
        return 0, {"error": str(e)}


def _hub_module_states(port: int, timeout_s: float = 10.0) -> Optional[dict[str, str]]:
    url = f"http://localhost:{port}/hub/status"
    req = urllib.request.Request(url, headers={"Accept": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            data = json.loads(resp.read())
    except Exception:
        return None
    return {m["moduleId"]: m.get("state", "") for m in (data.get("modules") or []) if m.get("moduleId")}


def _fan_out_post(
    pool: ThreadPoolExecutor,
    port: int,
    module_ids: list[str],
    verb: str,
) -> tuple[list[str], list[tuple[str, int, dict]]]:
    """POST /hub/modules/<id>/<verb> concurrently. Returns (accepted, failures)."""
    futures = {
        mid: pool.submit(_hub_post, port, f"/hub/modules/{mid}/{verb}")
        for mid in module_ids
    }
    accepted: list[str] = []
    failures: list[tuple[str, int, dict]] = []
    for mid, fut in futures.items():
        status, body = fut.result()
        if status in (200, 202):
            accepted.append(mid)
        else:
            failures.append((mid, status, body))
    return accepted, failures


# Any of these means the module is done transitioning and will not reach the
# target state on its own (without a retry we control).
_FAILED_STATES = {"crashed", "unprovisioned"}


def _wait_for_state(
    port: int,
    module_ids: list[str],
    target_state: str,
    timeout_s: float,
    label: str,
) -> tuple[list[str], list[str], dict[str, str]]:
    """Poll hub status until every module hits target_state, fails, or times out.

    Returns (stuck, failed, last_states). 'stuck' = still mid-transition when
    we timed out. 'failed' = hit an _FAILED_STATES value such as 'crashed'.
    """
    deadline = time.monotonic() + timeout_s
    remaining = set(module_ids)
    failed: list[str] = []
    last_states: dict[str, str] = {mid: "" for mid in module_ids}

    while remaining and time.monotonic() < deadline:
        states = _hub_module_states(port)
        if states is not None:
            for mid in list(remaining):
                s = states.get(mid, "")
                last_states[mid] = s
                if s == target_state:
                    remaining.discard(mid)
                elif s in _FAILED_STATES and target_state not in _FAILED_STATES:
                    remaining.discard(mid)
                    failed.append(mid)
        done = len(module_ids) - len(remaining)
        print(f"[{label}] {done}/{len(module_ids)} '{target_state}' ({len(failed)} failed)...", end="\r")
        time.sleep(2.0)

    print()
    return list(remaining), failed, last_states


# ---------------------------------------------------------------------------
# Snapshot copy (run AFTER hub shutdown so there's no concurrent writer)
# ---------------------------------------------------------------------------

def _copy_snapshot(src_root: Path, dst_root: Path, module_ids: list[str]) -> tuple[int, int, int]:
    """Copy src_root/<mid>/* to dst_root/<mid>/* for each module.

    Returns (modules_copied, files_copied, bytes_copied).
    Replaces only the specific per-module subdirs; never touches siblings.
    """
    dst_root.mkdir(parents=True, exist_ok=True)
    modules_copied = 0
    files_copied = 0
    bytes_copied = 0

    for i, mid in enumerate(module_ids, 1):
        src = src_root / mid
        dst = dst_root / mid
        if not src.is_dir():
            print(f"[snapshot] WARNING: source missing for {mid}: {src}")
            continue
        if dst.exists():
            _rmtree_robust(dst)
        shutil.copytree(src, dst, symlinks=False, dirs_exist_ok=False)
        modules_copied += 1
        for root, _dirs, files in os.walk(dst):
            for f in files:
                p = Path(root) / f
                try:
                    bytes_copied += p.stat().st_size
                except OSError:
                    pass
                files_copied += 1
        if i % 25 == 0 or i == len(module_ids):
            print(f"[snapshot] {i}/{len(module_ids)} modules copied "
                  f"({files_copied:,} files, {bytes_copied/1024/1024:.1f} MB)")

    return modules_copied, files_copied, bytes_copied


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--hub-data-dir", default="E:/Dev/zen-perf-seed/hub-a",
                        help="Hub --data-dir (default: E:/Dev/zen-perf-seed/hub-a)")
    parser.add_argument("--snapshot-dir", default="E:/Dev/zen-perf-seed/s3-snapshot",
                        help="Destination for per-module server-state trees (default: E:/Dev/zen-perf-seed/s3-snapshot)")
    parser.add_argument("--port", type=int, default=8558,
                        help="Hub HTTP port (default: 8558)")
    parser.add_argument("--module-count", type=int, default=1000,
                        help="Number of modules to snapshot (default: 300)")
    parser.add_argument("--workers", type=int, default=50,
                        help="Concurrent HTTP workers (default: 50)")
    parser.add_argument("--poll-timeout", type=float, default=1800.0,
                        help="Max seconds to wait for provision or hibernate to finish (default: 1800)")
    parser.add_argument("--zenserver-dir",
                        help="Directory containing zenserver executable (auto-detected by default)")
    args = parser.parse_args()

    hub_data_dir = Path(args.hub_data_dir).resolve()
    snapshot_dir = Path(args.snapshot_dir).resolve()

    # Safety: snapshot-dir is the preserved output. It must not overlap the
    # hub data-dir in either direction so nothing the hub writes can clobber
    # the preserved state.
    if (snapshot_dir == hub_data_dir
            or snapshot_dir in hub_data_dir.parents
            or hub_data_dir in snapshot_dir.parents):
        sys.exit(f"[setup] snapshot-dir ({snapshot_dir}) and hub-data-dir ({hub_data_dir}) must be disjoint")

    s3_uri = os.environ.get("S3_URI", _DEFAULT_S3_URI)
    aws_profile = os.environ.get("AWS_PROFILE", _DEFAULT_AWS_PROFILE)
    aws_region = os.environ.get("AWS_REGION", _DEFAULT_AWS_REGION)
    if not s3_uri:
        sys.exit("[setup] S3 URI not set. Set ZEN_PERF_S3_URI (or S3_URI) to a bucket like s3://your-bucket/")
    if not aws_profile:
        sys.exit("[setup] AWS profile not set. Set ZEN_PERF_AWS_PROFILE (or AWS_PROFILE) to your SSO profile name")

    hub_log = hub_data_dir / "hub.log"

    zenserver_exe = _find_zenserver(args.zenserver_dir)
    zen_exe = _find_zen(zenserver_exe)
    zenserver_mode = "release" if "release" in zenserver_exe.parts else ("debug" if "debug" in zenserver_exe.parts else "?")
    print(f"[setup] build mode:   {zenserver_mode}")
    print(f"[setup] zenserver:    {zenserver_exe}")
    print(f"[setup] zen cli:      {zen_exe}")
    print(f"[setup] S3 URI:       {s3_uri}")
    print(f"[setup] profile:      {aws_profile}")
    print(f"[setup] hub-data-dir: {hub_data_dir}")
    print(f"[setup] snapshot-dir: {snapshot_dir}")

    session, frozen = _get_session(aws_profile)
    print(f"[aws] credentials resolved (key prefix {frozen.access_key[:6]}..., session-token={'yes' if frozen.token else 'no'})")

    bucket, prefix = _parse_s3_uri(s3_uri)
    print(f"[s3] listing folders under bucket='{bucket}' prefix='{prefix}'...")
    module_ids = _list_module_ids(session, bucket, prefix, aws_region, args.module_count)
    print(f"[s3] selected {len(module_ids)} module folders (UUID-shaped)")
    if len(module_ids) < args.module_count:
        print(f"[s3] WARNING: asked for {args.module_count} modules, only {len(module_ids)} matched the UUID filter")

    if not module_ids:
        sys.exit("[s3] no module folders found, aborting")

    aws_env = {
        "AWS_ACCESS_KEY_ID": frozen.access_key,
        "AWS_SECRET_ACCESS_KEY": frozen.secret_key,
    }
    if frozen.token:
        aws_env["AWS_SESSION_TOKEN"] = frozen.token

    hub_data_dir.mkdir(parents=True, exist_ok=True)
    config_path = hub_data_dir / "hydration_config.json"
    config_path.write_text(
        json.dumps({
            "type": "s3",
            "settings": {"uri": s3_uri, "region": aws_region},
        }),
        encoding="ascii",
    )
    hub_extra_args = [
        f"--hub-hydration-target-config={config_path}",
        "--hub-enable-dehydration=false",
    ]

    hub_proc: Optional[subprocess.Popen] = None
    hub_log_handle = None
    exit_code = 0

    try:
        hub_instance_limit = max(args.module_count + 10, 500)
        hub_proc, hub_log_handle = _start_hub(
            zenserver_exe, hub_data_dir, args.port, hub_log,
            hub_instance_limit, hub_extra_args, aws_env,
        )
        _wait_for_hub(hub_proc, args.port)

        t_start = time.monotonic()

        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            # --- Provision ---
            print(f"[provision] firing {len(module_ids)} provision requests...")
            t0 = time.monotonic()
            accepted, failures = _fan_out_post(pool, args.port, module_ids, "provision")
            print(f"[provision] accepted={len(accepted)}, failed={len(failures)} (fan-out {time.monotonic()-t0:.1f}s)")
            for mid, status, body in failures[:10]:
                print(f"[provision] FAILED {mid}: status={status} body={body}")
            if not accepted:
                sys.exit("[provision] nothing accepted, aborting")

            stuck, failed, last_states = _wait_for_state(args.port, accepted, "provisioned", args.poll_timeout, "provision")
            prov_done = len(accepted) - len(stuck) - len(failed)
            print(f"[provision] complete: {prov_done}/{len(accepted)} provisioned, {len(failed)} failed, {len(stuck)} stuck "
                  f"({time.monotonic()-t0:.1f}s)")
            if failed:
                for mid in failed[:10]:
                    print(f"[provision]   FAILED {mid}: last state='{last_states.get(mid, '')}'")
                exit_code = 1
            if stuck:
                for mid in stuck[:10]:
                    print(f"[provision]   STUCK {mid}: last state='{last_states.get(mid, '')}'")
                exit_code = 1
            accepted = [m for m in accepted if m not in set(stuck) and m not in set(failed)]

            if not accepted:
                sys.exit("[provision] nothing successfully provisioned, aborting")

            # --- Hibernate ---
            print(f"[hibernate] firing {len(accepted)} hibernate requests...")
            t0 = time.monotonic()
            hib_accepted, hib_failures = _fan_out_post(pool, args.port, accepted, "hibernate")
            print(f"[hibernate] accepted={len(hib_accepted)}, failed={len(hib_failures)} (fan-out {time.monotonic()-t0:.1f}s)")
            for mid, status, body in hib_failures[:10]:
                print(f"[hibernate] FAILED {mid}: status={status} body={body}")
                exit_code = 1

            stuck_hib, failed_hib, last_states_hib = _wait_for_state(args.port, hib_accepted, "hibernated", args.poll_timeout, "hibernate")
            hib_done = len(hib_accepted) - len(stuck_hib) - len(failed_hib)
            print(f"[hibernate] complete: {hib_done}/{len(hib_accepted)} hibernated, {len(failed_hib)} failed, {len(stuck_hib)} stuck "
                  f"({time.monotonic()-t0:.1f}s)")
            if failed_hib or stuck_hib:
                exit_code = 1
                for mid in (failed_hib + stuck_hib)[:10]:
                    print(f"[hibernate]   not-hibernated {mid}: last state='{last_states_hib.get(mid, '')}'")

        # --- Copy snapshots while hub is still running. All instances are
        #     hibernated (no writers), watchdog-hibernated-timeout is 86400s
        #     (no auto-deprovision), hub is only touching its own metadata
        #     outside servers/<mid>/. Safe. ---
        copy_src = hub_data_dir / "servers"
        to_copy = [m for m in hib_accepted if m not in set(stuck_hib) and m not in set(failed_hib)]
        print(f"[snapshot] copying {len(to_copy)} module trees from {copy_src} -> {snapshot_dir}")
        t0 = time.monotonic()
        modules_copied, files_copied, bytes_copied = _copy_snapshot(copy_src, snapshot_dir, to_copy)
        print(f"[snapshot] copied {modules_copied} modules, {files_copied:,} files, {bytes_copied/1024/1024:.1f} MB "
              f"({time.monotonic()-t0:.1f}s)")

        if modules_copied < len(to_copy):
            print(f"[snapshot] WARNING: only {modules_copied}/{len(to_copy)} trees copied")
            exit_code = 1

        # --- Graceful hub shutdown via 'zen down' ---
        _zen_down_hub(zen_exe, hub_proc)
        hub_proc = None
        if hub_log_handle is not None:
            hub_log_handle.close()
            hub_log_handle = None

        print(f"[summary] stage A total elapsed: {time.monotonic()-t_start:.1f}s, exit={exit_code}")
        print(f"[summary] snapshot available at: {snapshot_dir}")

    finally:
        if hub_proc is not None and hub_proc.poll() is None:
            # Fallback: something aborted before the 'zen down' path ran.
            print("[hub] force-terminating (zen down was not invoked)")
            hub_proc.terminate()
            try:
                hub_proc.wait(timeout=60)
            except subprocess.TimeoutExpired:
                hub_proc.kill()
                hub_proc.wait()
        if hub_log_handle is not None:
            hub_log_handle.close()

    return exit_code


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