import pickle
import subprocess
import os
import re
import networkx as nx
import json

# call graph generator:
# https://github.com/vitsalis/PyCG

def generate_call_graph(project_path):
    json_path = project_path + '/project_call_graph.json'
    if os.path.exists(json_path):
        with open(json_path, 'r') as f:
            json_data = json.load(f)
            return json_data

    """Generate call graph by using PyCG"""
    os.chdir(project_path)

    # Find all Python files in current directory and subdirectories
    python_files = []
    for root, _, files in os.walk("."):
        for file in files:
            if file.endswith(".py"):
                python_files.append(os.path.join(root, file))

    python_files_str = " ".join(python_files)

    # Run pycg command
    command = f"pycg --max-iter 1 {python_files_str} --fasten -o {project_path}/project_call_graph.json"
    subprocess.run(command, shell=True)

    # load the result json file
    with open(json_path, 'r') as f:
        json_data = json.load(f)

    return json_data


def read_function_content(source_file, first_line, last_line):
    """Read the content of a function from a source file."""
    with open(source_file, 'r', errors='ignore') as f:
        lines = f.readlines()
    return "".join(lines[first_line - 1:last_line])  # Line numbering starts from 1


def define_node_type(namespace):
    # Pattern for functions
    function_pattern = r".+\(\)$"
    # Pattern for folders
    folder_pattern = r"/[\w\d.]+/"
    # Pattern for classes
    class_pattern = r"/[\w\d.]+/[\w\d.]+"

    if re.search(function_pattern, namespace):
        node_type = 'function'
    elif re.fullmatch(folder_pattern, namespace):
        node_type = 'file'
    elif re.fullmatch(class_pattern, namespace):
        node_type = 'class'
    else:
        print("It doesn't match any known categories.")
        print(namespace)
        node_type = None

    return node_type


def flatten_node_info(call_graph_data):
    """Extract and format node info"""
    os.chdir(project_path)
    # Mapping of node id to details
    node_map = {}
    # call_graph_data = json_data
    # Gather internal nodes and their function content
    for module, details in call_graph_data['modules']['internal'].items():
        source_file = details['sourceFile']
        for node_id, namespace_detail in details['namespaces'].items():
            metadata = namespace_detail['metadata']
            first_line = metadata['first']
            last_line = metadata['last']

            if first_line is None or last_line is None:
                content = None
            else:
                content = read_function_content(source_file, first_line, last_line)

            node_map[node_id] = {
                'sourceFile': source_file,
                'first': first_line,
                'last': last_line,
                'namespace': namespace_detail['namespace'],
                'content': content,
                'type': define_node_type(namespace_detail['namespace'])
            }

    return node_map


def format_internal_call_graph(node_map, raw_call_graph, project_path):
    """format call graph"""
    results = []
    folder_content = {}
    base_path = os.path.basename(project_path)
    # generate folder -> function edges
    for node_id, node in node_map.items():
        if node['type'] == 'file':
            folder_content[node['sourceFile'].split('/')[-1]] = node['content']
            continue

        relationship = {
            "source": {
                "label": node['sourceFile'].split('/')[-1],
                "file": base_path + '/' + node['sourceFile'],
                "type": 'file',
                "content": None
            },
            "target": {
                "label": re.sub(r'\(\)', '', node['namespace'].split('/')[-1]),
                "file": base_path + '/' + node['sourceFile'],
                "type": node['type'],
                "content": node['content']
            }
        }

        results.append(relationship)

    for edge in results:
        if edge['source']['type'] == 'file':
            edge['source']['content'] = folder_content[edge['source']['label']]

    # generate function -> function edges
    for source_node, target_node, _ in raw_call_graph['graph']['internalCalls']:

        source_node = node_map[source_node]
        target_node = node_map[target_node]

        if source_node['type'] == 'file':
            relationship = {
                "source": {
                    "label": source_node['sourceFile'].split('/')[-1],
                    "file": base_path + '/' + source_node['sourceFile'],
                    "type": source_node['type'],
                    "content": source_node['content']
                },
                "target": {
                    "label": re.sub(r'\(\)', '', target_node['namespace'].split('/')[-1]),
                    "file": base_path + '/' + target_node['sourceFile'],
                    "type": target_node['type'],
                    "content": target_node['content']
                }
            }
        else:

            relationship = {
                "source": {
                    "label": re.sub(r'\(\)', '', source_node['namespace'].split('/')[-1]),
                    "file": base_path + '/' + source_node['sourceFile'],
                    "type": source_node['type'],
                    "content": source_node['content']
                },
                "target": {
                    "label": re.sub(r'\(\)', '', target_node['namespace'].split('/')[-1]),
                    "file": base_path + '/' + target_node['sourceFile'],
                    "type": target_node['type'],
                    "content": target_node['content']
                }
            }
        results.append(relationship)

    return results


def generate_folder_structure(source_dir, base_dir=None):
    if base_dir is None:
        base_dir = os.path.dirname(source_dir)
    relationships = []

    for item in os.listdir(source_dir):

        if item.startswith('.') or item in ['dist', 'vendor']:
            continue

        item_path = os.path.join(source_dir, item)

        # Calculate the relative paths for source and target
        rel_source_path = os.path.relpath(source_dir, base_dir)
        rel_item_path = os.path.relpath(item_path, base_dir)

        relationship = {
            "source": {
                "label": os.path.basename(source_dir),
                "file": rel_source_path,
                "type": 'folder',
                "content": None
            },
            "target": {
                "label": item,
                "file": rel_item_path,
                "type": 'file',
                "content": None
            }
        }

        if os.path.isfile(item_path):
            with open(item_path, 'r', encoding='utf-8', errors='ignore') as file:
                file_content = file.read()
            relationship['target']['content'] = '"' + file_content + '"'

            file_extension = os.path.splitext(item_path)[1].lower()
            image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tif', '.tiff', '.webp', '.svg']
            if file_extension in image_extensions:
                relationship['target']['type'] = 'image'
        else:
            relationship['target']['type'] = 'folder'

        relationships.append(relationship)

        if os.path.isdir(item_path):
            relationships += generate_folder_structure(item_path, base_dir)

    return relationships


def convert_json_to_digraph(call_graph, folder_structure):
    # combine all
    whole = call_graph + folder_structure

    # Create a directed graph
    G = nx.DiGraph()

    # Add edges to the directed graph from the combined json data

    for edge in whole:
        source_file = edge['source']['file']
        target_file = edge['target']['file']

        source_label = edge['source']['label'].split('/')[-1]
        target_label = edge['target']['label'].split('/')[-1]

        source_label = f"{source_label}"
        target_label = f"{target_label}"

        G.add_edge(source_label, target_label)

        # add node type
        G.nodes[source_label]['type'] = edge['source']['type']
        G.nodes[target_label]['type'] = edge['target']['type']

        # add node content
        G.nodes[source_label]['content'] = edge['source']['content']
        G.nodes[target_label]['content'] = edge['target']['content']

    return G


if __name__ == '__main__':
    base_path = '/path/to/directory'
    project_name = 'AI-Project'

    project_path = base_path + '/' + project_name
    save_path = base_path + '/' + project_name + '.png'

    folder_structure = generate_folder_structure(project_path)

    # generate raw call graph using the existing package
    json_data = generate_call_graph(project_path)

    # flatten the raw call graph
    node_map = flatten_node_info(json_data)

    # format interal call graph
    call_graph = format_internal_call_graph(node_map, json_data, project_path)

    G = convert_json_to_digraph(call_graph, folder_structure)

    # save graph object to file
    pickle.dump(G, open(base_path + '/' + project_name + '.pickle', 'wb'))