#!/usr/bin/env python3
"""
Generate an SSH config file for all exe.dev VMs.

For each SSH key, runs `ssh -i <key> exe.dev ls --json` to discover VMs,
then writes ~/.ssh/exe-dev-config with Host entries so that
`ssh <vm_name>` connects to <vm_name>.exe.xyz with the correct key.

Usage:
    exe-ssh-config-generator [key1 key2 ...]

If no keys are specified, uses the default keys:
    ~/.ssh/id_ed25519, ~/.ssh/id_ed25519-two, ~/.ssh/id_rsa
(only those that exist on disk).
"""

import json
import os
import subprocess
import sys

CONFIG_FILE = os.path.expanduser("~/.ssh/exe-dev-config")
KNOWN_HOSTS_FILE = os.path.expanduser("~/.ssh/exe-dev-known-hosts")
SSH_CONFIG = os.path.expanduser("~/.ssh/config")
INCLUDE_LINE = "Include ~/.ssh/exe-dev-config"

DEFAULT_KEYS = [
    "~/.ssh/id_ed25519",
    "~/.ssh/id_ed25519-two",
    "~/.ssh/id_rsa",
]


def expand_key(path):
    return os.path.expanduser(path)


def parse_dest(ssh_dest):
    """Split an ssh_dest into (hostname, user).

    Modern VMs report a bare hostname like 'foo.exe.xyz' (user is None).
    Legacy VMs (never migrated to name-based routing) report a full SSH
    destination like 'vm+foo@vm.exe.xyz', which routes by username through
    the old proxy and must be split into User + HostName for a valid
    ssh_config entry — a 'user@host' string is not a resolvable HostName.
    """
    if "@" in ssh_dest:
        user, host = ssh_dest.rsplit("@", 1)
        return host, user
    return ssh_dest, None


def discover_vms(key_path):
    """Run ssh -i <key> exe.dev ls --json and return list of VM dicts."""
    try:
        result = subprocess.run(
            [
                "ssh",
                "-i", key_path,
                "-F", "/dev/null",
                "-o", "IdentityAgent=none",
                "-o", "IdentitiesOnly=yes",
                "-o", "PreferredAuthentications=publickey",
                "-o", "PubkeyAuthentication=yes",
                "-o", "PasswordAuthentication=no",
                "-o", "KbdInteractiveAuthentication=no",
                "-o", "ChallengeResponseAuthentication=no",
                "-o", "StrictHostKeyChecking=no",
                "-o", "UserKnownHostsFile=/dev/null",
                "-o", "GlobalKnownHostsFile=/dev/null",
                "-o", "BatchMode=yes",
                "-o", "LogLevel=ERROR",
                "exe.dev", "ls", "--json",
            ],
            capture_output=True,
            text=True,
            timeout=30,
        )
    except subprocess.TimeoutExpired:
        print(f"  timeout connecting with {key_path}", file=sys.stderr)
        return []

    if result.returncode != 0:
        print(f"  ssh failed with {key_path}: {result.stderr.strip()}", file=sys.stderr)
        return []

    try:
        data = json.loads(result.stdout)
    except json.JSONDecodeError:
        print(f"  invalid JSON from {key_path}", file=sys.stderr)
        return []

    return data.get("vms", [])


def fetch_host_key():
    """Fetch the SSH host key from exe.dev (shared by all VMs via the SSH proxy)."""
    try:
        result = subprocess.run(
            ["ssh-keyscan", "exe.dev"],
            capture_output=True,
            text=True,
            timeout=10,
        )
    except subprocess.TimeoutExpired:
        print("  timeout fetching host key from exe.dev", file=sys.stderr)
        return None

    for line in result.stdout.splitlines():
        if line and not line.startswith("#"):
            # Return just the key type and key data (strip the hostname)
            parts = line.split(None, 2)
            if len(parts) == 3:
                return parts[1], parts[2]  # key_type, key_data
    print("  no host key found from exe.dev", file=sys.stderr)
    return None


def write_known_hosts(vm_key_map, host_key):
    """Write a known_hosts file with the exe.dev host key for every VM."""
    key_type, key_data = host_key
    lines = [
        "# Auto-generated by exe-ssh-config-generator",
        "# All exe.dev VMs share the same host key (SSH proxy).",
        "",
    ]
    # Key the entries by the actual host we connect to (the HostName), not the
    # raw ssh_dest — legacy 'user@host' dests share one host (vm.exe.xyz), and
    # ssh verifies the host key against the HostName, not the user@host string.
    seen = set()
    # Also include exe.dev itself
    lines.append(f"exe.dev {key_type} {key_data}")
    seen.add("exe.dev")
    for vm_name in sorted(vm_key_map):
        ssh_dest, _ = vm_key_map[vm_name]
        host, _user = parse_dest(ssh_dest)
        if host in seen:
            continue
        seen.add(host)
        lines.append(f"{host} {key_type} {key_data}")

    with open(KNOWN_HOSTS_FILE, "w") as f:
        f.write("\n".join(lines) + "\n")
    print(f"wrote {len(seen)} entries to {KNOWN_HOSTS_FILE}")


def generate_config(vm_key_map, have_known_hosts):
    """Generate SSH config content for all VMs."""
    lines = [
        "# Auto-generated by exe-ssh-config-generator",
        "# Do not edit manually; re-run the generator to update.",
        "",
    ]

    for vm_name in sorted(vm_key_map):
        ssh_dest, key_path = vm_key_map[vm_name]
        host, user = parse_dest(ssh_dest)
        lines.append(f"Host {vm_name} {ssh_dest}")
        lines.append(f"    HostName {host}")
        if user:
            lines.append(f"    User {user}")
        lines.append(f"    IdentityFile {key_path}")
        lines.append(f"    IdentitiesOnly yes")
        lines.append(f"    IdentityAgent none")
        lines.append(f"    PreferredAuthentications publickey")
        lines.append(f"    PubkeyAuthentication yes")
        lines.append(f"    PasswordAuthentication no")
        lines.append(f"    KbdInteractiveAuthentication no")
        if have_known_hosts:
            lines.append(f"    UserKnownHostsFile {KNOWN_HOSTS_FILE}")
            lines.append(f"    StrictHostKeyChecking yes")
        lines.append("")

    return "\n".join(lines)


def ensure_include():
    """Ensure ~/.ssh/config includes our generated config file."""
    if not os.path.exists(SSH_CONFIG):
        with open(SSH_CONFIG, "w") as f:
            f.write(INCLUDE_LINE + "\n\n")
        print(f"  created {SSH_CONFIG} with include")
        return

    with open(SSH_CONFIG, "r") as f:
        content = f.read()

    lines = content.splitlines(True)
    # Find where the Include line is (if present) and where Host/Match blocks start.
    include_idx = None
    first_host_idx = None
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped == INCLUDE_LINE:
            include_idx = i
        if first_host_idx is None and (stripped.startswith("Host ") or stripped.startswith("Match ")):
            first_host_idx = i

    if include_idx is not None and (first_host_idx is None or include_idx < first_host_idx):
        # Include exists and is before any Host/Match block — nothing to do.
        return

    if include_idx is not None:
        # Include exists but is inside a Host/Match block — remove it and re-prepend.
        del lines[include_idx]
        # Also remove a trailing blank line if one was left behind.
        if include_idx < len(lines) and lines[include_idx].strip() == "":
            del lines[include_idx]

    # Prepend the include so it takes priority
    with open(SSH_CONFIG, "w") as f:
        f.write(INCLUDE_LINE + "\n\n" + "".join(lines))
    print(f"  added include to {SSH_CONFIG}")


def main():
    if len(sys.argv) > 1:
        keys = sys.argv[1:]
    else:
        keys = [expand_key(k) for k in DEFAULT_KEYS if os.path.exists(expand_key(k))]

    if not keys:
        print("error: no SSH keys found or specified", file=sys.stderr)
        sys.exit(1)

    print(f"using keys: {', '.join(keys)}")

    # vm_name -> (ssh_dest, key_path)
    # First key to claim a VM wins.
    vm_key_map = {}

    for key in keys:
        print(f"querying exe.dev with {key} ...")
        vms = discover_vms(key)
        new = 0
        for vm in vms:
            name = vm.get("vm_name")
            dest = vm.get("ssh_dest")
            if name and dest and name not in vm_key_map:
                vm_key_map[name] = (dest, key)
                new += 1
        print(f"  found {len(vms)} VMs ({new} new)")

    if not vm_key_map:
        print("no VMs found", file=sys.stderr)
        sys.exit(1)

    print("fetching host key from exe.dev ...")
    host_key = fetch_host_key()
    have_known_hosts = False
    if host_key:
        write_known_hosts(vm_key_map, host_key)
        have_known_hosts = True
    else:
        print("  warning: skipping host key pinning", file=sys.stderr)

    config = generate_config(vm_key_map, have_known_hosts)
    with open(CONFIG_FILE, "w") as f:
        f.write(config)
    print(f"wrote {len(vm_key_map)} entries to {CONFIG_FILE}")

    ensure_include()
    print("done")


if __name__ == "__main__":
    main()