"""
Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
"""

import json
import logging

from graph_notebook.network.EventfulNetwork import EventfulNetwork, DEFAULT_GRP
from networkx import MultiDiGraph

logging.basicConfig()
logger = logging.getLogger(__file__)

DEFAULT_LABEL_MAX_LENGTH = 10
ENTITY_KEY = "~entityType"
ID_KEY = "~id"
START_KEY = "~start"
END_KEY = "~end"
PROPERTIES_KEY = "~properties"
VERTEX_TYPE_KEY = "~entityType"
EDGE_TYPE_KEY = "~type"
LABEL_KEY = "~labels"
NODE_ENTITY_TYPE = 'node'
REL_ENTITY_TYPE = 'relationship'


class OCNetwork(EventfulNetwork):
    """OCNetwork extends the EventfulNetwork class and uses the add_results method to parse any response that returns
    nodes/relationships as part (or all) of the response.  Currently this only works with HTTPS response format but in
    the future, we will work to support Bolt based responses as well.
    """

    def __init__(self, graph: MultiDiGraph = None, callbacks=None, label_max_length=DEFAULT_LABEL_MAX_LENGTH,
                 edge_label_max_length=DEFAULT_LABEL_MAX_LENGTH, group_by_property=LABEL_KEY,
                 display_property=LABEL_KEY, edge_display_property=EDGE_TYPE_KEY,
                 tooltip_property=None, edge_tooltip_property=None,
                 ignore_groups=False):
        if graph is None:
            graph = MultiDiGraph()
        super().__init__(graph, callbacks, label_max_length, edge_label_max_length, group_by_property,
                         display_property, edge_display_property, tooltip_property, edge_tooltip_property,
                         ignore_groups)

    def get_node_property_value(self, node: dict, props: dict, title, custom_property):
        try:
            if isinstance(custom_property, dict):
                if isinstance(custom_property[title], tuple):
                    if custom_property[title][0] in props and \
                            isinstance(props[custom_property[title][0]], list) and \
                            len(props[custom_property[title][0]]) >= 2:
                        label = props[custom_property[title][0]][custom_property[title][1]]
                    else:
                        label = title
                elif custom_property[title] in props:
                    label = props[custom_property[title]]
                elif LABEL_KEY in props:
                    label = props[LABEL_KEY]
                else:
                    label = props
            elif isinstance(custom_property, tuple) and custom_property[0] in props:
                if isinstance(props[custom_property[0]], list) and len(props[custom_property[0]]) >= 2:
                    label = props[custom_property[0]][custom_property[1]]
                else:
                    label = title
            elif custom_property in [ID_KEY, 'id']:
                label = node[ID_KEY]
            elif custom_property in [LABEL_KEY, 'label']:
                label = node[LABEL_KEY]
            elif custom_property in [VERTEX_TYPE_KEY, 'type']:
                label = node[VERTEX_TYPE_KEY]
            elif custom_property in props:
                label = props[custom_property]
            else:
                label = title
        except (KeyError, IndexError) as e:
            logger.debug(e)
            label = title

        return label

    def get_edge_property_value(self, data: dict, rel: dict, custom_property):
        if custom_property is not EDGE_TYPE_KEY:
            try:
                if isinstance(custom_property, dict):
                    if isinstance(custom_property[data['label']], tuple) and \
                            custom_property[data['label']][0] in data['properties']:
                        if isinstance(data['properties'][custom_property[data['label']][0]], list) and \
                                len(data['properties'][custom_property[data['label']][0]]) >= 2:
                            display_label = str(data['properties'][custom_property[data['label']][0]]
                                                [custom_property[data['label']][1]])
                        else:
                            display_label = rel[EDGE_TYPE_KEY]
                    else:
                        display_label = data['properties'][custom_property[rel[EDGE_TYPE_KEY]]]
                elif isinstance(custom_property, tuple) and \
                        custom_property[0] in data['properties']:
                    display_label = str(data['properties'][custom_property[0]]
                                        [custom_property[1]])
                else:
                    display_label = data['properties'][custom_property]
            except (KeyError, IndexError, TypeError) as e:
                logger.debug(e)
                display_label = rel[EDGE_TYPE_KEY]
        else:
            display_label = rel[EDGE_TYPE_KEY]

        return display_label

    def parse_node(self, node: dict):
        """This parses the node parameter and adds the node to the network diagram

        Args:
            node (dict): The node dictionary to parse
        """
        # generate placeholder tooltip from label; if not present, amalgamate node property values instead
        if LABEL_KEY in node.keys():
            title_plc = node[LABEL_KEY][0]
        else:
            title_plc = ""
            for key in node:
                title_plc += str(node[key])

        if not isinstance(self.group_by_property, dict):  # Handle string format group_by
            try:
                if self.group_by_property in [LABEL_KEY, 'labels'] and len(node[LABEL_KEY]) > 0:
                    group = node[LABEL_KEY][0]
                elif self.group_by_property in [ID_KEY, 'id']:
                    group = node[ID_KEY]
                elif self.group_by_property in node[PROPERTIES_KEY]:
                    group = node[PROPERTIES_KEY][self.group_by_property]
                else:
                    group = DEFAULT_GRP
            except KeyError:
                group = DEFAULT_GRP
        else:  # handle dict format group_by
            try:
                if str(node[LABEL_KEY][0]) in self.group_by_property and len(node[LABEL_KEY]) > 0:
                    key = node[LABEL_KEY][0]
                    if self.group_by_property[key]['groupby'] in [LABEL_KEY, 'labels']:
                        group = node[LABEL_KEY][0]
                    else:
                        group = node[PROPERTIES_KEY][self.group_by_property[key]['groupby']]
                elif ID_KEY in self.group_by_property:
                    group = node[ID_KEY]
                else:
                    group = DEFAULT_GRP
            except KeyError:
                group = DEFAULT_GRP

        props = self.flatten(node)
        label = self.get_node_property_value(node, props, title_plc, self.display_property)
        title, label = self.strip_and_truncate_label_and_title(label, self.label_max_length)
        if self.tooltip_property and self.tooltip_property != self.display_property:
            title, label_plc = self.strip_and_truncate_label_and_title(
                self.get_node_property_value(node, props, title_plc, self.tooltip_property))
        data = {'properties': props, 'label': label, 'title': title, 'group': group}
        if self.ignore_groups:
            data['group'] = DEFAULT_GRP
        self.add_node(node[ID_KEY], data)
    
    def parse_rel(self, rel):
        data = {'properties': self.flatten(rel), 'label': rel[EDGE_TYPE_KEY], 'title': rel[EDGE_TYPE_KEY]}
        display_label = self.get_edge_property_value(data, rel, self.edge_display_property)
        edge_title, edge_label = self.strip_and_truncate_label_and_title(display_label, self.edge_label_max_length)
        if self.edge_tooltip_property and self.edge_tooltip_property != self.edge_display_property:
            edge_title = self.get_edge_property_value(data, rel, self.edge_tooltip_property)
        data['title'] = edge_title
        data['label'] = edge_label
        self.add_edge(from_id=rel[START_KEY], to_id=rel[END_KEY], edge_id=rel[ID_KEY], label=edge_label,
                      title=edge_title, data=data)

    def process_result(self, res: dict):
        """Determines the type of element passed in and processes it appropriately

        Args:
            res (dict): The dictionary to parse
        """
        if ENTITY_KEY in res:
            if res[ENTITY_KEY] == NODE_ENTITY_TYPE:
                self.parse_node(res)
            else:
                self.parse_rel(res)

    def add_results(self, results):
        """Adds the results parameter to the network

        Args:
            results (Object): Determines the type of the object and processes it appropriately
        """
        for res in results["results"]:
            if type(res) is dict:
                for k in res.keys():
                    if type(res[k]) is dict:
                        self.process_result(res[k]) 
                    elif type(res[k]) is list:
                        for res_sublist in res[k]:
                            try:
                                self.process_result(res_sublist)
                            except TypeError as e:
                                logger.debug(f'Property {res_sublist} in list results set is invalid, skipping')
                                logger.debug(f'Error: {e}')
                                continue
