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

from tqdm import tqdm
from pathlib import Path
import subprocess
import json


PAGINATION_STEP=3
PAGES=800


GRAPHQL_QUERY = """
query($endCursor: String) {
  repository(owner: "NixOS", name: "nixpkgs") {
    pullRequests(
      first: """ + str(PAGINATION_STEP) + """,
      after: $endCursor,
      orderBy: {
        field: CREATED_AT,
        direction: DESC
    	},
      labels: ["8.has: module (update)"],
    ) {
      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:
    try:
        return subprocess.check_output(["gh", *args], text=True)
    except subprocess.CalledProcessError as e:
        print(e)
        return None


def find_prs_without_new_module_tag():
    hasNextPage = True
    endCursor = None

    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}', "-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']]
        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'])]

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

        pbar.update(PAGINATION_STEP)

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

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"]} - {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("")
    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__":
    find_prs_without_new_module_tag()