import os
import re
import json
import shutil
import argparse
import subprocess
from tqdm import tqdm
from multiprocessing import Process, current_process

import pydot
import networkx as nx
from networkx.readwrite import json_graph

from utils.json_utils import save_jsonl, load_jsonl


# Global configuration for file paths and directories
PROJECT_ROOT = os.getcwd()
OUTPUT_BASE_DIR = f"{PROJECT_ROOT}/datasets/processed/primevul"
OUTPUT_CHUNKS_DIR = f"{OUTPUT_BASE_DIR}/data_chunks/cpg"
OUTPUT_GRAPHS_DIR = f"{OUTPUT_BASE_DIR}/data_graphs/cpg"

# Joern tool paths
JOERN_PARSE_PATH = os.path.join(PROJECT_ROOT, "utils/joern-cli/joern-parse")
JOERN_EXPORT_PATH = os.path.join(PROJECT_ROOT, "utils/joern-cli/joern-export")


def process_data(chunk_index):
    # Get the process ID for naming the output files and temporary directories
    process_id = current_process().name
    
    # Read the data samples
    data_samples = []
    with open(input_jsonl_path, "r") as file:
        for i, line in enumerate(file):
            if i < chunk_index[0]: # Skip lines before the start index
                continue
            if i >= chunk_index[1]: # Stop reading after the end index
                break
            try: # Load the JSON object within chunk_index
                data_samples.append(json.loads(line.strip()))
            except json.JSONDecodeError as e:
                print(f"Error decoding JSON: {e}")
    
    # Create a list to store the graphs
    raw_graphs = []
    processed_graphs = []
    
    # Initialize the temporary directories and create them if they do not exist
    temp_dir = f"{OUTPUT_BASE_DIR}/temporary/{process_id}" # temp_dir is the directory to store temporary files (including temp.c, cpg.bin, and temp folder)
    os.makedirs(temp_dir, exist_ok=True)
    
    # Set the current working directory to the temp_dir so joern-parse and joern-export can put the temp files in temp_dir
    os.chdir(temp_dir)
    
    # Define the variable used by joern-parse to process and save the output into temp_dir
    temp_file = os.path.join(temp_dir, "temp.c") # define the path to the directory of temp.c file
    graph_type = "cpg" # the type of graph to export, can be "cpg", "cfg", or "ast"
    representation = f"--repr={graph_type}"
    output_dir = f"--out=temp"
    output_format = "--format=dot"
    
    for data_sample in tqdm(data_samples, desc=f"processing chunk {process_id}"):
        # Write the C code to a temp.c file for each data sample
        with open(temp_file, "w") as file:
            file.write(data_sample["func"])
        
        # Run the command to parse the input file into cpg.bin under the home directory
        parse_result = subprocess.run([JOERN_PARSE_PATH, temp_file], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        # Run the command to export the parsed graph into a .dot file
        export_result = subprocess.run([JOERN_EXPORT_PATH, representation, output_dir, output_format], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        # Load the .dot file using pydot
        (graph,) = pydot.graph_from_dot_file(os.path.join(temp_dir, "temp/temp.c/_global_.dot"))
        
        # Convert to a NetworkX graph
        nx_graph = nx.nx_pydot.from_pydot(graph)
        
        # Save the raw graph (with all attributes) to the list of raw graphs
        raw_graphs.append(json_graph.node_link_data(nx_graph))
        
        # Iterate over each node in the graph and keep only the "CODE" and "label" attributes
        for node in nx_graph.nodes:
            # Get the attributes of the current node
            attributes = nx_graph.nodes[node]
            # Only keep "CODE" and "label", remove others
            for attr in list(attributes):  # Convert to list to avoid changing dict size during iteration
                if attr not in ["CODE", "label"]:
                    del nx_graph.nodes[node][attr]
        
        # Iterate over all edges and keep only the "label" attribute
        for u, v, attr in nx_graph.edges(data=True):
            label_value = attr.get("label")  # Get the "label" attribute if it exists
            # Clear all attributes and re-add the "label" if it exists
            attr.clear()
            if label_value is not None:
                attr["label"] = label_value
        
        # Append the processed graph to the list of processed graphs
        processed_graphs.append(json_graph.node_link_data(nx_graph))
        
        # Delete the temp.c file, cpg.bin, and temp folder under the temp_dir
        os.remove(temp_file)
        os.remove(os.path.join(temp_dir, "cpg.bin"))
        shutil.rmtree(os.path.join(temp_dir, "temp"))
    
    # Save the results
    os.makedirs(os.path.join(OUTPUT_CHUNKS_DIR, "raw"), exist_ok=True)
    os.makedirs(os.path.join(OUTPUT_CHUNKS_DIR, "processed"), exist_ok=True)
    save_jsonl(os.path.join(OUTPUT_CHUNKS_DIR, f"raw/{process_id}.jsonl"), raw_graphs)
    save_jsonl(os.path.join(OUTPUT_CHUNKS_DIR, f"processed/{process_id}.jsonl"), processed_graphs)


if __name__ == "__main__":
    # Create an argument parser
    parser = argparse.ArgumentParser(description="Run different sets of experiments")
    
    parser.add_argument(
        "--file_name", type=str,
        choices=["primevul_test.jsonl", "primevul_train.jsonl", "primevul_valid.jsonl"],
        help="The file name of the dataset to process.",
    )
    
    # Parse the arguments
    args = parser.parse_args()
    
    # Define the input and output file paths based on the provided file name
    input_jsonl_path = os.path.join(PROJECT_ROOT, f"datasets/raw/primevul/{args.file_name}")
    
    # Calculate the number of data samples
    with open(input_jsonl_path, "r") as file:
        number_of_data_samples = sum(1 for _ in file)
    
    # Split the data_samples into chunks
    num_chunks = os.cpu_count()
    chunk_size = number_of_data_samples // num_chunks
    data_chunks_index = [[i, min(i + chunk_size, number_of_data_samples)] for i in range(0, number_of_data_samples, chunk_size)]
    
    # Create and start processes for each chunk
    processes = []
    for i, chunk_index in enumerate(data_chunks_index):
        p = Process(target=process_data, args=(chunk_index,))
        processes.append(p)
        p.start()
    
    # Wait for all processes to finish
    for p in processes:
        p.join()
    
    print("Processing completed!")
    
    # Merge the data chunks into memory lists, then save to avoid overwriting
    for graph_type in ["raw", "processed"]:
        # Initialize an empty list to hold the merged data from all chunks
        merged_data = []
        chunk_dir = os.path.join(OUTPUT_CHUNKS_DIR, graph_type)
        chunk_files = [f for f in os.listdir(chunk_dir) if f.endswith('.jsonl')]
        chunk_files.sort(key=lambda x: int(re.search(r'\d+', x).group()) if re.search(r'\d+', x) else 0)
        
        # Load each chunk file and extend the merged_data list
        for chunk_file in tqdm(chunk_files, desc=f"Loading {graph_type} chunks"):
            chunk_file_path = os.path.join(chunk_dir, chunk_file)
            merged_data.extend(load_jsonl(chunk_file_path))
        
        # Save the fully extended list to distinct files
        output_jsonl_path = os.path.join(OUTPUT_GRAPHS_DIR, f"{graph_type}/{args.file_name}")
        os.makedirs(os.path.dirname(output_jsonl_path), exist_ok=True)
        save_jsonl(output_jsonl_path, merged_data)
        print(f"Merged {graph_type} saved to: {output_jsonl_path}")
    
    # Delete the temporary directory after merging
    if os.path.exists(OUTPUT_CHUNKS_DIR):
        shutil.rmtree(OUTPUT_CHUNKS_DIR)
        shutil.rmtree(os.path.join(OUTPUT_BASE_DIR, "temporary"))
        print(f"Deleted temporary directory.")