# SPDX-FileCopyrightText: (c) 2024 EvalPlus Team
#
# SPDX-License-Identifier: Apache-2.0
import json
import os
import subprocess
from enum import Enum
from pathlib import Path
from typing import List, Tuple

import git
import pygraphviz as pgv
import tempdir
from fire import Fire
from pygraphviz import Node
from tqdm.auto import tqdm

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


class CrateType(Enum):
    LIB = ("lib",)
    BIN = "bin"


# Given a namespace, find its location in the file directory
def find_path(node: Node, temp_dir: str, entrypoint: str, type: CrateType) -> str:
    module_path = node.get_name()
    temp_dir = Path(temp_dir)
    entrypoint = Path(entrypoint)
    paths = module_path.split("::")

    current_path = entrypoint
    best_path = None

    for path in paths[1:]:
        current_path = current_path / path
        if (temp_dir / current_path / "mod.rs").exists():
            best_path = current_path / "mod.rs"
            continue
        elif (temp_dir / current_path).with_suffix(".rs").exists():
            best_path = current_path.with_suffix(".rs")
        break

    if not (best_path):
        if type == CrateType.LIB:
            return os.path.join(entrypoint, "lib.rs")
        elif type == CrateType.BIN:
            return os.path.join(entrypoint, "main.rs")
    return str(best_path)


# Find all buildable packages within the entrypoint
def find_packages(temp_dir: str, entrypoint: str) -> List[Tuple[str, str, CrateType]]:
    command_list = f"cargo metadata -q --no-deps --format-version=1".split()
    output = subprocess.check_output(command_list, cwd=temp_dir)
    decoded_output = output.decode("utf-8")

    result = json.loads(decoded_output)
    packages = result["packages"]
    results = []
    for package in packages:
        for target in package["targets"]:
            if not (os.path.join(temp_dir, entrypoint) in target["src_path"]):
                continue
            if target["kind"][0] == "lib":
                results.append((target["name"], target["src_path"], CrateType.LIB))
            elif target["kind"][0] == "bin":
                results.append((target["name"], target["src_path"], CrateType.BIN))
    return results


def is_cargo_modules_available() -> bool:
    try:
        command_list = ["cargo-modules", "--help"]
        subprocess.run(command_list, capture_output=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False


# dataset_path is the dataset generated by dataset_ensemble_clone.py
def main():
    if not (is_cargo_modules_available()):
        print("cargo-modules tool not found, exiting...")
        return

    with open("scripts/cherrypick/lists.json") as f:
        lists = json.load(f)

    repos = lists["rust"]
    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)
            packages = find_packages(temp_dir, entrypoint)
            mapping = {}
            dependency = {}
            edges = []
            for package_name, src_path, crate_type in packages:
                if "_" in package_name:
                    package_name = "-".join(package_name.split("_"))
                analysis_param = ""

                # Run cargo-modules tool
                if crate_type == CrateType.BIN:
                    analysis_param = f"--bin {package_name}"
                else:
                    analysis_param = f"--lib"
                command_list = f"cargo modules dependencies --package {package_name} {analysis_param} --cfg-test --no-sysroot --no-traits --no-types".split()
                entrypoint = os.path.dirname(os.path.relpath(src_path, temp_dir))

                # cd `temp_dir`` and capture the output json
                output = subprocess.check_output(command_list, cwd=temp_dir)
                decoded_output = output.decode("utf-8")

                # Parse graph and get mapping
                graph = pgv.AGraph(string=decoded_output)
                for node in graph.nodes():
                    node_name = node.get_name()
                    mapping[node_name] = find_path(
                        node, temp_dir, entrypoint, crate_type
                    )
                    if mapping[node_name] in dependency:
                        continue
                    dependency[mapping[node_name]] = set()

                # Save edges for later
                for start_node, end_node in graph.edges():
                    edges.append((start_node.get_name(), end_node.get_name()))

            # Parse every edge for dependencies
            for start_name, end_name in edges:
                if (
                    not (start_name in mapping and end_name in mapping)
                    or mapping[start_name] == mapping[end_name]
                ):
                    continue
                dependency[mapping[start_name]].add(mapping[end_name])

            for key in dependency:
                dependency[key] = list(dependency[key])

            repo["dependency"] = dependency

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


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