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

import json
import os
from pathlib import Path
from typing import List, Union

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

from scripts.curate.utility import lang2suffix

POTENTIAL_INCLUDE_DIR_NAMES = [
    "include",
    "src",
    "lib",
    "libs",
    "library",
    "libraries",
    "inc",
]

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


def get_potential_include_dirs(temp_dir: str) -> List[Path]:
    include_dir_list = [Path(temp_dir)]
    for potential_include_dir_name in POTENTIAL_INCLUDE_DIR_NAMES:
        for dir in Path(temp_dir).rglob(potential_include_dir_name):
            include_dir_list.append(dir)

    return include_dir_list


def search_include(
    entrypoint: str,
    temp_dir: str,
    cpp_file: Path,
    include: str,
    additional_libs: List[Path] = [],
) -> Union[Path, None]:
    for possible_path in [cpp_file.parent] + additional_libs:
        if not (entrypoint in str(possible_path)):
            continue
        if (possible_path / include).exists():
            return (
                (possible_path / include)
                .resolve()
                .relative_to(Path(temp_dir))
                .as_posix()
            )


def get_dependencies(temp_dir: str, entrypoint: str, cpp_file: Path):
    user_defined_includes = []
    with open(cpp_file) as f:
        content = f.read()
        for line in content.split("\n"):
            line = line.strip()
            if line.startswith("#include"):
                include = line.split()[1]
            elif line.split()[:2] == ["#", "include"]:
                include = line.split()[2]
            else:
                continue

            if (include.startswith('"') or include.startswith("<")) and "." in include:
                include = include[1:-1]
                rela_include = search_include(
                    entrypoint,
                    temp_dir,
                    cpp_file,
                    include,
                    additional_libs=get_potential_include_dirs(temp_dir),
                )
                if rela_include:
                    user_defined_includes.append(rela_include)

    return user_defined_includes


# 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["cpp"]
    repos = lists["cpp"]
    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}")

        if repo.get("dependency"):
            print(f"Skipping {repo_name} as it already has dependency field.")
            continue

        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)

            output_dependency = {}
            abs_prefix = os.path.join(temp_dir, entrypoint)
            for cpp_ext in lang_suffix:
                for cpp_file in Path(abs_prefix).rglob("*" + cpp_ext):
                    dependencies = get_dependencies(temp_dir, entrypoint, cpp_file)
                    if len(dependencies) > 0:
                        if "" in dependencies:
                            dependencies.remove("")
                        output_dependency[
                            cpp_file.relative_to(Path(temp_dir)).as_posix()
                        ] = dependencies
                    else:
                        output_dependency[
                            cpp_file.relative_to(Path(temp_dir)).as_posix()
                        ] = []

            repo["dependency"] = output_dependency

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


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