#!/usr/bin/env python3
"""Bezalel Tailscale bootstrap scaffold.

Refs: timmy-home #535

Safe by default:
- builds a remote bootstrap shell script
- can write that script to disk
- can print the SSH command needed to execute it
- only runs remote SSH when --apply is explicitly passed
"""

from __future__ import annotations

import argparse
import json
import shlex
import subprocess
from pathlib import Path
from typing import Any

DEFAULT_HOST = "159.203.146.185"
DEFAULT_HOSTNAME = "bezalel"
DEFAULT_PEERS = {
    "mac": "100.124.176.28",
    "ezra": "100.126.61.75",
}


def build_remote_script(
    *,
    auth_key: str,
    ssh_public_key: str,
    peers: dict[str, str] | None = None,
    hostname: str = DEFAULT_HOSTNAME,
) -> str:
    peer_map = peers or DEFAULT_PEERS
    lines = [
        "#!/usr/bin/env bash",
        "set -euo pipefail",
        "curl -fsSL https://tailscale.com/install.sh | sh",
        f"tailscale up --authkey {shlex.quote(auth_key)} --ssh --hostname {shlex.quote(hostname)}",
        "install -d -m 700 ~/.ssh",
        f"touch ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys",
        f"grep -qxF {shlex.quote(ssh_public_key)} ~/.ssh/authorized_keys || printf '%s\\n' {shlex.quote(ssh_public_key)} >> ~/.ssh/authorized_keys",
        "tailscale status --json",
    ]
    for name, ip in peer_map.items():
        lines.append(f"ping -c 1 {shlex.quote(ip)} >/dev/null && echo 'PING_OK:{name}:{ip}'")
    return "\n".join(lines) + "\n"


def parse_tailscale_status(payload: dict[str, Any]) -> dict[str, Any]:
    self_block = payload.get("Self") or {}
    peers = payload.get("Peer") or {}
    return {
        "self": {
            "hostname": self_block.get("HostName"),
            "dns_name": self_block.get("DNSName"),
            "tailscale_ips": list(self_block.get("TailscaleIPs") or []),
        },
        "peers": {
            peer.get("HostName") or peer_key: list(peer.get("TailscaleIPs") or [])
            for peer_key, peer in peers.items()
        },
    }


def build_ssh_command(host: str, remote_script_path: str = "/tmp/bezalel_tailscale_bootstrap.sh") -> list[str]:
    return ["ssh", host, f"bash {shlex.quote(remote_script_path)}"]


def write_script(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content)


def run_remote(host: str, remote_script_path: str) -> subprocess.CompletedProcess[str]:
    return subprocess.run(build_ssh_command(host, remote_script_path), capture_output=True, text=True, timeout=120)


def parse_peer_args(items: list[str]) -> dict[str, str]:
    peers = dict(DEFAULT_PEERS)
    for item in items:
        name, ip = item.split("=", 1)
        peers[name.strip()] = ip.strip()
    return peers


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Prepare or execute Tailscale bootstrap for the Bezalel VPS.")
    parser.add_argument("--host", default=DEFAULT_HOST)
    parser.add_argument("--hostname", default=DEFAULT_HOSTNAME)
    parser.add_argument("--auth-key", help="Tailscale auth key")
    parser.add_argument("--auth-key-file", type=Path, help="Path to file containing the Tailscale auth key")
    parser.add_argument("--ssh-public-key", help="SSH public key to append to authorized_keys")
    parser.add_argument("--ssh-public-key-file", type=Path, help="Path to the SSH public key file")
    parser.add_argument("--peer", action="append", default=[], help="Additional peer as name=ip")
    parser.add_argument("--script-out", type=Path, default=Path("/tmp/bezalel_tailscale_bootstrap.sh"))
    parser.add_argument("--remote-script-path", default="/tmp/bezalel_tailscale_bootstrap.sh")
    parser.add_argument("--apply", action="store_true", help="Execute the generated script over SSH")
    parser.add_argument("--json", action="store_true")
    return parser.parse_args()


def _read_secret(value: str | None, path: Path | None) -> str | None:
    if value:
        return value.strip()
    if path and path.exists():
        return path.read_text().strip()
    return None


def main() -> None:
    args = parse_args()
    auth_key = _read_secret(args.auth_key, args.auth_key_file)
    ssh_public_key = _read_secret(args.ssh_public_key, args.ssh_public_key_file)
    peers = parse_peer_args(args.peer)

    if not auth_key:
        raise SystemExit("Missing Tailscale auth key. Use --auth-key or --auth-key-file.")
    if not ssh_public_key:
        raise SystemExit("Missing SSH public key. Use --ssh-public-key or --ssh-public-key-file.")

    script = build_remote_script(auth_key=auth_key, ssh_public_key=ssh_public_key, peers=peers, hostname=args.hostname)
    write_script(args.script_out, script)

    payload: dict[str, Any] = {
        "host": args.host,
        "hostname": args.hostname,
        "script_out": str(args.script_out),
        "remote_script_path": args.remote_script_path,
        "ssh_command": build_ssh_command(args.host, args.remote_script_path),
        "peer_targets": peers,
        "applied": False,
    }

    if args.apply:
        result = run_remote(args.host, args.remote_script_path)
        payload["applied"] = True
        payload["exit_code"] = result.returncode
        payload["stdout"] = result.stdout
        payload["stderr"] = result.stderr

    if args.json:
        print(json.dumps(payload, indent=2))
        return

    print("--- Bezalel Tailscale Bootstrap ---")
    print(f"Host: {args.host}")
    print(f"Local script: {args.script_out}")
    print("SSH command: " + " ".join(payload["ssh_command"]))
    if args.apply:
        print(f"Exit code: {payload['exit_code']}")


if __name__ == "__main__":
    main()