#!/usr/bin/env python3
import argparse
import json
import ipaddress 

from subprocess import check_output

def _get_available_machines() -> list:
    output = check_output(['nix', 'flake', 'show', '--json'])
    parsed_output = json.loads(output)
    machines = parsed_output.get('nixosConfigurations', dict()).keys()
    return list(machines)

def _validate_ip(ip: str) -> bool:
    try:
        ipaddress.ip_address(ip)
        return True
    except ValueError:
        return False

def _check_ssh_connection(ip: str) -> bool:
    try:
        check_output(['ssh', f'root@{ip}', 'echo', 'Connected'])
        return True
    except Exception:
        return False

def bootstrap_machine(ip: str):
    check_output(
        [
            'nix',
            'run',
            'github:nix-community/nixos-anywhere',
            '--',
            '--flake',
            '".#bootstrap"',
            '--target-host',
            f'root@{ip}',
            '--build-on-remote'
        ]
    )


def get_ssh_key(ip: str) -> str:
    """
    This function uses machines ssh-keyscan to get the ssh key and then get the ed25519 key
    """


def get_machine_config(machine_name: str) -> dict:
    output = check_output(['nix', 'eval', '--json', f'.#nixosConfigurations.{machine_name}.config.kropcloud'])
    return json.loads(output)


def main() -> int:
    parser = argparse.ArgumentParser(description='Install a machine')
    parser.add_argument('machine_name', type=str, help='The name of the machine to install')
    parser.add_argument('machine_ip', type=str, help='The ip of the machine to install')
    args = parser.parse_args()

    machine_name = args.machine_name
    if not machine_name in _get_available_machines():
        raise ValueError(f'Machine {machine_name} not found, available machines are: {_get_available_machines()}')

    if _validate_ip(args.machine_ip):
        raise ValueError(f'Invalid IP address {args.machine_ip}')

    machine_config = get_machine_config(machine_name)
    print(machine_config)
    # We are bootstraping the machine first because we need their ssh keys
    bootstrap_machine()

    while not _check_ssh_connection():
        time.sleep(5)

    # # connect and get ssh keys

    ssh_key = get_ssh_key()

    # Add the ssh key to keys in secrets/secrets.nix
    # and rekey the secrets

    # install_machine()

    return 0

if __name__ == '__main__':
    raise SystemExit(main())