from cognitive_architecture.infrastructure.databases.graph.get_graph_client import get_graph_client
from cognitive_architecture.shared.data_models import GraphDBType
from networkx import Graph


async def extract_node_descriptions(data):
    descriptions = []

    for node_id, attributes in data:
        if 'description' in attributes and 'unique_id' in attributes:
            descriptions.append({'node_id': attributes['unique_id'], 'description': attributes['description'], 'layer_uuid': attributes['layer_uuid'], 'layer_decomposition_uuid': attributes['layer_decomposition_uuid'] })
    return descriptions


async def add_node_connection(node_descriptions):
    grouped_data = {}

    for item in node_descriptions:
        uuid = item['layer_decomposition_uuid']

        if uuid not in grouped_data:
            grouped_data[uuid] = []

        grouped_data[uuid].append(item)

    return grouped_data


def connect_nodes_in_graph(graph: Graph, relationship_dict: dict) -> Graph:
    """
    For each relationship in relationship_dict, check if both nodes exist in the graph based on node attributes.
    If they do, create a connection (edge) between them.

    :param graph: A NetworkX graph object
    :param relationship_dict: A dictionary containing relationships between nodes
    """
    for id, relationships in relationship_dict.items():
        for relationship in relationships:
            searched_node_attr_id = relationship['searched_node_id']
            print(searched_node_attr_id)
            score_attr_id = relationship['original_id_for_search']
            score = relationship['score']

            # Initialize node keys for both searched_node and score_node
            searched_node_key, score_node_key = None, None

            # Find nodes in the graph that match the searched_node_id and score_id from their attributes
            for node, attrs in graph.nodes(data=True):
                if 'unique_id' in attrs:  # Ensure there is an 'id' attribute
                    if attrs['unique_id'] == searched_node_attr_id:
                        searched_node_key = node
                    elif attrs['unique_id'] == score_attr_id:
                        score_node_key = node

                # If both nodes are found, no need to continue checking other nodes
                if searched_node_key and score_node_key:
                    break

            # Check if both nodes were found in the graph
            if searched_node_key is not None and score_node_key is not None:
                print(searched_node_key)
                print(score_node_key)
                # If both nodes exist, create an edge between them
                # You can customize the edge attributes as needed, here we use 'score' as an attribute
                graph.add_edge(searched_node_key, score_node_key, weight=score,
                               score_metadata=relationship.get('score_metadata'))

    return graph


def graph_ready_output(results):
    relationship_dict = {}

    for result_tuple in results:

        uuid, scored_points_list, desc, node_id = result_tuple
        # Unpack the tuple

        # Ensure there's a list to collect related items for this uuid
        if uuid not in relationship_dict:
            relationship_dict[uuid] = []

        for scored_points in scored_points_list:  # Iterate over the list of ScoredPoint lists
            for scored_point in scored_points:  # Iterate over each ScoredPoint object
                if scored_point.score > 0.9:  # Check the score condition
                    # Append a new dictionary to the list associated with the uuid
                    relationship_dict[uuid].append({
                        'collection_name_uuid': uuid,
                        'searched_node_id': scored_point.id,
                        'score': scored_point.score,
                        'score_metadata': scored_point.payload,
                        'original_id_for_search': node_id,
                    })
    return relationship_dict




if __name__ == '__main__':
    graph_client = get_graph_client(GraphDBType.NETWORKX)
    add_node_connection(graph_client)