"""Native payload staging through PM's existing package authority."""
from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import time
from contextlib import contextmanager
from dataclasses import asdict
from pathlib import Path

from pm.install import _lockfile
from pm.lock import Facts
from pm.registry import get_package, walk
from pm.store import current_target

def _bundle_package_names() -> list[str]:
    names = [
        n
        for n in _lockfile().names()
        if not get_package(n).internal or n == "uv"
    ]
    if "python" not in names:
        names.append("python")
    return names


def _arch_guard(store_dir: Path) -> list[str]:
    """Every staged binary must be built for this machine's target — a
    payload staged with a mismatched interpreter or PATH tool ships an
    artifact that cannot run. Reads facts, probes each entry binary."""
    from pm.package import machine_matches_binary

    facts = Facts(store_dir / "facts.json")
    problems = []
    target = current_target()
    for name in _lockfile().names():
        package = get_package(name)
        fact = facts.get(name)
        if fact is None or "entry" not in fact:
            continue
        binary = package.binary(store_dir / fact["entry"], target)
        if binary is None or not binary.is_file():
            continue
        verdict = machine_matches_binary(binary, target)
        # A package that declares this target as emulated (x64 binary run
        # under Windows ARM64 built-in emulation) is fine with the x64 PE.
        if verdict is False and target not in package.emulated_arch_targets:
            problems.append(f"{name}: {binary.name} is not a {target} binary")
    return problems



from pm.uv_cache_prune import lock_package_names, prune_uv_cache_to_lock

__all__ = ["lock_package_names", "prune_uv_cache_to_lock", "stage_uv_cache"]


def stage_uv_cache(source: Path, destination: Path) -> None:
    """Ship everything a per-install venv rebuild resolves from offline.

    sdist source trees (including Rust target/ outputs) are build-only bulk —
    offline rebuilds install the built wheel entries, never recompile from
    source — so they stay out. Built wheel ZIPs (including the copies uv keeps
    beside sdist metadata.msgpack shards) ship: without them an offline
    rebuild falls back to wanting the sdist download and fails closed.
    """
    def ignore(directory: str, names: list[str]) -> set[str]:
        path = Path(directory)
        parts = path.relative_to(source).parts
        if not parts or not parts[0].startswith("sdists-v"):
            return set()
        # Source trees live under a revision selected by a sibling pointer.
        if "src" in names and (path / "src").is_dir() and any(
            (path.parent / pointer).is_file() for pointer in ("revision.http", "revision.rev")
        ):
            return {"src"}
        return set()

    shutil.copytree(source, destination, ignore=ignore)


def stage_pm_runtime(root: Path, python: Path, repo: Path, *, offline: bool = False,
                     cache: Path | None = None) -> None:
    """Publish the same PM dependency graph as source installs, ready offline."""
    from pm import stage_manager_runtime
    from scripts.bundles.payload import seal_pm_runtime

    destination = root / "pm-runtime"
    if destination.exists():
        shutil.rmtree(destination)
    stage_manager_runtime(python=python, destination=destination, project=repo / "pm", offline=offline, cache=cache)
    seal_pm_runtime(root, python)


def stage_native(args) -> int:
    """Isolate HOME and PM state, but retain the provider's reusable build cache."""
    from pm.packages import uv_cache_dir

    out = Path(args.out).resolve()
    out.mkdir(parents=True, exist_ok=True)
    (out / "manifest.json").unlink(missing_ok=True)
    root = Path(__file__).resolve().parents[2]
    # Resolve before HOME isolation so the build warms the cache CI saves.
    cache = Path(getattr(args, "cache", None) or os.environ.get("UV_CACHE_DIR") or uv_cache_dir()).resolve()
    base_env = dict(os.environ)
    if current_target() == "win32-arm64":
        from pm.native_build import prepare_windows_environment

        base_env = prepare_windows_environment(source=root, state=out.parent / ".build-deps", env=base_env)
    # Rustup resolves its installed toolchain under HOME unless these are explicit.
    # Preserve the build host's compiler and Cargo cache before isolating app state.
    for key, directory in (("CARGO_HOME", ".cargo"), ("RUSTUP_HOME", ".rustup")):
        base_env.setdefault(key, str(Path.home() / directory))
    with tempfile.TemporaryDirectory(prefix=".build-", dir=out) as work:
        env = {**base_env, "HOME": work, "USERPROFILE": work,
               "HERMES_HOME": str(Path(work) / ".hermes"),
               "HERMES_RUNTIME_DIR": str(out / "tools"),
               "HERMES_PYTHON_SRC_ROOT": str(root),
               "XDG_CACHE_HOME": str(Path(work) / "cache"),
               "XDG_CONFIG_HOME": str(Path(work) / "config"),
               "UV_CACHE_DIR": str(cache),
               "PYTHONPATH": os.pathsep.join([str(root), *filter(None, sys.path)])}
        command = [sys.executable, "-B", "-m", "scripts.bundles.native", "--out", str(out),
                   "--ref", args.ref or "HEAD", "--source", str(root)]
        for name, product in getattr(args, "frontends", {}).items():
            command += [f"--{name}", str(product)]
        return subprocess.run(command, cwd=root, env=env, check=False).returncode


@contextmanager
def _phase(label: str):
    """Name each staging phase and its duration: several run minutes without output."""
    print(f"  {label}…", flush=True)
    started = time.monotonic()
    yield
    print(f"  ✓ {label} ({time.monotonic() - started:.1f}s)", flush=True)


def prune_staged_store(store_dir: Path, names: list[str]) -> None:
    """Prune this build's store, never the ambient user's tool store."""
    facts = Facts(store_dir / "facts.json", strict=True)
    facts.retain({package.name for package in walk(names)})
    keep = facts.entries_in_use()
    for entry in store_dir.iterdir():
        if entry.is_dir() and not entry.name.startswith(".") and entry.name not in keep:
            shutil.rmtree(entry)


def prepare_native(*, out: Path, ref: str, source: Path, cache: Path,
                   tools: Path | None = None, env: dict | None = None) -> Path:
    """Prepare final payload dependencies inside the caller's isolated PM process.

    The caller supplies its compiler environment; only the full standalone
    wrapper provisions machine prerequisites and isolates HOME.
    """
    from scripts.bundles.native_prepared import preparation_lock, prepared_path

    out = Path(out).absolute()
    if out != out.resolve():
        raise ValueError("symlinked native output")
    out.mkdir(parents=True, exist_ok=True)
    with preparation_lock(out):
        prepared_path(out).unlink(missing_ok=True)
        (out / "manifest.json").unlink(missing_ok=True)
        return _prepare_native(out=out, ref=ref, source=Path(source).resolve(),
                               cache=Path(cache).absolute(), tools=tools, env=env)


def _prepare_native(*, out: Path, ref: str, source: Path, cache: Path,
                    tools: Path | None, env: dict | None) -> Path:
    from pm import paths
    from pm.package import InstallError

    store_dir = out / "tools"
    store_dir.mkdir(parents=True, exist_ok=True)
    repo_dir = out / "hermes-agent"
    from scripts.bundles.payload import INERT_SNAPSHOT_DIRS, snapshot
    revision = subprocess.check_output(
        ["git", "rev-parse", "--verify", f"{ref}^{{commit}}"], cwd=source,
        text=True, encoding="utf-8").strip()
    with _phase(f"staging repo snapshot ({ref})"):
        snapshot(source, revision, repo_dir, exclude=INERT_SNAPSHOT_DIRS)
    # PM's provider code reads its adjacent lock. Never combine that tool graph
    # with a revision selecting different pins.
    if (repo_dir / "pm/lock.json").read_bytes() != paths.lockfile_path().read_bytes():
        raise ValueError("selected revision's PM lock differs from the builder; use a checkout at that revision")
    build_env = os.environ if env is None else env

    names = [
        n for n in _bundle_package_names()
        if get_package(n).missing_reason(current_target()) is None
    ]
    from pm import prepare_tools, stage_tools

    with _phase(f"preparing {len(names)} pinned tools"):
        prepare_tools(names, out=Path(tools) if tools is not None else store_dir,
                      target=current_target(), cache=cache)
    if tools is not None:
        with _phase("copying verified tools into the payload store"):
            stage_tools(names, source_store=Path(tools), out=store_dir, target=current_target())

    # Prune the staged store BEFORE the venv sync and packaging: drop the
    # fetch-<sha> download-cache archives (needed only at install time — dead
    # weight in the shipped payload AND in the CI cache that restores this
    # dir) and any orphaned package versions left over from an older lock
    # the cache carried in. A lean staged store = a lean CI cache.

    # Only this build's store is ours to prune; machine-wide partials are not.
    # Cached facts may still name packages removed from the current selection.
    # Retain the dependency closure before using facts as the deletion roots.
    with _phase("pruning the staged tool store"):
        prune_staged_store(store_dir, names)


    facts = Facts(store_dir / "facts.json")
    python_fact = facts.get("python")
    if python_fact is None:
        raise InstallError("venv", "no staged interpreter to build on")
    python_bin = get_package("python").binary(
        store_dir / python_fact["entry"], current_target()
    )

    if python_bin is None:
        raise FileNotFoundError("staged Python executable is missing")
    with _phase("pm-runtime (independent locked dependencies)"):
        stage_pm_runtime(out, python_bin, repo_dir, cache=cache)

    # Build + sync INSIDE the staged repo: the editable project install
    # must point at the payload's own tree, not this checkout.
    venv_dir = out / "venv"
    if venv_dir.exists():
        shutil.rmtree(venv_dir)
    env = dict(build_env)
    from pm import build_environment

    # Cold native wheels need a larger budget than interactive installs.
    build_environment(source=repo_dir, python=python_bin, out=venv_dir,
                      env=env, cache=cache, all_extras=True, sealed=True, explicit=True,
                      timeout=2 * 60 * 60)
    print("✓ venv (all extras, on the staged interpreter)")

    # Inventory the staged interpreter before publishing the bundle contract.
    from pm.features import installed_extras, write_features

    features = installed_extras(repo_dir, venv_dir, python_exe=python_bin)
    write_features(features, out)

    # The payload venv was built from this exact lock and feature set. Publish
    # that state so the first runtime sync can reuse the shipped environment
    # instead of rebuilding it before any plugin dependencies are requested.
    from pm.packages import Venv

    facts.record_state("venv", Venv(repo_dir).expected_stamp(features, plugin_dirs=[]), features)
    print(f"✓ enabled-features.json ({len(features)} extras recorded)")

    # Ship the full uv cache (build-only sdist sources and wheel ZIPs are
    # already excluded by stage_uv_cache). Any per-install venv rebuild —
    # plugin extras, feature changes — resolves entirely from the shipped
    # wheels with zero network access, on every profile of this install.
    payload_cache = out / "uv-cache"
    if payload_cache.exists():
        shutil.rmtree(payload_cache, ignore_errors=True)
    src_cache = cache
    if src_cache.is_dir():
        print(f"  uv-cache: copying {src_cache} → payload...", flush=True)
        stage_uv_cache(src_cache, payload_cache)
        pruned = prune_uv_cache_to_lock(payload_cache, repo_dir)
        print(f"✓ uv-cache (lock-scoped: pruned {pruned} stale entries; offline rebuilds resolve from shipped wheels)", flush=True)
    else:
        raise InstallError("uv-cache", "runtime dependency cache is missing")

    bad = _arch_guard(store_dir)
    for line in bad:
        print(f"✗ arch: {line}")
    if bad:
        raise InstallError("tools", "native architecture verification failed")
    from scripts.bundles.payload import record_tools
    recorded = {name: fact["entry"] for name in names if (fact := facts.get(name)) and "entry" in fact}
    record_tools(out, paths.lockfile_path(), current_target(), recorded)
    from scripts.build.inputs import AgentInputs, RESOURCE_ENV, dependency_site
    from scripts.bundles.native_prepared import publish_prepared

    site = dependency_site(venv_dir, python_fact["version"], current_target())
    (site / "hermes-agent.pth").write_text(
        Path(os.path.relpath(repo_dir, site)).as_posix() + "\n", encoding="utf-8")
    from scripts.bundles.payload import relativize_links
    relativize_links(out)
    from scripts.bundles.bytecode import bake_bytecode
    baked = bake_bytecode(out, python_bin)
    print(f"✓ baked bytecode ({baked['modules']} modules, unchecked-hash, read-only caches)")
    inputs = AgentInputs(
        project=repo_dir / "pyproject.toml", code=repo_dir, repo="hermes-agent",
        placement="contained", target=current_target(), python=python_bin,
        site_packages=site, environment=venv_dir,
        tools=store_dir, pm_runtime=out / "pm-runtime", ref=ref,
        resources={name: repo_dir / name for name in RESOURCE_ENV},
        features=out / "enabled-features.json",
    )
    return publish_prepared(out, source, revision, inputs)


def finish_native(prepared: Path, frontends: dict[str, Path]) -> int:
    """Consume verified job-local paths. No dependency acquisition or repair."""
    from scripts.build.agent import assemble
    from scripts.build.inputs import AgentInputs
    from scripts.bundles.native_prepared import load_prepared, preparation_lock

    prepared = Path(prepared).absolute()
    if not prepared.name.endswith(".prepared.json") or not prepared.is_file():
        raise ValueError("native preparation is missing or invalid; run preparation again")
    out = prepared.with_name(prepared.name.removesuffix(".prepared.json"))
    with preparation_lock(out):
        (out / "manifest.json").unlink(missing_ok=True)
        inputs = load_prepared(prepared)
        values = asdict(inputs)
        values["frontends"] = {name: Path(path).absolute() for name, path in frontends.items()}
        print("assembling verified payload", flush=True)
        assemble(AgentInputs.from_dict(values), out)
    print(f"✓ manifest ({out / 'manifest.json'})")
    return 0


def _stage_native(args) -> int:
    from pm import paths
    from pm.features import FeatureProbeError
    from pm.package import InstallError

    try:
        prepared = prepare_native(
            out=Path(args.out), ref=args.ref or "HEAD",
            source=getattr(args, "source", None) or paths.repo_root(),
            cache=Path(getattr(args, "cache", None) or os.environ["UV_CACHE_DIR"]),
            tools=getattr(args, "tools", None),
        )
        return finish_native(prepared, getattr(args, "frontends", {}))
    except (InstallError, FeatureProbeError) as exc:
        print(f"✗ {exc}")
        return 1




def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", required=True)
    parser.add_argument("--ref", default="HEAD")
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--tui", type=Path)
    parser.add_argument("--web", type=Path)
    args = parser.parse_args()
    args.frontends = {name: path for name in ("tui", "web") if (path := getattr(args, name)) is not None}
    return _stage_native(args)


if __name__ == "__main__":
    raise SystemExit(main())