#!/usr/bin/env python3
"""Versioned macOS telemetry startup: plan/apply/rollback, no worker job launch.

The prepared stack owns credentials and volumes. SSH provides only a loopback
MQTT forward through an existing strictly pinned Tailscale SSH profile.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import plistlib
import re
import subprocess
import time
import urllib.request
from pathlib import Path

from manage_mission_core_launch_agent import _wait_for_health, _write_atomic
from migrate_worker_tunnel_logs import reload_agent

from k1link.launchd_logs import launchd_log_path, prepare_launchd_log
from k1link.web.compute_contour_network import _replace_environment_value

STARTUP = "com.nodedc.telemetry-startup.local"
TUNNEL = "com.nodedc.telemetry-tunnel.local"
CORE = "com.nodedc.mission-core.local"
DOCKER = "/Applications/Docker.app/Contents/Resources/bin/docker"


def command(args, timeout=30):
    # Never echo Compose output: expanded environments can contain credentials.
    return subprocess.run(args, capture_output=True, timeout=timeout, check=False)


def receiver_ready():
    try:
        with urllib.request.urlopen("http://127.0.0.1:18030/health", timeout=3) as r:
            d = json.loads(r.read())
        return (
            d.get("ok") is True
            and d.get("mqtt_connected") is True
            and d.get("database_reachable") is True
        )
    except (OSError, ValueError):
        return False


def broker_loopback():
    result = command(
        [
            DOCKER,
            "inspect",
            "--format",
            "{{json .NetworkSettings.Ports}}",
            "ndc-mission-core-mqtt-broker",
        ],
        10,
    )
    try:
        bindings = json.loads(result.stdout).get("1883/tcp", [])
        return result.returncode == 0 and bindings == [{"HostIp": "127.0.0.1", "HostPort": "1883"}]
    except (ValueError, AttributeError):
        return False


def reconcile(stack: Path):
    """One bounded retry, repeated by launchd after delayed Docker/network start."""
    if receiver_ready() and broker_loopback():
        return "ready"
    if command([DOCKER, "info", "--format", "{{.ServerVersion}}"], 10).returncode:
        # Desktop's own login setting is not a dependency manager. No VM/resource
        # setting is changed and no user inference job is started here.
        command(["/usr/bin/open", "-g", "-a", "/Applications/Docker.app"], 10)
        return "waiting-for-docker"
    compose = [
        DOCKER,
        "compose",
        "--project-directory",
        str(stack),
        "--env-file",
        str(stack / ".env"),
        "-f",
        str(stack / "compose.yaml"),
    ]
    if command([*compose, "config", "--quiet"], 20).returncode:
        return "configuration-unavailable"
    if command(
        [
            *compose,
            "up",
            "-d",
            "--no-build",
            "--pull",
            "never",
            "--wait",
            "--wait-timeout",
            "45",
            "broker",
            "timescale",
            "normalizer",
        ],
        60,
    ).returncode:
        return "waiting-for-receiver"
    return "ready" if receiver_ready() and broker_loopback() else "waiting-for-receiver"


def private_file(path):
    if path.is_symlink() or not path.is_file() or path.stat().st_uid != os.getuid():
        raise ValueError("Expected owned regular configuration file")
    return path.read_bytes()


def plan(stack: Path, ssh_alias: str, repository: Path, agents: Path):
    if not stack.is_absolute() or stack.resolve() != stack or not repository.is_absolute():
        raise ValueError("Canonical absolute installation paths required")
    if re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9_.-]{0,80}", ssh_alias) is None:
        raise ValueError("Invalid SSH profile")
    for name in (
        "compose.yaml",
        "runtime/agents.json",
        "runtime/mosquitto/acl",
        "runtime/mosquitto/passwords",
    ):
        private_file(stack / name)
    env = private_file(stack / ".env")
    public = dict(
        line.split("=", 1)
        for line in env.decode().splitlines()
        if "=" in line and not line.startswith("#")
    )
    if (
        public.get("MISSIONCORE_MQTT_PORT", "1883") != "1883"
        or public.get("MISSIONCORE_TELEMETRY_QUERY_PORT", "18030") != "18030"
    ):
        raise ValueError("This deployment profile requires MQTT 1883 and query 18030")
    # Verify the existing trust boundary; never enroll a host or accept a new key.
    ssh = command(["/usr/bin/ssh", "-G", ssh_alias])
    settings = dict(line.split(" ", 1) for line in ssh.stdout.decode().splitlines() if " " in line)
    if (
        ssh.returncode
        or settings.get("stricthostkeychecking") not in ("true", "yes")
        or not settings.get("hostname", "").endswith(".ts.net")
    ):
        raise ValueError("A strictly pinned Tailscale SSH profile is required")
    python = repository / ".venv/bin/python"
    if not python.is_file() or not Path(DOCKER).is_file():
        raise ValueError("Prepared application Python and Docker Desktop are required")
    common = dict(
        RunAtLoad=True, AbandonProcessGroup=False, ProcessType="Background", ExitTimeOut=15
    )
    startup = dict(
        common,
        Label=STARTUP,
        StartInterval=30,
        ProgramArguments=[
            str(python),
            str(repository / "scripts/manage_telemetry_startup.py"),
            "reconcile",
            "--stack-root",
            str(stack),
        ],
        StandardOutPath="/dev/null",
        StandardErrorPath=str(launchd_log_path("telemetry-startup.log")),
    )
    tunnel = dict(
        common,
        Label=TUNNEL,
        KeepAlive=True,
        ThrottleInterval=10,
        ProgramArguments=[
            "/usr/bin/ssh",
            "-o",
            "BatchMode=yes",
            "-o",
            "StrictHostKeyChecking=yes",
            "-o",
            "ExitOnForwardFailure=yes",
            "-o",
            "ConnectTimeout=8",
            "-o",
            "ServerAliveInterval=10",
            "-o",
            "ServerAliveCountMax=3",
            "-N",
            "-T",
            "-R",
            "127.0.0.1:1883:127.0.0.1:1883",
            ssh_alias,
        ],
        StandardOutPath="/dev/null",
        StandardErrorPath=str(launchd_log_path("telemetry-tunnel.log")),
    )
    core_path = agents / (CORE + ".plist")
    core = plistlib.loads(private_file(core_path))
    if core.get("Label") != CORE or core.get("WorkingDirectory") != str(repository):
        raise ValueError("Canonical Mission Core installation changed")
    core["EnvironmentVariables"]["MISSIONCORE_TELEMETRY_PLANE_ROOT"] = str(stack)
    desired = {
        stack / ".env": _replace_environment_value(
            env.decode(), "MISSIONCORE_MQTT_BIND_ADDRESS", "127.0.0.1"
        ).encode()
    }
    for label, doc in ((STARTUP, startup), (TUNNEL, tunnel), (CORE, core)):
        path = agents / (label + ".plist")
        if path.exists() and plistlib.loads(private_file(path)).get("Label") != label:
            raise ValueError("Foreign LaunchAgent at target path")
        desired[path] = plistlib.dumps(doc, sort_keys=True)
    changes = []
    for path, data in desired.items():
        before = private_file(path) if path.exists() else None
        changes.append(
            dict(
                path=str(path),
                before=hashlib.sha256(before).hexdigest() if before is not None else None,
                after=hashlib.sha256(data).hexdigest(),
            )
        )
    document = dict(
        schema_version="missioncore.telemetry-startup-plan/v1",
        changes=changes,
        stack_root=str(stack),
        ssh_profile=ssh_alias,
        mqtt="loopback-over-ssh-tailscale",
        startup="macOS user login; retry every 30s",
        worker_jobs_started=False,
    )
    document["artifact_sha256"] = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    document["compose_sha256"] = hashlib.sha256((stack / "compose.yaml").read_bytes()).hexdigest()
    document["sha256"] = hashlib.sha256(json.dumps(document, sort_keys=True).encode()).hexdigest()
    return document, desired


def restore(backup: Path, agents: Path):
    manifest = json.loads(private_file(backup / "manifest.json"))
    for item in manifest["changes"]:
        path = Path(item["path"])
        current = hashlib.sha256(private_file(path)).hexdigest() if path.exists() else None
        if current not in (item["before"], item["after"]):
            raise ValueError("Installed configuration changed; refusing rollback")
    for label in (STARTUP, TUNNEL):
        command(["launchctl", "bootout", f"gui/{os.getuid()}/{label}"])
    for i, item in enumerate(manifest["changes"]):
        path = Path(item["path"])
        if item["before"] is None:
            path.unlink(missing_ok=True)
        else:
            old = private_file(backup / str(i))
            if hashlib.sha256(old).hexdigest() != item["before"]:
                raise ValueError("Rollback payload changed")
            _write_atomic(path, old)
    for label in (STARTUP, TUNNEL, CORE):
        path = agents / (label + ".plist")
        if path.exists():
            reload_agent(path, label)
    if not _wait_for_health(45):
        raise RuntimeError("Core health not accepted after rollback")
    return {
        "restored": True,
        "container_state": "not changed; declared endpoint restored on next reconciliation",
    }


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("action", choices=["plan", "apply", "reconcile", "rollback"])
    p.add_argument("--stack-root", type=Path)
    p.add_argument("--repository-root", type=Path, default=Path(__file__).resolve().parents[1])
    p.add_argument("--ssh-alias", default="mission-gpu")
    p.add_argument("--expected-sha256")
    p.add_argument("--backup", type=Path)
    a = p.parse_args()
    agents = Path.home() / "Library/LaunchAgents"
    if a.action == "rollback":
        if a.backup is None:
            p.error("--backup required")
        print(json.dumps(restore(a.backup, agents)))
        return
    if a.stack_root is None:
        p.error("--stack-root required")
    if a.action == "reconcile":
        try:
            phase = reconcile(a.stack_root)
        except (OSError, subprocess.SubprocessError, ValueError):
            phase = "receiver-unavailable"
        # No repeating stdout log; health remains independently observable.
        print(json.dumps({"phase": phase}))
        return
    document, desired = plan(a.stack_root, a.ssh_alias, a.repository_root, agents)
    if a.action == "plan":
        print(json.dumps(document, indent=2))
        return
    if a.expected_sha256 != document["sha256"]:
        raise ValueError("Startup plan changed before apply")
    backup = (
        a.repository_root
        / ".runtime/mission-core/telemetry-service-backups"
        / (str(time.time_ns()))
    )
    backup.mkdir(parents=True, mode=0o700)
    _write_atomic(backup / "manifest.json", json.dumps(document).encode())
    for i, path in enumerate(desired):
        if path.exists():
            _write_atomic(backup / str(i), private_file(path))
    for name in ("telemetry-startup.log", "telemetry-tunnel.log"):
        prepare_launchd_log(launchd_log_path(name))
    try:
        for path, data in desired.items():
            _write_atomic(path, data)
        for label in (STARTUP, TUNNEL, CORE):
            reload_agent(agents / (label + ".plist"), label)
        if not _wait_for_health(45):
            raise RuntimeError("Canonical Core failed health acceptance")
    except BaseException:
        restore(backup, agents)
        raise
    print(
        json.dumps(
            {
                "installed": True,
                "backup": str(backup),
                "plan_sha256": document["sha256"],
                "telemetry_acceptance": "pending fresh worker sample",
            }
        )
    )


if __name__ == "__main__":
    main()