#!/usr/bin/env python3
"""
axiom-zero-connect — One-Time Remote Frontend Setup
====================================================
Run this ONCE on the operator's desktop/laptop to connect
the GTK frontend to a remote Axiom Zero daemon.

Usage:
    axiom-zero-connect <server-ip-or-hostname>

What it does:
    1. SSHes into the server and copies the 3 TLS cert files
       (~/.axiom-zero/client.pem, client.key, ca.pem)
    2. Writes ~/.axiom-zero/connection.json with the server address
    3. Verifies the TLS connection succeeds
    4. Prints "Connected!" — from now on, axiom-zero-dashboard
       connects to this server automatically on every launch

If IP changes later:
    Just re-run:  axiom-zero-connect <new-ip>
    The connection.json is updated. No reinstall needed.
"""

import os
import sys
import json
import ssl
import socket
import subprocess
import getpass
from pathlib import Path

AXIOM_DIR   = Path.home() / ".axiom-zero"
CONN_FILE   = AXIOM_DIR / "connection.json"
CLIENT_CERT = AXIOM_DIR / "client.pem"
CLIENT_KEY  = AXIOM_DIR / "client.key"
CA_CERT     = AXIOM_DIR / "ca.pem"
PORT        = 7443
SERVER_CERT_DIR = "/var/lib/axiom-zero"


def main():
    if len(sys.argv) < 2:
        print("Usage: axiom-zero-connect <server-ip-or-hostname>")
        print("")
        print("Example:")
        print("  axiom-zero-connect 192.168.1.50")
        print("  axiom-zero-connect myserver.example.com")
        sys.exit(1)

    server = sys.argv[1]
    AXIOM_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)

    print(f"\n⚡ AXIOM ZERO — Remote Frontend Setup")
    print(f"   Connecting to server: {server}")
    print()

    # ── Step 1: Copy TLS certs via SCP ───────────────────────────────
    print("[1/3] Copying TLS certificates from server...")
    ssh_user = input(f"  SSH username for {server} [root]: ").strip() or "root"

    files_to_copy = [
        (f"{SERVER_CERT_DIR}/client.pem", str(CLIENT_CERT)),
        (f"{SERVER_CERT_DIR}/client.key", str(CLIENT_KEY)),
        (f"{SERVER_CERT_DIR}/ca.pem",     str(CA_CERT)),
    ]

    for remote_path, local_path in files_to_copy:
        cmd = ["scp", f"{ssh_user}@{server}:{remote_path}", local_path]
        result = subprocess.run(cmd)
        if result.returncode != 0:
            print(f"\n  ✖ Could not copy {remote_path}")
            print(f"    Make sure the Axiom Zero daemon is installed and running on {server}")
            print(f"    and that SSH is accessible.")
            sys.exit(1)
        os.chmod(local_path, 0o600)
        print(f"  ✔ {os.path.basename(local_path)}")

    # ── Step 2: Write connection.json ────────────────────────────────
    print(f"\n[2/3] Saving connection config → {CONN_FILE}")
    cfg = {
        "server_host": server,
        "server_port": PORT,
        "tls_verify": True,
        "ca_cert": str(CA_CERT),
        "connected_at": __import__('datetime').datetime.now().isoformat()
    }
    CONN_FILE.write_text(json.dumps(cfg, indent=2))
    CONN_FILE.chmod(0o600)
    print(f"  ✔ Saved")

    # ── Step 3: Test connection ──────────────────────────────────────
    print(f"\n[3/3] Testing TLS connection to {server}:{PORT}...")
    try:
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
        ctx.load_verify_locations(str(CA_CERT))
        ctx.load_cert_chain(certfile=str(CLIENT_CERT), keyfile=str(CLIENT_KEY))
        with socket.create_connection((server, PORT), timeout=5) as raw:
            with ctx.wrap_socket(raw, server_hostname=server) as tls:
                peer_cert = tls.getpeercert()
                print(f"  ✔ TLS handshake successful")
                print(f"  ✔ Server certificate verified")
    except ConnectionRefusedError:
        print(f"  ⚠ Connection refused on port {PORT}")
        print(f"    The daemon may still be starting. Try: axiom-zero-dashboard")
        print(f"    (it will reconnect automatically once the daemon is up)")
    except ssl.SSLError as e:
        print(f"  ✖ TLS error: {e}")
        sys.exit(1)
    except OSError as e:
        print(f"  ⚠ Could not reach {server}:{PORT} — {e}")
        print(f"    Check firewall: sudo ufw allow {PORT}/tcp")

    print(f"""
╔══════════════════════════════════════════════════════════╗
║  ✅  AXIOM ZERO REMOTE FRONTEND CONFIGURED               ║
╠══════════════════════════════════════════════════════════╣
║  Server  : {server:<47} ║
║  Port    : {PORT:<47} ║
║  Certs   : ~/.axiom-zero/                                ║
╠══════════════════════════════════════════════════════════╣
║  Launch the dashboard anytime:                           ║
║    axiom-zero-dashboard                                  ║
║                                                          ║
║  If the server IP changes, just re-run:                  ║
║    axiom-zero-connect <new-ip>                           ║
╚══════════════════════════════════════════════════════════╝
""")


if __name__ == "__main__":
    main()