# SPDX-FileCopyrightText: (c) 2024 EvalPlus Team
#
# SPDX-License-Identifier: Apache-2.0

import json
import os
import subprocess

import git
import tempdir
from fire import Fire
from tqdm.auto import tqdm

from scripts.curate.dataset_ensemble_clone import get_files_to_include
from scripts.curate.utility import lang2suffix

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))


# dataset_path is the dataset generated by dataset_ensemble_clone.py
def main():
    with open("scripts/cherrypick/lists.json") as f:
        lists = json.load(f)

    lang_suffix = lang2suffix["python"]
    repos = lists["python"]
    for repo in tqdm(repos):
        repo_name = repo["repo"]
        commit_sha = repo["commit_sha"]
        entrypoint = repo["entrypoint_path"]
        print(f"Visiting https://github.com/{repo_name}/tree/{commit_sha}")

        with tempdir.TempDir() as temp_dir:
            gh_repo = git.Repo.clone_from(
                f"https://github.com/{repo_name}.git",
                temp_dir,
            )
            gh_repo.git.checkout(commit_sha)
            command_list = f"pydeps {entrypoint} --show-deps --no-show".split()
            # cd `temp_dir`` and capture the output json
            output = subprocess.check_output(command_list, cwd=temp_dir)
            dependencies = json.loads(output)

            output_dependency = {}
            mod2path = {}

            abs_prefix = os.path.join(temp_dir, entrypoint)
            for mod, v in dependencies.items():
                if v["path"] and v["path"].startswith(abs_prefix):
                    mod2path[mod] = v["path"]

            for mod, v in dependencies.items():
                if v["path"] and v["path"].startswith(abs_prefix):
                    if "imports" in v:
                        relative_path = os.path.relpath(v["path"], temp_dir)
                        output_dependency[relative_path] = [
                            os.path.relpath(mod2path[imp], temp_dir)
                            for imp in v["imports"]
                            if imp in mod2path
                        ]

            files_to_include = get_files_to_include(gh_repo, entrypoint, lang_suffix)
            for path, _ in files_to_include:
                if path not in output_dependency:
                    output_dependency[path] = []

            assert output_dependency, "Empty output_dependency"
            repo["dependency"] = output_dependency

    with open(os.path.join(CURRENT_DIR, "data", "python.json"), "w") as f_out:
        json.dump({"python": repos}, f_out)


if __name__ == "__main__":
    Fire(main)
