import argparse
import json
import os
import subprocess

GET_INPUTS_CMD = [
    "nix-instantiate",
    "--eval",
    "--json",  # This parser is stupid, better provide it with pre-eaten stuff
    "--expr",
    "builtins.fromJSON (builtins.toJSON (import ./flake.nix).inputs)",
]


def process_flake(flakeUri: str) -> None:
    # get full path
    flakeUri = os.path.normpath(flakeUri)
    flakeFile = os.path.join(flakeUri, "flake.nix")
    if not os.path.isfile(flakeFile):
        raise FileNotFoundError(f"Flake not found: {flakeUri}")
    # import dependencies
    p = subprocess.run(GET_INPUTS_CMD, cwd=flakeUri, stdout=subprocess.PIPE)
    deps = json.loads(p.stdout)
    p.check_returncode()
    # for each dependency
    for dep_name, dep in deps.items():
        dep_url = dep["url"]
        # if not local path, continue
        if not (
            dep_url.startswith("path:")
            or dep_url.startswith("git+file:")
        ):
            continue
        if dep.get("flake", True):
            #  get flake file corresponding
            dep_path = dep_url.split(":")[1]
            if not dep_path.startswith("/"):
                dep_path = os.path.join(flakeUri, dep_path)
            process_flake(dep_path)
        # update lockfile
        cmd = [
            "nix",
            "--extra-experimental-features",
            "nix-command",
            "--extra-experimental-features",
            "flakes",
            "flake",
            "update",
            dep_name,
        ]
        p = subprocess.run(cmd, cwd=flakeUri)
        p.check_returncode()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Recursively update lockfiles "
        "of flakes located on the system"
    )
    parser.add_argument("flake", help="Starting flake", default="/")
    args = parser.parse_args()
    process_flake(args.flake)