#!/usr/bin/env python

# Script to rebuild my system(s) when something changes. I'm running
# this from a cron job on a server.
#
# Usage: nix-conf [update|build] [nix build options]
#
# With no sub-command, update is implied.
#
# Example: nix-conf build --log-format raw

import argparse
import os
import shutil
import socket
import subprocess
import sys
from collections.abc import Callable
from shlex import quote
from subprocess import STDOUT, CalledProcessError, check_output
from typing import IO, Any, List
from os.path import expanduser
from pathlib import Path

verbose = True

HAS_TTY = sys.stdout.isatty()


def color_text(code: int, file: IO[Any] | None = None) -> Callable[[str], None]:
    def wrapper(text: str) -> None:
        if HAS_TTY:
            print(f"\x1b[{code}m{text}\x1b[0m", file=file)
        else:
            print(text, file=file)

    return wrapper


warn = color_text(31, file=sys.stderr)
info = color_text(32, file=sys.stderr)


def print_verbose(*args, **kwargs):
    if verbose:
        info(*args, **kwargs)


def system(cmd: str):
    check_call(cmd, shell=True)


def run(*args, **kwargs):
    print_verbose("$ " + (args[0] if isinstance(args[0], str) else " ".join([quote(a) for a in args[0]])))
    return subprocess.run(*args, **kwargs)


def check_call(*args, **kwargs):
    print_verbose("$ " + (args[0] if isinstance(args[0], str) else " ".join([quote(a) for a in args[0]])))
    return subprocess.check_call(*args, **kwargs)


def was_head_changed(cmd):
    old_tree = check_output("git rev-parse HEAD^{tree}", shell=True).decode().strip()
    if isinstance(cmd, str):
        system(cmd)
    else:
        cmd()
    new_tree = check_output("git rev-parse HEAD^{tree}", shell=True).decode().strip()
    if old_tree != new_tree:
        print(f"!!! Tree hash in {os.getcwd()} changed from {old_tree} to {new_tree}")
        return True
    return False


def exit_if_dirty():
    if run("git diff-index --quiet HEAD --", shell=True).returncode != 0:
        warn(f"{os.getcwd()} dirty!")
        sys.exit(1)


def cd(path):
    print_verbose(f"$ cd {path}")
    os.chdir(expanduser(path))


def update_nix_conf():
    cd("~/nix/conf")
    flake_lock_dirty = False
    if run("git diff --quiet flake.lock", shell=True).returncode != 0:
        check_call("git stash push -- flake.lock", shell=True)  # Dirty flake.lock can prevent git pull
        flake_lock_dirty = True

    if was_head_changed("git pull --quiet https://github.com/wentasah/nix-conf.git"):
        # Reexec potentially newer version of ourselves
        os.putenv("WAS_RERUN", "true")
        os.execv(sys.argv[0], sys.argv)

    if flake_lock_dirty and run("git stash pop --quiet", shell=True).returncode != 0:
        # Don't try to resolve the conflict. In the worst case, we'll
        # run an unnecessary rebuild.
        system("git restore flake.lock")
        system("git stash drop")

    # Git repo was not changed - report the result based on the value of RERUN variable.
    return os.getenv("WAS_RERUN") == "true"


def merge_nixpkgs_topics_into_master():
    cd("~/nix/nixpkgs")

    system("git checkout --quiet master")
    system("git reset --keep nixos-unstable")

    commit_msg='automatic merge by nix-conf'
    if os.path.isfile(".git/machete"):
        # On my laptop - merge all non-root branches managed by git-machete
        with open(".git/machete", 'r') as f:
            to_merge = [line.strip().split(' ', 1)[0] for line in f.readlines() if line.startswith((' ', '\t'))]
        commit_msg += "\n\nMerged branches:\n" + "\n".join([f"- {branch}" for branch in to_merge])
    else:
        # CI: Recreate the merge from the HEAD of my repo
        ## git fetch . test
        if commit_msg in check_output("git show wentasah/master", shell=True).decode():
            to_merge = check_output("git cat-file -p wentasah/master | sed -ne '/^parent / s///p'", shell=True).decode().splitlines()
            commit_msg = check_output("git show --format='%B' --no-patch wentasah/master", shell=True).decode()  # Use the same commit message
        else:
            remote = check_output('git config remote.wentasah.url', shell=True)
            warn(f"HEAD of {remote} is not a merge commit created by nix-conf")
            sys.exit(1)

    check_call(['git', 'merge', '--quiet', '--no-ff', '--no-commit'] + to_merge)
    check_call(['git', 'commit', '-m', commit_msg], env=os.environ | {
        "GIT_AUTHOR_DATE": check_output("git show nixos-unstable --format='%aD'", shell=True).decode(),
        "GIT_COMMITTER_DATE": check_output("git show nixos-unstable --format='%cD'", shell=True).decode(),
    })


def update_clonned_nixpkgs():
    # Update stable
    cd("~/nix/nixpkgs-stable")
    exit_if_dirty()
    system("git pull --quiet")

    # Update unstable + merge my changes
    cd("~/nix/nixpkgs")
    exit_if_dirty()

    def update_nixpkgs_and_merge_topics():
        old_wentasah_master = check_output("git rev-parse wentasah/master", shell=True).decode().strip()
        system("git fetch --quiet wentasah")
        new_wentasah_master = check_output("git rev-parse wentasah/master", shell=True).decode().strip()

        system("git fetch --quiet origin")
        old_nixos_unstable = check_output("git rev-parse nixos-unstable", shell=True).decode().strip()
        new_nixos_unstable = check_output("git rev-parse origin/nixos-unstable", shell=True).decode().strip()

        if old_wentasah_master == new_wentasah_master and \
           old_nixos_unstable == new_nixos_unstable:
            # No relevant refs changed - no need to recreate master by merging topic branches
            print("No nixpkgs change detected")
            return

        system("git update-ref refs/heads/nixos-unstable origin/nixos-unstable")
        merge_nixpkgs_topics_into_master()

    if was_head_changed(update_nixpkgs_and_merge_topics):
        print("nixpkgs upstream commit: ", check_output("git rev-parse origin/nixos-unstable", shell=True).decode().strip())
        return True
    return False


def update_home_manager():
    cd("~/nix/home-manager")
    was_head_changed("git pull --quiet")


def update_flake_input(flake_input: str):
    cd("~/nix/conf")
    old_hash = check_output("git hash-object flake.lock", shell=True)
    check_call(["nix", "flake", "update", flake_input])
    new_hash = check_output("git hash-object flake.lock", shell=True)
    if old_hash != new_hash:
        print(f"!!! Flake input {flake_input} changed")
        return True
    return False


class FailedBuildCounter:
    def __init__(self, host):
        self.filename = expanduser(f'~/nix/out/{host}/flake.failed.cnt')

    def reset(self):
        # Reset failed build counter
        try:
            os.remove(self.filename)
        except FileNotFoundError:
            pass

    def get(self):
        try:
            return int(open(self.filename).readline())
        except (FileNotFoundError, ValueError):
            return 0

    def increase(self):
        cnt = self.get()
        with open(self.filename, 'w') as f:
            print(cnt + 1, file=f)


def updated(host):
    updated = [
        update_nix_conf(),
        update_clonned_nixpkgs(),
        update_flake_input("ros2nix")
    ]
    # Update nixpkgs-stable, but don't trigger rebuild of unstable
    # stuff (i.e. result=0). `nix flake check` will then rebuild stable
    update_flake_input("nixpkgs-stable")
    if any(updated):
        FailedBuildCounter(host).reset()
        # Update selected overlays only if something else is updated
        update_flake_input("emacs-overlay")
        update_flake_input("home-manager")
        update_flake_input("home-manager-stable")
        update_flake_input("nix-index-database")

    return any(updated)


def failed_recently(host):
    cnt = FailedBuildCounter(host).get()
    return cnt > 0 and cnt <= 2


def print_switch_instructions(host, override_input: List[str]):
    print("Diff from current configuration:")
    # nix store diff-closures /run/current-system ~/nix/out/"$HOST"/flake
    system(f'dix /run/current-system ~/nix/out/{host}/flake')
    print(f"""
Run one of the following to switch the configuration:
  sudo nix-conf switch
  sudo nixos-rebuild switch {" ".join(override_input)}
    """)


def build(host, nix_args):
    cd("~/nix")

    override_input = []
    match host:
        case "steelpick": override_input += ['--override-input', 'nixpkgs', f'{os.getcwd()}/nixpkgs']

    os.makedirs(f"out/{host}", exist_ok=True)
    nix = 'nom' if os.isatty(sys.stdin.fileno()) else "nix"

    cmd = [
        nix,
        "build",
        *nix_args,
        "--out-link",
        expanduser(f"~/nix/out/{host}/flake"),
        *override_input,
        expanduser(
            f"~/nix/conf#nixosConfigurations.{host}.config.system.build.toplevel"
        ),
    ]
    status = run(cmd).returncode

    if status != 0:
        sys.exit(status)

    if check_output("hostname").decode().strip() == host:
        print_switch_instructions(host, override_input)

def publish_ci_update():
        cd("~/nix/nixpkgs")
        if run(["git", "diff", "--quiet", "wentasah/ci"]).returncode == 0:
            return False

        system("git push --force wentasah HEAD:ci")
        old_commit = check_output("git rev-parse history", shell=True).decode().strip()
        new_commit = check_output(f'git commit-tree -p {old_commit} -p HEAD "HEAD^{{tree}}" -m "update history branch"', shell=True).decode().strip()
        system(f'git update-ref refs/heads/history "{new_commit}" "{old_commit}"')
        system("git push wentasah history")
        return True


def notify_ci_update(attr):
        cd("~/nix/out/ci")
        if not os.path.exists(f'gcroots.old/result-{attr}'):
            warn(f"warning: {os.getcwd()}/gcroots.old/result-{attr} does not exist - not checking for changes")
            return
        if os.readlink(f'gcroots.old/result-{attr}') != os.readlink(f'gcroots/result-{attr}'):
            diff = check_output(f"dix gcroots.old/result-{attr} gcroots/result-{attr}", shell=True).decode().splitlines()
            diff_md = "\n".join(["    " + line for line in diff])
            try:
                check_call(["matrix-commander-rs", "--markdown", "-m", f"nix-conf CI update completed\n\n{diff_md}"])
            except CalledProcessError:
                # The above can fail with:
                # ERROR: matrix-commander: E150: room_send failed with error 'RoomSendError: M_TOO_LARGE event too large'.
                # Try sending a shorter message
                check_call(["matrix-commander-rs", "--markdown", "-m", f"nix-conf CI update completed (diff omitted - probably too long)"])


def switch(host, method='switch'):
    O = f'/home/{os.getenv("SUDO_USER", os.getenv("USER"))}/nix/out/{host}/flake'
    system(f'nix-env -p /nix/var/nix/profiles/system --set "{O}"')
    system(f'{O}/bin/switch-to-configuration {method}')


def push():
    cd("~/nix/nixpkgs")
    system("git push --force-with-lease wentasah master:master")
    system("git push wentasah nixos-unstable:nixos-unstable")
    cd("~/nix/conf")
    system("nix flake update nixpkgs --commit-lock-file")
    system("git push")


def add_commits(branch: str, commits: list[str]):
    cd("~/nix/nixpkgs")
    system(f'git checkout -B {branch} nixos-unstable')
    system(f'git cherry-pick {" ".join(commits)}')
    system(f'git-machete add')
    merge_nixpkgs_topics_into_master()


def add_pr(pr: int, branch: str):
    cd("~/nix/nixpkgs")
    system('git fetch origin master staging')
    system(f'git fetch origin refs/pull/{pr}/head')
    if run('git merge-base --is-ancestor nixos-unstable FETCH_HEAD', shell=True).returncode == 0:
        # Rebase PR's commits on nixos-unstable
        commits = check_output(f"gh pr view --json commits {pr} | jq --raw-output '.commits[].oid'", shell=True) \
            .decode().splitlines()
        add_commits(branch, commits)
    else:
        # Merge PR's commits
        system(f'git branch {branch} FETCH_HEAD')
        system(f'git-machete add {branch} --onto nixos-unstable')
        merge_nixpkgs_topics_into_master()
    system(f'git-machete anno --branch="{branch}" "PR #{pr} rebase=no push=no"')


def update(host, nix_args):
    if updated(host):
        build(host, nix_args)
    else:
        print("Nothing to do for update.")


def ci():
    if updated("ci") or failed_recently("ci"):
        cd("~/nix")

        shutil.rmtree("out/ci/gcroots.old", ignore_errors=True)
        try:
            shutil.copytree("out/ci/gcroots", "out/ci/gcroots.old", symlinks=True)
        except (FileNotFoundError, shutil.SameFileError):
            pass

        cd("~/nix/conf")
        check_call('nix flake update nixpkgs --override-input nixpkgs ~/nix/nixpkgs', shell=True)
        cmd = "nix-fast-build -j4 --no-nom --attic-cache wsh --out-link ~/nix/out/ci/gcroots/result"
        if not os.isatty(sys.stdin.fileno()):
            cmd = f"set -o pipefail; {cmd} --no-nom 2>&1 | tee ~/nix/out/ci/log"
        else:
            Path('../out/ci/log').touch() # ensure path exists

        status = run(cmd, shell=True).returncode

        if status == 0:
            shutil.copy("flake.lock", f'../out/ci/flake.lock')
            open(f'../out/ci/error', 'w') # truncate
            FailedBuildCounter("ci").reset()
        else:
            # Be careful to not delete the error file to not break hardlinks to it from icinga
            system("grep -E '(^error:|ERROR:nix_fast_build)' ../out/ci/log > ../out/ci/error")
            FailedBuildCounter("ci").increase()
            sys.exit(status)
    else:
        print("Nothing to do for update.")

    if publish_ci_update():
        notify_ci_update("x86_64-linux.steelpick")


def copy(host):
    update_clonned_nixpkgs()
    system(f"rsync ritchie:nix/out/ci/flake.lock ~/nix/conf")
    O = check_output(f'ssh ritchie readlink nix/out/ci/gcroots/result-x86_64-linux.{host}', shell=True).decode().strip()
    #system(f'nix copy -j16 -L --from ssh://ritchie "{O}"')
    system(f'nom build "{O}"')  # Copy packages from my binary cache
    dest = expanduser(f'~/nix/out/{host}/flake')
    try: os.remove(dest)
    except FileNotFoundError: pass
    os.symlink(O, dest)
    print_switch_instructions(host, [])  # TODO override


def get_nix_args(args):
    nix_args = args.nix_args
    if args.local:
        nix_args = ['--builders', '', '-j4'] + nix_args
    elif socket.gethostname() == 'steelpick':
        nix_args = ['--builders', 'ssh-ng://ritchie x86_64-linux,i686-linux - 16 2 nixos-test,benchmark,big-parallel,kvm'] + nix_args
    return nix_args


os.environ['GIT_AUTHOR_NAME'] = 'nix-conf'
os.environ['GIT_AUTHOR_EMAIL'] = 'nix-conf@localhost'
os.environ['GIT_COMMITTER_NAME'] = 'nix-conf'
os.environ['GIT_COMMITTER_EMAIL'] = 'nix-conf@localhost'

parser = argparse.ArgumentParser(
    prog="nix-conf",
    description="Tool to help maintaining my Nix and NixOS environment",
)
parser.add_argument(
    "--host",
    default=os.getenv("HOST", check_output("hostname").decode().strip()),
    help="For which host to build the configuration",
)

parser.add_argument("-v", "--verbose", action='store_true',
                    help="Print executed commands (default)")
parser.add_argument("-q", "--quiet", action='store_true',
                    help="Don't print executed commands")

parent_parser = argparse.ArgumentParser(add_help=False)
parent_parser.add_argument("-l", "--local", action='store_true',
                           help="Disable remote builders")
subparsers = parser.add_subparsers(required=True)

parser_update = subparsers.add_parser('update', parents=[parent_parser])
parser_update.add_argument("nix_args", nargs='*')
parser_update.set_defaults(func=lambda args: update(args.host, get_nix_args(args)))

parser_ci = subparsers.add_parser('ci', parents=[parent_parser])
parser_ci.set_defaults(func=lambda args: ci())

parser_build = subparsers.add_parser('build', parents=[parent_parser])
parser_build.add_argument("nix_args", nargs='*')
parser_build.set_defaults(func=lambda args: build(args.host, get_nix_args(args)))

parser_copy = subparsers.add_parser('copy')
parser_copy.set_defaults(func=lambda args: copy(args.host))

parser_merge = subparsers.add_parser('merge')
parser_merge.set_defaults(func=lambda args: merge_nixpkgs_topics_into_master())

parser_switch = subparsers.add_parser('switch')
parser_switch.set_defaults(func=lambda args: switch(args.host))

parser_boot = subparsers.add_parser('boot')
parser_boot.set_defaults(func=lambda args: switch(args.host, method='boot'))

parser_push = subparsers.add_parser('push')
parser_push.set_defaults(func=lambda args: push())

parser_add_pr = subparsers.add_parser('add-pr')
parser_add_pr.add_argument("pr_num", type=int)
parser_add_pr.add_argument("branch", type=str)
parser_add_pr.set_defaults(func=lambda args: add_pr(args.pr_num, args.branch))

parser_add_commits = subparsers.add_parser('add-commits')
parser_add_commits.add_argument("branch", type=str)
parser_add_commits.set_defaults(func=lambda args: add_commits(args.branch))

args = parser.parse_args()
verbose = not args.quiet
args.func(args)
