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

from argparse import ArgumentParser
from pathlib import Path
from socket import gethostname

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


def _upd(
    subcommand: str = "switch",
    pull: bool = False,
    lock: bool = False,
    rewrite_hardware_configuration: bool = False,
    purge_vscode_extensions: bool = False,
    extra_args: list = [],
):
    path = Path("/etc/nixos")
    if subcommand not in ("switch", "boot", "test", "build-vm"):
        raise ValueError(f"{subcommand} is not a valid subcommand!")
    if path.exists():
        files_to_delete = {
            "Visual Studio Code user settings": ".config/Code/User/settings.json.bak",
            "fontconfig": ".config/fontconfig/conf.d/10-hm-fonts.conf.bak",
        }
        if pull:
            print(
                f"{c.BLUE}Pulling {c.YELLOW}NixOS{c.BLUE} configuration from remote{c.END}"
            )
            run(["git", "pull"], cwd=path)
        if lock:
            print(f"{c.BLUE}Updating {c.YELLOW}Nix Flake{c.BLUE} lock file{c.END}")
            run(["nix", "flake", "update", *extra_args], cwd=path)
        if rewrite_hardware_configuration:
            print(
                f"{c.BLUE}Updating {c.YELLOW}NixOS{c.BLUE} hardware configuration file for {c.YELLOW}{gethostname()}{c.BLUE}{c.END}"
            )
            run(
                [
                    "sudo",
                    "nixos-generate-config",
                    "--dir",
                    ".",
                ],
                cwd=path / "hosts",
            )
            print(
                f"{c.BLUE}Deleting redundant {c.YELLOW}NixOS{c.BLUE} configuration file{c.END}"
            )
            run(["sudo", "rm", "configuration.nix"], cwd=path / "hosts")
            print(
                f"{c.BLUE}Moving {c.YELLOW}NixOS{c.BLUE} hardware configuration file{c.END}"
            )
            run(
                [
                    "sudo",
                    "mv",
                    "hardware-configuration.nix",
                    "{hostname}.nix".format(hostname=gethostname()),
                ],
                cwd=path / "hosts",
            )
            print(
                f"{c.BLUE}Adding {c.YELLOW}NixOS{c.BLUE} hardware configuration file for {c.YELLOW}{gethostname()}{c.BLUE} to git{c.END}"
            )
            run(
                ["git", "add", "hosts/{hostname}.nix".format(hostname=gethostname())],
                cwd=path,
            )
        for file_description, file_path in files_to_delete.items():
            print(
                f"{c.BLUE}Deleting {c.YELLOW}{file_description}{c.BLUE} backup file{c.END}"
            )
            run(["rm", file_path], exit_on_error=False, cwd="/home/cswimr")
        if purge_vscode_extensions:
            print(
                f"{c.BLUE}Killing {c.YELLOW}Visual Studio Code{c.BLUE} processes{c.END}"
            )
            run(["killall", "code"], exit_on_error=False)
            print(
                f"{c.BLUE}Purging {c.YELLOW}Visual Studio Code{c.BLUE} extensions{c.END}"
            )
            run(["rm", "-rf", ".vscode/extensions"], cwd="/home/cswimr")
        print(f"{c.BLUE}Rebuilding {c.YELLOW}NixOS{c.BLUE} configuration{c.END}")
        # TODO: Remove --impure once the Starship module is merged - see ../../nixos/shell.nix for more information
        extra_args.append("--impure")
        if "--impure" in extra_args:
            print(f"{c.RED}WARNING: The --impure flag is set!{c.END}")
        run(["sudo", "nixos-rebuild", subcommand, *extra_args], cwd=path)


if __name__ == "__main__":
    parser = ArgumentParser(
        prog="upd.py",
        description="Update your NixOS system via a convenient wrapper script.",
        epilog="Example usage: upd.py",
    )
    parser.add_argument(
        "--subcommand",
        help="Which subcommand of `nixos-rebuild` to run.",
        default="switch",
    )
    parser.add_argument(
        "--pull",
        help="If this is set, the script will run `git pull` in the flake directory before updating.",
        action="store_true",
    )
    parser.add_argument(
        "--lock",
        help="If this is set, `nix flake update` will be ran before the system is updated.",
        action="store_true",
    )
    parser.add_argument(
        "--rewrite-hardware-configuration",
        help="If this is set, `nixos-generate-config` will be ran before the system is updated.",
        action="store_true",
    )
    parser.add_argument(
        "--purge-vscode-extensions",
        help="If this is set, extensions in the `~/.vscode/extensions` folder will be deleted before the system is updated.",
        action="store_true",
    )

    args = parser.parse_args()

    _upd(
        subcommand=args.subcommand,
        pull=args.pull,
        lock=args.lock,
        rewrite_hardware_configuration=args.rewrite_hardware_configuration,
        purge_vscode_extensions=args.purge_vscode_extensions,
        # extra_args = args.extra_args
    )