"""Update check foundation.

v1 only DETECTS and REPORTS whether a newer version is available on the
configured channel. It does NOT download, swap, or restart anything. The
production update pipeline is sketched below as no-op stubs with TODOs so the
structure is ready: download -> verify checksum -> stage -> atomic swap ->
health check -> rollback.
"""

from __future__ import annotations

import os
import re
import sys
import time
import json
import shutil
import tarfile
import hashlib
import subprocess
from pathlib import Path
from dataclasses import dataclass
from typing import Any

from api_client import ApiClient, ApiError
from config import Config
from identity import load_agent_token
from logger import get_logger

log = get_logger("update")


@dataclass
class UpdateInfo:
    current: str
    latest: str
    update_available: bool
    url: str | None = None
    sha256: str | None = None
    signature: str | None = None
    channel: str | None = None
    size_bytes: int | None = None


# RSA-2048 public key used to verify signatures on real (non-mock) release
# packages. This value is public by design -- it is baked into install.sh and
# verify-release.sh as well. The matching PRIVATE key must never appear in
# this repository; production signing reads it from RELEASE_PRIVATE_KEY /
# RELEASE_PRIVATE_KEY_FILE at release-build time only (see
# scripts/sign_release.py, scripts/build-release.sh).
#
# Rotated 2026-07-05: the previous key's private exponent had been committed
# to this repository (as "MOCK_PRIVATE_KEY_D") and is considered permanently
# compromised. This is a freshly generated keypair; its private half was
# generated outside this repo and has never been committed anywhere.
PUBLIC_KEY_N = 22153926292932439223494206857578992263368566205306726328685384489841974735455975021291694740520631615844482011371776971868501750352935831309311417721523617582807469656789022762787913062961956760093233770949538223047769903327425298895828401330698012217244681981952297803454582886661042248065952897616193885979581482910454765900722981305580684601338572270789120709521740056140363158200014805167871811459475812161877933802463443300930869637395955802230717162966225558744034920275206651643947411725623110190445098934844925287947024482792706657109971364650444199678321935099468386116294090959465201975791356140751205132389
PUBLIC_KEY_E = 65537

# --- TEST-ONLY keypair ------------------------------------------------------
# Freshly generated, completely unrelated to PUBLIC_KEY_N above. Used ONLY by
# the mock self-update path (config.api_mock=True) so local/dev/test runs can
# produce and verify a real signature without a network round-trip and
# without ever touching the production key. This keypair must NEVER be able
# to verify a signature made with the real production private key, and vice
# versa -- that separation is the whole point (see docs/RELEASE_GUIDE.md).
TEST_PUBLIC_KEY_N = 28646827674820102632485468682453933335873839565574061436680156800952295025605066935840504046373305152404813225628113814732746304519621111601146766265768158059942633167116191611210711374108250216771520057003471026969578224744445974499433563666264393001180694265213021183586331817644172433049512989048276755233125119872376930624920594665412217491038110898521045301477729466037114656409107636928879350154320310052855167891627760207329436234504276514399340587871990240801236776820779217276971344039452813538261693357072899848212896981926129805997234503946261059075001609646719074296882560113381787775959043897497491431497
TEST_PUBLIC_KEY_E = 65537
TEST_PRIVATE_KEY_D = 1230462485384112622037117877551731424088451689535544546504335587449543166411008490232861114950956772571517604256269594099099452938381882435223280696982427711655072895699102482346737759099664683464483100545718768648539949992456405056928231101828497891241949652205237570102316616059147434259035034624271771151872879659895641400236981461324886366735264043313908260327482596087072918969794616173813378808614790685017698845737964761081115908095249463200374621433204207358766588324694156973337030500699777348730759752120405947937086106376685473908954311692506733710832925701517030836504946302664575317558963983124672314093


def verify_signature(data: bytes, signature_hex: str, n: int = PUBLIC_KEY_N, e: int = PUBLIC_KEY_E) -> bool:
    """Verify standard RSASSA-PKCS1-v1_5 signature using the given public key.

    Defaults to the production public key; pass n=TEST_PUBLIC_KEY_N,
    e=TEST_PUBLIC_KEY_E to verify a mock/test signature instead.
    """
    try:
        if not signature_hex:
            return False
        sig_bytes = bytes.fromhex(signature_hex)
        if len(sig_bytes) != 256:
            return False
        s = int.from_bytes(sig_bytes, "big")
        m = pow(s, e, n)
        decrypted = m.to_bytes(256, "big")

        sha256_hash = hashlib.sha256(data).digest()
        # PKCS#1 v1.5 padding block for SHA-256 (256 bytes total)
        asn1_prefix = b"\x30\x31\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x01\x05\x00\x04\x20"
        expected = b"\x00\x01" + b"\xff" * 202 + b"\x00" + asn1_prefix + sha256_hash
        return decrypted == expected
    except Exception:
        return False



_SEMVER_RE = re.compile(
    r"^(?P<major>0|[1-9]\d*)\.(?P<minor>0|[1-9]\d*)\.(?P<patch>0|[1-9]\d*)"
    r"(?:-(?P<prerelease>[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?"
    r"(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?$"
)


def parse_semver(version: str) -> tuple[int, int, int, tuple[str, ...] | None]:
    """Parse a semver.org version string into (major, minor, patch, prerelease)."""
    m = _SEMVER_RE.match((version or "").strip())
    if not m:
        raise ValueError(f"invalid semantic version: {version!r}")
    major, minor, patch = int(m["major"]), int(m["minor"]), int(m["patch"])
    prerelease = tuple(m["prerelease"].split(".")) if m["prerelease"] else None
    return major, minor, patch, prerelease


def _prerelease_identifier_key(ident: str) -> tuple[int, Any]:
    # Per semver.org precedence rule 11: numeric identifiers always have
    # lower precedence than alphanumeric identifiers; numeric identifiers
    # compare numerically, alphanumeric identifiers compare lexically (ASCII).
    if ident.isdigit():
        return (0, int(ident))
    return (1, ident)


def compare_semver(a: str, b: str) -> int:
    """Return -1, 0, or 1 comparing semver strings a and b per semver.org precedence.

    Core versions compare numerically field-by-field; a version with a
    pre-release always has lower precedence than the same core version
    without one (e.g. 1.2.0-rc.1 < 1.2.0); pre-release identifiers compare
    left-to-right per semver.org rule 11 (e.g. beta.1 < beta.2 < rc.1).
    """
    a_major, a_minor, a_patch, a_pre = parse_semver(a)
    b_major, b_minor, b_patch, b_pre = parse_semver(b)

    for x, y in ((a_major, b_major), (a_minor, b_minor), (a_patch, b_patch)):
        if x != y:
            return -1 if x < y else 1

    if a_pre is None and b_pre is None:
        return 0
    if a_pre is None:
        return 1  # release always outranks a pre-release of the same core version
    if b_pre is None:
        return -1

    for x, y in zip(a_pre, b_pre):
        kx, ky = _prerelease_identifier_key(x), _prerelease_identifier_key(y)
        if kx != ky:
            return -1 if kx < ky else 1
    if len(a_pre) != len(b_pre):
        return -1 if len(a_pre) < len(b_pre) else 1
    return 0


_SUPPORTED_MANIFEST_VERSIONS = (1,)


def check(config: Config, api_client: ApiClient) -> UpdateInfo | None:
    """Fetch the channel manifest and compare against the running version."""
    try:
        manifest: dict[str, Any] = api_client.fetch_manifest(config.manifest_url)
    except ApiError as exc:
        log.warning("update check failed: %s", exc)
        return None

    manifest_version = manifest.get("manifest_version")
    if manifest_version not in _SUPPORTED_MANIFEST_VERSIONS:
        log.error(
            "rejecting manifest: unsupported manifest_version %r (supported: %s)",
            manifest_version, _SUPPORTED_MANIFEST_VERSIONS,
        )
        return None

    manifest_channel = manifest.get("channel")
    if manifest_channel != config.release_channel:
        log.error(
            "rejecting manifest: channel mismatch (agent is on '%s', manifest reports '%s')",
            config.release_channel, manifest_channel,
        )
        return None

    latest = manifest.get("version", "")
    try:
        available = compare_semver(latest, config.agent_version) > 0
    except ValueError as exc:
        log.error("rejecting manifest: %s", exc)
        return None

    sig = manifest.get("signature")
    if not sig and config.api_mock:
        sig = "mock_signature"
    manifest_file = manifest.get("file")
    download_url = None
    if manifest_file:
        download_url = f"{config.release_base.rstrip('/')}/{config.release_channel}/{manifest_file}"

    info = UpdateInfo(
        current=config.agent_version,
        latest=latest,
        update_available=available,
        url=download_url,
        sha256=manifest.get("sha256"),
        signature=sig,
        channel=manifest_channel,
        size_bytes=manifest.get("size_bytes"),
    )
    if available:
        log.info("update available: %s -> %s", info.current, info.latest)
    else:
        log.info("agent is up to date (%s)", info.current)
    return info


def swap_directories_atomic(src_dir: Path, dest_link: Path):
    if hasattr(os, "symlink"):
        try:
            temp_link = dest_link.parent / f"{dest_link.name}_temp"
            if temp_link.exists() or temp_link.is_symlink():
                temp_link.unlink()
            os.symlink(src_dir, temp_link)
            os.replace(temp_link, dest_link)
            return
        except OSError as e:
            log.warning("Symlink creation failed (%s), falling back to directory swap", e)
            
    if dest_link.exists() or dest_link.is_symlink():
        if dest_link.is_symlink():
            dest_link.unlink()
        else:
            shutil.rmtree(dest_link)
    shutil.copytree(src_dir, dest_link)


def clean_old_versions(config: Config) -> None:
    """Clean up old extracted version directories under opt/versions/."""
    try:
        retention = getattr(config, "update_retention", 3)
        versions_dir = config.opt_dir / "versions"
        if not versions_dir.exists():
            return
            
        current_symlink = config.opt_dir / "current"
        current_target = None
        if current_symlink.is_symlink():
            current_target = current_symlink.readlink().name
            
        # List all directories under opt/versions
        dirs = []
        for p in versions_dir.iterdir():
            if p.is_dir():
                if p.name.endswith(".backup"):
                    continue
                # Do not delete the currently active version folder
                if current_target and p.name == current_target:
                    continue
                # Do not delete the running version we are updating from/to
                if p.name == config.agent_version:
                    continue
                dirs.append((p, p.stat().st_mtime))
                
        # Sort oldest first
        dirs.sort(key=lambda x: x[1])
        
        # Keep at most `retention` directories (excluding the current/running ones)
        if len(dirs) > retention:
            for p, _ in dirs[:-retention]:
                try:
                    shutil.rmtree(p)
                    log.info("Cleaned up old version directory: %s", p)
                except Exception as e:
                    log.warning("Failed to remove old version directory %s: %s", p, e)
    except Exception as exc:
        log.warning("Failed to clean up old version directories: %s", exc)


def download_with_timeout(url: str, dest_path: Path, timeout: int = 30) -> None:
    """Download a file from url to dest_path with a timeout using urllib."""
    import urllib.request
    log.info("Downloading file with timeout %ds...", timeout)
    req = urllib.request.Request(url, headers={"User-Agent": "wolfpanel-agent"})
    with urllib.request.urlopen(req, timeout=timeout) as response:
        with open(dest_path, "wb") as f:
            while chunk := response.read(8192):
                f.write(chunk)


def self_update(config: Config, *args, **kwargs) -> None:
    """Perform self-update of the agent."""
    info = None
    api_client = None
    payload = None

    if len(args) >= 1:
        first = args[0]
        if isinstance(first, UpdateInfo):
            info = first
        else:
            api_client = first
            
    if len(args) >= 2:
        payload = args[1]

    job_id = payload.get("job_id") if payload else None
    token = load_agent_token(config) if api_client else None

    # Step 2: Report status "running" to backend for this job_id
    if api_client and job_id:
        try:
            log.info("Reporting status 'running' to backend for job_id=%s", job_id)
            api_client.post_command_result(
                job_id,
                {"status": "running", "result": {"message": "Update in progress..."}},
                token
            )
        except Exception as exc:
            log.warning("Failed to report running status: %s", exc)

    # Step 3: Fetch latest.json (if not already fetched from CLI)
    if not info:
        if not api_client:
            from api_client import get_api_client
            api_client = get_api_client(config)
            token = load_agent_token(config)
        info = check(config, api_client)
        if not info:
            raise Exception("Manifest check returned no info")

    # Defense in depth: check() already validates the manifest channel against
    # config.release_channel, but self_update() can also be handed an `info`
    # built elsewhere (e.g. CLI), so re-assert it here before doing anything.
    if info.channel and info.channel != config.release_channel:
        msg = (
            f"channel mismatch: agent is on channel '{config.release_channel}' "
            f"but manifest is for channel '{info.channel}'; refusing to update"
        )
        log.error(msg)
        if api_client and job_id:
            try:
                api_client.post_command_result(job_id, {"status": "failed", "result": {"error": msg}}, token)
            except Exception:
                pass
        raise Exception(msg)

    if config.api_mock and info and not info.signature:
        info.signature = "mock_signature"

    if not info.url or not info.sha256 or not info.signature:
        msg = "Manifest missing download_url, checksum, or signature"
        log.error(msg)
        if api_client and job_id:
            try:
                api_client.post_command_result(job_id, {"status": "failed", "result": {"error": msg}}, token)
            except Exception:
                pass
        raise Exception(msg)

    # Step 4: If already on latest version: report success, exit early
    if not info.update_available:
        log.info("Agent is already up to date (v%s)", info.current)
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "succeeded", "result": {"success": True, "message": f"Agent is up to date (v{info.current})"}},
                    token
                )
            except Exception as exc:
                log.warning("Failed to report success status: %s", exc)
        return

    # Step 5: Download new agent archive to a secure path inside cache_dir (root restricted)
    temp_dir = config.cache_dir / "wolfpanel-agent-update"

    try:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        temp_dir.mkdir(parents=True, exist_ok=True)
        temp_tar = temp_dir / "agent.tar.gz"

        log.info("Downloading update archive from %s to %s", info.url, temp_tar)
        if config.api_mock:
            # Generate mock package for testing
            mock_src = temp_dir / "mock_agent_src"
            mock_src.mkdir(parents=True, exist_ok=True)
            (mock_src / "src").mkdir(parents=True, exist_ok=True)
            
            # Simple main.py
            (mock_src / "src" / "main.py").write_text("""
import os, sys
flag_file = os.environ.get("WOLFPANEL_UPDATE_FLAG_FILE")
if flag_file:
    with open(flag_file, "w") as f:
        f.write("OK")
""", encoding="utf-8")
            (mock_src / "VERSION").write_text(info.latest, encoding="utf-8")
            (mock_src / "README.md").write_text("mock", encoding="utf-8")
            (mock_src / "systemd").mkdir(parents=True, exist_ok=True)
            (mock_src / "systemd" / "wolfpanel-agent.service").write_text("mock", encoding="utf-8")
            
            with tarfile.open(temp_tar, "w:gz") as tar:
                tar.add(mock_src / "src", arcname="src")
                tar.add(mock_src / "VERSION", arcname="VERSION")
                tar.add(mock_src / "README.md", arcname="README.md")
                tar.add(mock_src / "systemd", arcname="systemd")
                
            with open(temp_tar, "rb") as f:
                tar_bytes = f.read()
                
            sha256 = hashlib.sha256()
            sha256.update(tar_bytes)
            info.sha256 = sha256.hexdigest()
            info.size_bytes = len(tar_bytes)

            # Sign mock package using the TEST-ONLY private key (never the
            # production key -- see the TEST_* constants at the top of this
            # module).
            sha256_hash = sha256.digest()
            asn1_prefix = b"\x30\x31\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x01\x05\x00\x04\x20"
            padded = b"\x00\x01" + b"\xff" * 202 + b"\x00" + asn1_prefix + sha256_hash
            m = int.from_bytes(padded, "big")
            sig_int = pow(m, TEST_PRIVATE_KEY_D, TEST_PUBLIC_KEY_N)
            info.signature = sig_int.to_bytes(256, "big").hex()
        else:
            download_timeout = getattr(config, "download_timeout", 30)
            download_with_timeout(info.url, temp_tar, download_timeout)
    except Exception as exc:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "failed", "result": {"error": f"Download failed: {exc}"}},
                    token
                )
            except Exception:
                pass
        raise Exception(f"Download failed: {exc}")

    # Step 6: Verify size, then SHA256 checksum
    try:
        actual_size = temp_tar.stat().st_size
        if info.size_bytes is not None and actual_size != info.size_bytes:
            raise Exception(f"Size mismatch: expected {info.size_bytes} bytes, got {actual_size}")

        sha256 = hashlib.sha256()
        with open(temp_tar, "rb") as f:
            while chunk := f.read(8192):
                sha256.update(chunk)
        actual_checksum = sha256.hexdigest()

        if actual_checksum != info.sha256:
            raise Exception(f"Checksum mismatch: expected {info.sha256}, got {actual_checksum}")
        log.info("Checksum verified successfully")

        # Step 6.5: Verify cryptographic signature. Mock/test updates are
        # signed with the TEST-ONLY key and must be verified against it, never
        # against the production public key.
        with open(temp_tar, "rb") as f:
            tar_data = f.read()
        if config.api_mock:
            sig_ok = verify_signature(tar_data, info.signature, n=TEST_PUBLIC_KEY_N, e=TEST_PUBLIC_KEY_E)
        else:
            sig_ok = verify_signature(tar_data, info.signature)
        if not sig_ok:
            raise Exception("Signature verification failed: release package signature is invalid!")
        log.info("Release signature verified successfully")
    except Exception as exc:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "failed", "result": {"error": f"Verification failed: {exc}"}},
                    token
                )
            except Exception:
                pass
        raise Exception(f"Verification failed: {exc}")

    # Step 7: Back up current binary/source to a known path
    current_symlink = config.opt_dir / "current"
    backup_dir = config.opt_dir / "versions" / f"{config.agent_version}.backup"
    
    current_target = None
    try:
        if backup_dir.exists():
            shutil.rmtree(backup_dir)
            
        if current_symlink.is_symlink():
            current_target = current_symlink.readlink()
            if not current_target.is_absolute():
                current_target = (config.opt_dir / current_target).resolve()
        elif current_symlink.exists():
            current_target = current_symlink.resolve()
            
        log.info("Backing up current version from %s to %s", current_target, backup_dir)
        if current_target and current_target.exists():
            shutil.copytree(current_target, backup_dir)
    except Exception as exc:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "failed", "result": {"error": f"Failed to create backup: {exc}"}},
                    token
                )
            except Exception:
                pass
        raise Exception(f"Failed to create backup: {exc}")

    # Step 8: Extract and install new version alongside current
    new_version_dir = config.opt_dir / "versions" / info.latest
    try:
        if new_version_dir.exists():
            shutil.rmtree(new_version_dir)
        new_version_dir.mkdir(parents=True, exist_ok=True)
        
        log.info("Extracting update to %s", new_version_dir)
        with tarfile.open(temp_tar, "r:gz") as tar:
            tar.extractall(path=new_version_dir)

        # Downgrade/package guard: the extracted package's own VERSION file
        # must match the manifest version exactly. If it doesn't, abort here
        # -- BEFORE the symlink is touched or the service is restarted -- so
        # a corrupt build or a tampered/misconfigured channel can never end
        # up live.
        package_version_file = new_version_dir / "VERSION"
        if not package_version_file.is_file():
            raise Exception("staged package is missing its VERSION file")
        package_version = package_version_file.read_text(encoding="utf-8").strip()
        if package_version != info.latest:
            raise Exception(
                f"package VERSION mismatch: manifest says {info.latest} but "
                f"extracted package is {package_version}; refusing to activate it"
            )
        log.info("staged package VERSION matches manifest: %s", package_version)

        # Clean up temp files
        if temp_dir.exists():
            shutil.rmtree(temp_dir)

        wrapper_path = new_version_dir / "wolfpanel-agent"
        wrapper_path.write_text(f"""#!/usr/bin/env bash
exec {sys.executable} "$(dirname "$0")/src/main.py" "$@"
""", encoding="utf-8")
        wrapper_path.chmod(0o755)
        log.info("Staged new version files at %s", new_version_dir)
        # Clean up old versions
        clean_old_versions(config)
    except Exception as exc:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        if new_version_dir.exists():
            shutil.rmtree(new_version_dir)
        if backup_dir.exists():
            shutil.rmtree(backup_dir)
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "failed", "result": {"error": f"Failed to stage new files: {exc}"}},
                    token
                )
            except Exception:
                pass
        raise Exception(f"Failed to stage new files: {exc}")

    # Step 9: Detect startup method
    state_file = config.var_dir / "update_state.json"
    
    use_systemd = False
    if os.environ.get("INVOCATION_ID"):
        use_systemd = True
    else:
        try:
            res = subprocess.run(["systemctl", "is-active", "wolfpanel-agent"], capture_output=True, text=True)
            if res.stdout.strip() == "active":
                use_systemd = True
        except Exception:
            pass
            
    update_state = {
        "job_id": job_id,
        "backup_dir": str(backup_dir),
        "current_symlink": str(current_symlink),
        "rollback_version": config.agent_version,
        "updating_to_version": info.latest,
        "temp_dir": str(temp_dir),
        "api_url": config.api_url,
        "token": token,
        "systemd": use_systemd
    }
    
    try:
        with open(state_file, "w") as f:
            json.dump(update_state, f)
        log.info("Wrote update_state.json: %s", update_state)
    except Exception as exc:
        if new_version_dir.exists():
            shutil.rmtree(new_version_dir)
        if backup_dir.exists():
            shutil.rmtree(backup_dir)
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "failed", "result": {"error": f"Failed to write update state: {exc}"}},
                    token
                )
            except Exception:
                pass
        raise Exception(f"Failed to write update state: {exc}")

    import command_runner
    if command_runner._worker_manager is not None:
        log.info("Self-update initiated: pausing new command polling and execution...")
        command_runner._worker_manager.updating = True
        
        # Wait for other active commands to finish execution
        start_wait = time.time()
        max_wait = getattr(config, "update_timeout", 60) - 10
        if max_wait < 10:
            max_wait = 30
        while time.time() - start_wait < max_wait:
            with command_runner._worker_manager.lock:
                other_active = [cid for cid in command_runner._worker_manager.active_commands if cid != job_id]
            if not other_active:
                break
            log.info("Waiting for other running commands to finish: %s", other_active)
            time.sleep(1.0)

    log.info("Swapping current directory path to new version...")
    try:
        swap_directories_atomic(new_version_dir, current_symlink)
    except Exception as exc:
        if new_version_dir.exists():
            shutil.rmtree(new_version_dir)
        if backup_dir.exists():
            shutil.rmtree(backup_dir)
        if state_file.is_file():
            state_file.unlink()
        if api_client and job_id:
            try:
                api_client.post_command_result(
                    job_id,
                    {"status": "failed", "result": {"error": f"Failed to swap symlink: {exc}"}},
                    token
                )
            except Exception:
                pass
        raise Exception(f"Failed to swap symlink: {exc}")

    # Spawn detached watcher process
    watcher_code = """
import os, sys, time, json, shutil, urllib.request
from pathlib import Path

state_file_path = sys.argv[1]
timeout = int(sys.argv[2]) if len(sys.argv) > 2 else 60
time.sleep(timeout)

if not os.path.exists(state_file_path):
    sys.exit(0)

try:
    with open(state_file_path, "r") as f:
        state = json.load(f)
except Exception:
    sys.exit(1)

job_id = state.get("job_id")
backup_dir = Path(state.get("backup_dir"))
current_symlink = Path(state.get("current_symlink"))
temp_dir = Path(state.get("temp_dir"))
api_url = state.get("api_url")
token = state.get("token")
use_systemd = state.get("systemd")

print("Timeout/failure detected. Restoring backup...", flush=True)
try:
    if backup_dir.exists():
        if current_symlink.is_symlink():
            current_symlink.unlink()
        elif current_symlink.exists():
            if current_symlink.is_dir():
                shutil.rmtree(current_symlink)
            else:
                current_symlink.unlink()
        
        try:
            os.symlink(backup_dir, current_symlink)
        except Exception:
            shutil.copytree(backup_dir, current_symlink)
except Exception as e:
    print(f"Failed to restore backup: {e}", flush=True)

try:
    if temp_dir.exists():
        shutil.rmtree(temp_dir)
except Exception:
    pass

try:
    os.unlink(state_file_path)
except Exception:
    pass

if api_url and job_id:
    try:
        url = f"{api_url}/api/v1/agent/commands/{job_id}/result"
        headers = {"Content-Type": "application/json", "User-Agent": "wolfpanel-agent"}
        if token:
            headers["X-Agent-Token"] = token
        payload = {"status": "failed", "result": {"error": "new version failed to connect"}}
        req = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=15) as res:
            pass
    except Exception as e:
        print(f"Failed to report failure to backend: {e}", flush=True)

if use_systemd:
    os.system("systemctl restart wolfpanel-agent")
else:
    os.spawnv(os.P_NOWAIT, sys.executable, [sys.executable, str(current_symlink / "src" / "main.py"), "run"])
"""
    watcher_argv = [sys.executable, "-c", watcher_code, str(state_file), str(config.update_timeout)]
    watcher_spawned = False
    if use_systemd and shutil.which("systemd-run"):
        # CRITICAL: a plain subprocess.Popen(..., start_new_session=True) only
        # detaches the watcher's process group/session -- it does NOT remove
        # it from wolfpanel-agent.service's systemd cgroup. Under the default
        # KillMode=control-group, "systemctl restart wolfpanel-agent" (below)
        # kills every process in that cgroup, including this watcher, before
        # it can ever wake up from its sleep(). Launching it as its own
        # transient systemd unit via systemd-run puts it in a sibling cgroup
        # that the restart cannot touch.
        unit_name = f"wolfpanel-agent-rollback-{job_id or 'update'}-{int(time.time())}"
        try:
            log.info("Spawning rollback watcher as independent systemd unit %s...", unit_name)
            subprocess.Popen(
                ["systemd-run", "--unit", unit_name, "--collect", "--quiet", "--", *watcher_argv],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            watcher_spawned = True
        except Exception as exc:
            log.warning("Failed to spawn watcher via systemd-run: %s", exc)

    if not watcher_spawned:
        try:
            log.info("Spawning detached rollback watcher process...")
            subprocess.Popen(
                watcher_argv,
                start_new_session=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
        except Exception as exc:
            log.warning("Failed to spawn detached watcher: %s", exc)

    if "pytest" in sys.modules:
        log.info("Running under pytest; skipping reload and process exit.")
        if state_file.is_file():
            state_file.unlink()
        return

    # Trigger restart based on startup detection
    if use_systemd:
        log.info("Running under systemd. Executing systemctl restart wolfpanel-agent...")
        try:
            subprocess.Popen(["systemctl", "restart", "wolfpanel-agent"])
            sys.exit(0)
        except Exception as exc:
            log.error("Failed to run systemctl restart: %s. Falling back to standalone spawn.", exc)
            use_systemd = False

    if not use_systemd:
        log.info("Running standalone. Spawning new version process...")
        try:
            subprocess.Popen([sys.executable, str(current_symlink / "src" / "main.py"), "run"])
            sys.exit(0)
        except Exception as exc:
            log.error("Failed to spawn new standalone process: %s", exc)
            raise exc
