#!/usr/bin/env python3
"""Plan/apply a loopback-only SSH link to the existing NAS Map Gateway.

No remote command, NAS configuration, provider credential or cache is changed.
The user LaunchAgent reconnects independently of the Core process.
"""
from __future__ import annotations

import argparse
import hashlib
import http.client
import json
import os
import plistlib
import re
import socket
import subprocess
import tempfile
import time
from pathlib import Path

LABEL = "com.nodedc.mission-core.map-gateway-link"
PORT = 18103


def payload(target: str, runtime: Path) -> bytes:
    if not re.fullmatch(r"[a-zA-Z0-9_][a-zA-Z0-9_.-]*@[a-zA-Z0-9][a-zA-Z0-9.-]*", target):
        raise ValueError("SSH target must be an explicit user@host")
    return plistlib.dumps({
        "Label": LABEL,
        "ProgramArguments": [
            "/usr/bin/ssh", "-N", "-T", "-o", "BatchMode=yes",
            "-o", "StrictHostKeyChecking=yes", "-o", "ExitOnForwardFailure=yes",
            "-o", "ConnectTimeout=10", "-o", "ServerAliveInterval=15",
            "-o", "ServerAliveCountMax=3", "-o", "PermitLocalCommand=no",
            "-o", "ForwardAgent=no", "-o", "ForwardX11=no",
            "-L", f"127.0.0.1:{PORT}:127.0.0.1:{PORT}", target,
        ],
        "RunAtLoad": True, "KeepAlive": True, "ThrottleInterval": 30,
        "ExitTimeOut": 10,
        "StandardOutPath": str(runtime / "stdout.log"),
        "StandardErrorPath": str(runtime / "stderr.log"),
    }, sort_keys=True)


def digest(value: bytes | None) -> str:
    return hashlib.sha256(value).hexdigest() if value is not None else "absent"


def health() -> bool:
    connection = http.client.HTTPConnection("127.0.0.1", PORT, timeout=2)
    try:
        connection.request("GET", "/healthz", headers={
            "x-nodedc-user-id": "mission-core-loopback-operator",
        })
        response = connection.getresponse()
        raw = response.read(65537)
        if response.status != 200 or len(raw) > 65536:
            return False
        value = json.loads(raw)
        return value.get("ok") is True and value.get("service") == "nodedc-map-gateway"
    except (OSError, ValueError, http.client.HTTPException):
        return False
    finally:
        connection.close()


def loaded() -> bool:
    return subprocess.run(["launchctl", "print", f"gui/{os.getuid()}/{LABEL}"],
                          capture_output=True, timeout=5, check=False).returncode == 0


def unload() -> None:
    subprocess.run(["launchctl", "bootout", f"gui/{os.getuid()}/{LABEL}"],
                   capture_output=True, timeout=20, check=False)
    deadline = time.monotonic() + 20
    while loaded():
        if time.monotonic() > deadline:
            raise RuntimeError("Map Gateway link did not finish unloading")
        time.sleep(.2)


def write_atomic(path: Path, value: bytes) -> None:
    fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=".map-link-")
    try:
        with os.fdopen(fd, "wb") as stream:
            stream.write(value)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    finally:
        Path(temporary).unlink(missing_ok=True)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["plan", "apply", "status"])
    parser.add_argument("--target")
    parser.add_argument("--expected-current-sha256")
    parser.add_argument("--expected-desired-sha256")
    args = parser.parse_args()
    if args.action == "status":
        print(json.dumps({"loaded": loaded(), "healthy": health()}))
        return
    if not args.target:
        parser.error("--target is required")
    runtime = Path.home() / "Library/Logs/NODEDC/MissionCore/map-gateway-link"
    agent = Path.home() / "Library/LaunchAgents" / f"{LABEL}.plist"
    if agent.is_symlink():
        raise ValueError("Refusing a symlink LaunchAgent")
    previous = agent.read_bytes() if agent.exists() else None
    if previous is not None and plistlib.loads(previous).get("Label") != LABEL:
        raise ValueError("Refusing a foreign LaunchAgent")
    desired = payload(args.target, runtime)
    plan = {"label": LABEL, "agent_path": str(agent), "target": args.target,
            "local_bind": f"127.0.0.1:{PORT}", "remote_bind": f"127.0.0.1:{PORT}",
            "current_sha256": digest(previous), "desired_sha256": digest(desired),
            "remote_mutation": False}
    if args.action == "plan":
        print(json.dumps(plan, indent=2))
        return
    if (args.expected_current_sha256 != digest(previous)
            or args.expected_desired_sha256 != digest(desired)):
        raise ValueError("Map Gateway link plan changed before apply")
    was_loaded = loaded()
    if previous == desired and was_loaded and health():
        print(json.dumps({**plan, "applied": False, "healthy": True}))
        return
    if not was_loaded:
        with socket.socket() as probe:
            probe.bind(("127.0.0.1", PORT))  # Refuse an unrelated local listener.
    runtime.mkdir(parents=True, exist_ok=True, mode=0o700)
    agent.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    if previous is not None:
        write_atomic(runtime / f"backup-{digest(previous)}.plist", previous)
    try:
        if was_loaded:
            unload()
        write_atomic(agent, desired)
        subprocess.run(["launchctl", "bootstrap", f"gui/{os.getuid()}", str(agent)],
                       capture_output=True, timeout=20, check=True)
        deadline = time.monotonic() + 40
        while not health():
            if time.monotonic() > deadline:
                raise RuntimeError("Map Gateway link did not pass health acceptance")
            time.sleep(.5)
    except BaseException:
        unload()
        if previous is not None:
            write_atomic(agent, previous)
            if was_loaded:
                subprocess.run(["launchctl", "bootstrap", f"gui/{os.getuid()}", str(agent)],
                               capture_output=True, timeout=20, check=True)
        else:
            agent.unlink(missing_ok=True)
        raise
    print(json.dumps({**plan, "applied": True, "healthy": True}, indent=2))


if __name__ == "__main__":
    main()