#!/usr/bin/env -S uv run
# /// script
# dependencies = []
# ///

import argparse
import subprocess
import tempfile
import shutil
from pathlib import Path

def run_cmd(cmd, capture=False):
    """Run a command and optionally capture output."""
    if capture:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, check=True)
        return result.stdout.strip()
    else:
        subprocess.run(cmd, shell=True, check=True)


def derive_public_key(private_key_path):
    """Derive public key from private key."""
    return run_cmd(f"ssh-keygen -y -f {private_key_path}", capture=True)


def main():
    parser = argparse.ArgumentParser(description='Deploy NixOS using nixos-anywhere')
    parser.add_argument('--flake', required=True, help='Name of the flake configuration')
    parser.add_argument('--target-ip', required=True, help='IP address of target VM')
    args = parser.parse_args()

    flake_name = args.flake
    target_ip = args.target_ip

    temp_dir = Path(tempfile.mkdtemp())

    try:
        initrd_ssh_dir = temp_dir / "persist" / "secrets"
        initrd_ssh_dir.mkdir(parents=True)

        host_ssh_dir = temp_dir / "persist" / "etc" / "ssh"
        host_ssh_dir.mkdir(parents=True)

        ssh_key = run_cmd(f"passage show hosts/{flake_name}/ssh/ssh_host_ed25519_key", capture=True)

        initrd_ssh_key_path = initrd_ssh_dir / "ssh_host_ed25519_key"
        initrd_ssh_key_path.write_text(ssh_key + '\n')
        initrd_ssh_key_path.chmod(0o600)

        host_ssh_key_path = host_ssh_dir / "ssh_host_ed25519_key"
        host_ssh_key_path.write_text(ssh_key + '\n')
        host_ssh_key_path.chmod(0o600)

        public_key = derive_public_key(host_ssh_key_path)

        initrd_ssh_pub_path = initrd_ssh_dir / "ssh_host_ed25519_key.pub"
        initrd_ssh_pub_path.write_text(public_key + '\n')
        initrd_ssh_pub_path.chmod(0o644)

        host_ssh_pub_path = host_ssh_dir / "ssh_host_ed25519_key.pub"
        host_ssh_pub_path.write_text(public_key + '\n')
        host_ssh_pub_path.chmod(0o644)

        disk_key = run_cmd(f"passage show hosts/{flake_name}/disk-encryption/passphrase", capture=True)
        disk_key_file = Path("/tmp/disk.key")
        disk_key_file.write_text(disk_key)
        disk_key_file.chmod(0o600)

        cmd = [
            "nix", "run", "github:nix-community/nixos-anywhere", "--",
            "--flake", f".#{flake_name}",
            "--build-on", "remote",
            "--disk-encryption-keys", "/tmp/disk.key", str(disk_key_file),
            "--target-host", f"root@{target_ip}",
            "--extra-files", str(temp_dir)
        ]

        subprocess.run(cmd, check=True)

    finally:
        # Clean up
        shutil.rmtree(temp_dir)
        if Path("/tmp/disk.key").exists():
            Path("/tmp/disk.key").unlink()

if __name__ == "__main__":
    main()