#! /usr/bin/env nix-shell
#! nix-shell -i python3 -p python3 gh python3Packages.tqdm

import argparse
import json
import subprocess
import time

from tqdm import tqdm

PAGINATION_STEP = 3
PAGES = 800
MAX_RETRIES = 8
RETRY_BACKOFF_SECONDS = 3


GRAPHQL_QUERY = (
    """
query($endCursor: String, $states: [PullRequestState!]) {
  repository(owner: "NixOS", name: "nixpkgs") {
    pullRequests(
      first: """
    + str(PAGINATION_STEP)
    + """,
      after: $endCursor,
      orderBy: {
        field: CREATED_AT,
        direction: DESC
    	},
      labels: ["8.has: module (update)"],
      states: $states,
    ) {
      edges {
        node {
          title
          number
          state
          createdAt
          author {
            login
          }
          files(first: 20) {
            edges {
              node {
                path
                changeType
              }
            }
          }
          labels(first: 20) {
            edges {
              node {
                name
              }
            }
          }
        }
      }
      pageInfo {
        endCursor
        hasNextPage
      }
    }
  }
}
"""
)


def _gh(args: list[str]) -> str | None:
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            return subprocess.check_output(
                ["gh", *args], text=True, stderr=subprocess.DEVNULL
            )
        except subprocess.CalledProcessError as e:
            if attempt == MAX_RETRIES:
                print(e)
                return None
            time.sleep(RETRY_BACKOFF_SECONDS * attempt)
    return None


def parse_args():
    parser = argparse.ArgumentParser()
    state_group = parser.add_mutually_exclusive_group()
    state_group.add_argument("--open", action="store_true", help="Only show open PRs")
    state_group.add_argument(
        "--merged", action="store_true", help="Only show merged PRs"
    )
    state_group.add_argument(
        "--closed", action="store_true", help="Only show closed (unmerged) PRs"
    )
    parser.add_argument(
        "--before",
        type=int,
        metavar="PR_NUMBER",
        help="Only show PRs with a number lower than this (i.e. created earlier)",
    )
    parser.add_argument(
        "--after",
        type=int,
        metavar="PR_NUMBER",
        help="Only show PRs with a number higher than this (i.e. created later)",
    )
    args = parser.parse_args()
    if args.before is not None and args.after is not None:
        assert args.after < args.before, "--after must be lower than --before"
    return args


def find_prs_without_new_module_tag(
    state_filter: str | None, before: int | None, after: int | None
):
    hasNextPage = True
    endCursor = None

    states_args = [] if state_filter is None else ["-F", f"states[]={state_filter}"]

    pbar = tqdm(total=PAGES * PAGINATION_STEP)
    for i in range(PAGES):
        prs = _gh(
            [
                "api",
                "graphql",
                "-F",
                f"endCursor={'null' if endCursor is None else endCursor}",
                *states_args,
                "-f",
                f"query={GRAPHQL_QUERY}",
            ]
        )
        if prs is None:
            print("gh exited with error")
            exit(1)
        data = json.loads(prs)["data"]["repository"]["pullRequests"]
        endCursor = data["pageInfo"]["endCursor"]
        hasNextPage = data["pageInfo"]["hasNextPage"]
        prs = [pr["node"] for pr in data["edges"]]
        page_numbers = [pr["number"] for pr in prs]
        prs = [
            {
                "createdAt": pr["createdAt"],
                "state": pr["state"],
                "title": pr["title"],
                "url": "https://github.com/NixOS/nixpkgs/pull/" + str(pr["number"]),
                "number": pr["number"],
                "author": pr["author"]["login"] if pr["author"] is not None else "???",
                "files": [
                    file["node"]
                    for file in pr.get("files", dict()).get("edges", tuple())
                ],
                "labels": [
                    label["node"]["name"]
                    for label in pr["labels"].get("edges", tuple())
                ],
            }
            for pr in prs
        ]
        prs = [
            pr
            for pr in prs
            if not any(label == "8.has: module (new)" for label in pr["labels"])
        ]
        prs = [
            pr
            for pr in prs
            if any(
                file["changeType"] == "ADDED"
                and file["path"].startswith("nixos/modules")
                for file in pr["files"]
            )
        ]

        # Results are ordered by number/createdAt descending, so once every PR
        # in a page is below the `after` bound, no later page can contain a match.
        if (
            after is not None
            and page_numbers
            and all(number <= after for number in page_numbers)
        ):
            hasNextPage = False

        if before is not None:
            prs = [pr for pr in prs if pr["number"] < before]
        if after is not None:
            prs = [pr for pr in prs if pr["number"] > after]

        for pr in reversed(prs):
            print_pr(pr)

        pbar.update(PAGINATION_STEP)

        if not hasNextPage:
            break


STATUS_COLORS = {
    "OPEN": "\033[1;32mOPEN\033[0m",
    "MERGED": "\033[1;35mMERGED\033[0m",
    "CLOSED": "\033[1;31mCLOSED\033[0m",
}

CHANGE_TYPE_COLORS = {
    "ADDED": "\033[1;32mADDED\033[0m",
    "MODIFIED": "\033[1;33mMODIFIED\033[0m",
    "RENAMED": "\033[1;36mRENAMED\033[0m",
    "DELETED": "\033[1;31mDELETED\033[0m",
}


def color_file_path(path: str, changeType: str) -> str:
    if path == "nixos/modules/module-list.nix":
        return f"\033[1;31m{path}\033[0m"
    elif (
        changeType == "ADDED"
        and path.startswith("nixos/modules/")
        and path.endswith(".nix")
    ):
        return f"\033[1;32m{path}\033[0m"
    elif path == "nixos/tests/all-tests.nix":
        return f"\033[1;31m{path}\033[0m"
    elif (
        changeType == "ADDED"
        and path.startswith("nixos/tests/")
        and path.endswith(".nix")
    ):
        return f"\033[1;34m{path}\033[0m"
    elif path.endswith(".md"):
        return f"\033[1;33m{path}\033[0m"
    else:
        return path


def color_backport_label(title: str) -> str:
    if title.startswith("[Backport release-"):
        return "\033[31m" + title + "\033[0m"
    else:
        return title


PROMETHEUS_EXPORTERS_DIR = "nixos/modules/services/monitoring/prometheus/exporters/"
PROMETHEUS_EXPORTERS_FILE = "nixos/modules/services/monitoring/prometheus/exporters.nix"


def warnings_for_pr(pr: dict[str, any]) -> list[str]:
    warnings = []

    has_new_module = any(
        file["changeType"] == "ADDED"
        and file["path"].startswith("nixos/modules/")
        and file["path"].endswith(".nix")
        for file in pr["files"]
    )

    is_prometheus_exporter = any(
        file["path"].startswith(PROMETHEUS_EXPORTERS_DIR) for file in pr["files"]
    ) or any(file["path"] == PROMETHEUS_EXPORTERS_FILE for file in pr["files"])

    module_list_touched = any(
        file["path"] == "nixos/modules/module-list.nix" for file in pr["files"]
    )
    if has_new_module and not module_list_touched and not is_prometheus_exporter:
        warnings.append(
            "New module added, but nixos/modules/module-list.nix was not changed"
        )

    has_new_test = any(
        file["changeType"] == "ADDED"
        and file["path"].startswith("nixos/tests/")
        and file["path"].endswith(".nix")
        for file in pr["files"]
    )
    all_tests_touched = any(
        file["path"] == "nixos/tests/all-tests.nix" for file in pr["files"]
    )
    if has_new_test and not all_tests_touched and not is_prometheus_exporter:
        warnings.append("New test added, but nixos/tests/all-tests.nix was not changed")

    release_notes_touched = any(
        file["path"].startswith("nixos/doc/manual/release-notes")
        for file in pr["files"]
    )
    if has_new_module and not release_notes_touched:
        warnings.append(
            "New module added, but nothing in nixos/doc/manual/release-notes was changed"
        )

    return warnings


def print_pr(pr: dict[str, any]):
    tqdm.write(
        f"[{STATUS_COLORS.get(pr['state'], pr['state'])}|\033[1m{pr['number']}\033[0m] {pr['author']} - {pr['createdAt'][:10]} - {color_backport_label(pr['title'])}"
    )
    tqdm.write("\033[1;34m" + pr["url"] + "\033[0m")
    tqdm.write("")
    tqdm.write("Files:")
    for file in pr["files"]:
        tqdm.write(
            f"  [{CHANGE_TYPE_COLORS.get(file['changeType'], file['changeType'])}] {color_file_path(file['path'], file['changeType'])}"
        )
    tqdm.write("")
    warnings = warnings_for_pr(pr)
    for warning in warnings:
        tqdm.write(f"\033[1;31m!!! {warning} !!!\033[0m")
    if warnings:
        tqdm.write("")
    if any(
        file["path"].startswith("nixos/tests/") and file["changeType"] == "ADDED"
        for file in pr["files"]
    ):
        tqdm.write(
            f'\033[1mgh pr edit {pr["number"]} --add-label "8.has: module (new),8.has: tests"\033[0m'
        )
    else:
        tqdm.write(
            f'\033[1mgh pr edit {pr["number"]} --add-label "8.has: module (new)"\033[0m'
        )
    tqdm.write("")
    tqdm.write(
        "------------------------------------------------------------------------------"
    )
    tqdm.write("")


if __name__ == "__main__":
    args = parse_args()
    state_filter = (
        "OPEN"
        if args.open
        else "MERGED"
        if args.merged
        else "CLOSED"
        if args.closed
        else None
    )
    find_prs_without_new_module_tag(state_filter, args.before, args.after)