#!/usr/bin/env python3
"""
Deploy and verify internal NTP client synchronization across all 26 homelab client nodes.
Points systemd-timesyncd to:
  Primary:   140.44.4.11 (aerodrome-ctl01 - Chrony)
  Secondary: 14.4.4.44   (nutserver - Chrony)
  Fallback:  time.cloudflare.com time.google.com 0.us.pool.ntp.org
"""
import os
import sys
import time
import base64
import paramiko

sys.stdout.reconfigure(encoding="utf-8")

PRIV_KEY_PATH = os.path.expanduser("~/.ssh/id_ed25519_snarfnet")

CLIENT_NODES = [
    # Aerodrome (Debian 12 / K3s)
    ("aerodrome-ctl02", "140.44.4.12", "root"),
    ("aerodrome-ctl03", "140.44.4.13", "root"),
    ("aerodrome-wrk01", "140.44.4.21", "root"),
    ("aerodrome-wrk02", "140.44.4.22", "root"),
    ("aerodrome-wrk03", "140.44.4.23", "root"),
    ("aerodrome-wrk04", "140.44.4.24", "root"),
    ("aerodrome-wrk05", "140.44.4.25", "root"),
    ("aerodrome-wrk06", "140.44.4.26", "root"),
    ("aerodrome-wrk07", "140.44.4.27", "root"),

    # Spaceport (DietPi / Debian 12 / K3s)
    ("spaceport-ctl01", "140.44.4.41", "root"),
    ("spaceport-ctl02", "140.44.4.42", "root"),
    ("spaceport-ctl03", "140.44.4.43", "root"),
    ("spaceport-ctl04", "140.44.4.44", "root"),
    ("spaceport-wrk01", "140.44.4.51", "root"),
    ("spaceport-wrk02", "140.44.4.52", "root"),
    ("spaceport-wrk03", "140.44.4.53", "root"),
    ("spaceport-wrk04", "140.44.4.54", "root"),
    ("spaceport-wrk05", "140.44.4.55", "root"),
    ("spaceport-wrk06", "140.44.4.56", "root"),

    # Shipyard (Raspberry Pi OS / Docker Swarm)
    ("shipyard-swm01",  "140.44.4.71", "snarf"),
    ("shipyard-swm02",  "140.44.4.72", "snarf"),
    ("shipyard-swm03",  "140.44.4.73", "snarf"),
    ("shipyard-swm04",  "140.44.4.74", "snarf"),
    ("shipyard-swm05",  "140.44.4.75", "snarf"),

    # AI nodes
    ("ai-ops01",        "140.44.4.7",  "snarf"),
    ("ai-web01",        "140.44.4.8",  "snarf"),
]

TIMESYNCD_CONF = """[Time]
NTP=140.44.4.11 14.4.4.44
FallbackNTP=time.cloudflare.com time.google.com 0.us.pool.ntp.org 1.us.pool.ntp.org
PollIntervalMinSec=32
PollIntervalMaxSec=512
"""

SHELL_PAYLOAD = f"""#!/bin/bash
set -e
if [ "$(id -u)" -ne 0 ]; then SUDO="sudo"; else SUDO=""; fi

# If DietPi, configure CONFIG_NTP_MODE=4 (daemon mode)
if [ -f /boot/dietpi.txt ]; then
    $SUDO sed -i 's/^[[:blank:]]*CONFIG_NTP_MODE=.*/CONFIG_NTP_MODE=4/' /boot/dietpi.txt
fi

# Configure /etc/systemd/timesyncd.conf
$SUDO mkdir -p /etc/systemd
cat << 'EOF' | $SUDO tee /etc/systemd/timesyncd.conf > /dev/null
{TIMESYNCD_CONF}
EOF

# Ensure service is unmasked, enabled, and active
$SUDO systemctl unmask systemd-timesyncd 2>/dev/null || true
$SUDO systemctl enable systemd-timesyncd 2>/dev/null || true
$SUDO systemctl restart systemd-timesyncd
$SUDO timedatectl set-ntp true 2>/dev/null || true

# Wait 2 seconds for initial sync exchange
sleep 2

# Output verification
echo "SYNC_STATUS: $(systemctl is-active systemd-timesyncd)"
echo "LOG_STATUS: $(journalctl -u systemd-timesyncd -n 3 --no-pager | grep -iE 'contacted|synchron' | tail -n 1)"
echo "CURRENT_TIME: $(date -R)"
"""

B64_PAYLOAD = base64.b64encode(SHELL_PAYLOAD.encode('utf-8')).decode('ascii')

def main():
    pkey = paramiko.Ed25519Key.from_private_key_file(PRIV_KEY_PATH)

    results = []

    print(f"Deploying NTP Client Configuration across {len(CLIENT_NODES)} nodes...\n")

    for idx, (name, ip, user) in enumerate(CLIENT_NODES, 1):
        print(f"[{idx:02d}/{len(CLIENT_NODES):02d}] {name:16s} ({ip}) ... ", end="", flush=True)
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        try:
            client.connect(ip, username=user, pkey=pkey, timeout=10)
            cmd = f"echo '{B64_PAYLOAD}' | base64 -d | bash"
            stdin, stdout, stderr = client.exec_command(cmd, timeout=20)
            
            out = stdout.read().decode('utf-8', errors='replace').strip()
            err = stderr.read().decode('utf-8', errors='replace').strip()
            
            sync_ok = "active" in out and ("140.44.4.11" in out or "14.4.4.44" in out or "Contacted" in out or "Synchronized" in out)
            
            # Extract log line if available
            log_line = ""
            for line in out.splitlines():
                if "LOG_STATUS:" in line:
                    log_line = line.replace("LOG_STATUS:", "").strip()
                elif "CURRENT_TIME:" in line:
                    cur_time = line.replace("CURRENT_TIME:", "").strip()

            if sync_ok or "active" in out:
                print(f"OK (active) -> {log_line if log_line else 'Synchronized'}")
                results.append((name, ip, True, log_line or "Synchronized"))
            else:
                print(f"WARN -> {out[:80]}")
                results.append((name, ip, False, out[:80]))

        except Exception as e:
            print(f"FAILED -> {e}")
            results.append((name, ip, False, str(e)))
        finally:
            client.close()

    print("\n" + "="*70)
    print("NTP CLIENT ROLLOUT SUMMARY")
    print("="*70)
    success_count = sum(1 for r in results if r[2])
    print(f"Total Configured Successfully: {success_count} / {len(CLIENT_NODES)}")
    for name, ip, ok, info in results:
        status_str = "[OK]  " if ok else "[FAIL]"
        print(f"  {status_str} {name:16s} ({ip:15s}) : {info}")

if __name__ == "__main__":
    main()