#!/usr/bin/env python3
from huggingface_hub import HfApi
import argparse
import os

api = HfApi()

GOD_TOKEN = os.environ.get("GOD_TOKEN")

def close_or_merge_pr(pr_link, merge_all_open):
    # Logic to close or merge the pull request on Hugging Face Hub
    repo_id = pr_link.replace("https://huggingface.co/", "")
    repo_id = repo_id.replace("https://hf.co/", "")
    repo_id = repo_id.replace("hf.co/", "")
    repo_id = repo_id.replace("huggingface.co/", "")

    discussion_num = int(repo_id.split("/")[-1])
    repo_id = "/".join(repo_id.split("/")[:2])
    discussion = api.get_discussion_details(repo_id, discussion_num=discussion_num)

    if discussion.status == "merged":
        print(f"PR: {pr_link} was merged")
    elif discussion.status == "closed":
        print(f"PR: {pr_link} was closed")
    elif discussion.status == "open":
        print(f"PR: {pr_link} was left open")

        if merge_all_open:
            api.merge_pull_request(repo_id, discussion_num, token=GOD_TOKEN)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Close or merge a pull request on the Hugging Face Hub")
    parser.add_argument("--pr_link", type=str, required=True, help="Link to the pull request")
    parser.add_argument("--merge_all_open", action="store_true", help="Merge all open pull requests")

    args = parser.parse_args()

    close_or_merge_pr(args.pr_link, args.merge_all_open)