#! /usr/bin/env nix-shell
#! nix-shell -i python3 -p python3 nix git

import itertools
import json
import os
import regex

from collections import defaultdict
from pathlib import Path
from textwrap import dedent

phase_hooks = [
    ('buildPhase', 'preBuild', 'postBuild'),
    ('checkInstallPhase', 'preCheckInstall', 'postCheckInstall'),
    ('checkPhase', 'preCheck', 'postCheck'),
    ('configurePhase', 'preConfigure', 'postConfigure'),
    ('distPhase', 'preDist', 'postDist'),
    ('installPhase', 'preInstall', 'postInstall'),
    ('patchPhase', 'prePatch', 'postPatch'),
    ('unpackPhase', 'preUnpack', 'postUnpack'),
]


def provide_paths(rootdir):
    for parent_dir, dirs, files in os.walk(rootdir):
        for basename in files:
            if not basename.endswith('.nix'):
                continue
            path = Path(parent_dir, basename)
            try:
                with open(path) as file:
                    pass
                yield path
            except FileNotFoundError as e:
                pass


class AbortEditException(Exception):
    pass


def try_to_fix_file_runhooks(path, dryrun) -> list[str]:
    with open(path) as file:
        raw_content = file.read()

    edited_hooks = []
    for phase_hook in phase_hooks:
        if (edit := insert_phase_if_not_exists_multiline(path, raw_content, *phase_hook)) is not None:
            edited_hooks.append(phase_hook[0])
            if not dryrun:
                with open(path, 'w') as file:
                    file.write(edit)
                    raw_content = edit
        if (edit := insert_phase_if_not_exists_singleline(path, raw_content, *phase_hook)) is not None:
            edited_hooks.append(phase_hook[0])
            if not dryrun:
                with open(path, 'w') as file:
                    file.write(edit)
                    raw_content = edit

    return edited_hooks


def insert_phase_if_not_exists_multiline(path, text, phase, pre_phase, post_phase) -> str | None:
    pattern = regex.compile(f"{phase} ?= ?''(.|\n)+?'';", regex.MULTILINE)
    match = regex.search(pattern, text)
    if match is None or 'runHook' in match.group():
        return None

    inner_lines = match.group().split('\n')
    indent = ''.join(
        itertools.takewhile(
            lambda x: x == ' ',
            next(line for line in inner_lines[1:-1] if line.strip() != '')
        )
    )
    replacement ='\n'.join(
        inner_lines[0:1]
        + [indent + f'runHook {pre_phase}', '']
        + inner_lines[1:-1]
        + ['', indent + f'runHook {post_phase}']
        + inner_lines[-1:]
    )

    result = text[0:match.start()] + replacement + text[match.end():]

    # Sanity check for string concatenations with optional string concats,
    # and (TODO) contained string interpolations which contain "'';"
    if regex.search(r"''(?!\$)", '\n'.join(inner_lines[1:-1])) is not None:
        colorizedResult = text[0:match.start()] + '\u001b[32m' + replacement + '\u001b[0m' + text[match.end():]
        print('-----\n' + colorizedResult + '\n------')
        if input(f"Accept {path} ({phase})? [N/y] ") not in 'Yy':
            print(f"Rejecting {path} ({phase})")
            raise AbortEditException
        print()

    return result


def insert_phase_if_not_exists_singleline(path, text, phase, pre_phase, post_phase) -> str | None:
    pattern = regex.compile(f"{phase} ?= ?\"(.+?)\";")
    match = regex.search(pattern, text)
    if match is None:
        return None

    base_indent_num = ''.join(reversed(text[:match.start()])).find('\n')
    phase_content = [
        f"{(base_indent_num + 2) * ' '}runHook {pre_phase}",
        "",
        f"{(base_indent_num + 2) * ' '}{match.group(1)}",
        "",
        f"{(base_indent_num + 2) * ' '}runHook {post_phase}",
    ]

    replacement = '\n'.join(
        [ f"{phase} = ''" ]
        + phase_content
        + [ f"{base_indent_num * ' '}'';" ]
    )

    result = text[0:match.start()] + replacement + text[match.end():]

    return result


def build_packages_path_dict(path_to_nixpkgs_root) -> dict[str, list[str]]:
    print("Generating package index...")
    os.system(f'nix-env -f "{path_to_nixpkgs_root}" -I nixpkgs="{path_to_nixpkgs_root}" -qa --meta --json --show-trace --arg config "import {path_to_nixpkgs_root}/pkgs/top-level/packages-config.nix" >/tmp/packages.json')

    with open("/tmp/packages.json") as file:
        packages_json = json.load(file)

    result = defaultdict(list)
    for key, package in packages_json.items():
        if 'position' in package['meta']:
            path = Path(package['meta']['position'].split(':')[0]).resolve()
            result[path].append(key)
    return dict(result)


def stage_and_commit_changes(package_path_dict, path, edited_hooks, dryrun):
    package_names = package_path_dict.get(path.resolve(), [])
    if len(package_names) == 0:
        while input(dedent(f"""
            Could not find package name for {path}.
            Please commit manually and press y to continue:
        """)) not in 'Yy':
            pass
    elif len(package_names) == 1:
        package_name = package_names[0]
        print(f'[GIT COMMIT] {package_name}: added missing `runHooks` for {", ".join(edited_hooks)}')
    else:
        # TODO: cat the file and let user choose a number
        #       if package count below threshold (including abort option)
        #       else ask user to commit manually
        input(f'{path} contains many packages: {", ".join(package_names)}, press anything to continue...')


if __name__ == '__main__':
    dryrun = True
    package_path_dict = build_packages_path_dict("../..")
    for path in provide_paths('.'):
        try:
            edited_hooks = try_to_fix_file_runhooks(path, dryrun)
            if len(edited_hooks) > 0:
                stage_and_commit_changes(package_path_dict, path, edited_hooks, dryrun)
        except AbortEditException:
            pass