#! /usr/bin/env nix-shell
#! nix-shell /etc/nixos/scripts/nix/python.nix -i python

from argparse import ArgumentParser
from pathlib import Path

from common.colors import Colors as c
from common.common import run


def _vm(name: str = "nixpkgs", args: list = []):
    if name == "nixpkgs":
        args.extend(["-I", "nixpkgs=/bulk/home/cswimr/Projects/nixpkgs"])
    path = Path(f"/etc/nixos/hosts/virtual-machines/{name}")
    if path.exists():
        if (path / "flake.nix").exists():
            args.extend(["--flake", f"{path}#nixos"])
        elif (path / "default.nix").exists():
            args.extend(["-I", f"nixos-config={path}/default.nix", "--no-flake"])
        print(f"{c.BLUE}Building virtual machine {c.YELLOW}{name}{c.END}")
        run(["nixos-rebuild", "build-vm", *args], cwd=path)
        print(f"{c.BLUE}Starting virtual vachine {c.YELLOW}{name}{c.END}")
        run(["./result/bin/run-nixos-vm"], cwd=path)
        print(
            f"{c.BLUE}Virtual machine {c.YELLOW}{name} {c.BLUE}has {c.RED}stopped.{c.END}"
        )
    else:
        raise FileNotFoundError(f"Virtual machine {name} does not exist.")


if __name__ == "__main__":
    parser = ArgumentParser(
        prog="vm.py",
        description="Run a NixOS virtual machine from the system flake.",
        epilog="Example usage: vm.py nixpkgs",
    )
    parser.add_argument("name", help="Which virtual machine to run.")

    args = parser.parse_args()

    _vm(name=args.name)