import datetime
import json
from lxml import etree
import networkx as nx
from mirmod import _table_to_object_lookup, miranda
import traceback


def import_graph(ko, graphml_file):
    G = nx.parse_graphml(graphml_file)
    nodes = {}
    wobs = {}
    for node in G.nodes(data=True):
        if "class" in node[1]:
            if node[1]["class"].upper() not in _table_to_object_lookup:
                raise Exception(f"Unknown class {node[1]['class'].upper()}")
            if node[1]["class"].upper() != "KNOWLEDGE_OBJECT":
                nodes[node[0]] = node[1]
        else:
            raise Exception(f'No node class type found for node "{node[0]}"')

    for node in nodes:
        print(f"Creating {node}")
        node_data = nodes[node]
        node_class = node_data["class"].upper()
        print(f"  class: {node_class}")
        try:
            if node_class == "CODE" or node_class == "CODE_BLOCK":
                wob = miranda.create_wob(
                    ko,
                    name=node_data["name"] if "name" in node_data else "Unnamed Code",
                    description=node_data["description"]
                    if "description" in node_data
                    else "",
                )
                if wob.id == -1:
                    raise Exception("Failed to create code block")
                for attr in node_data:
                    if attr in ["name", "description", "id", "metadata_id"]:
                        continue
                    print("  Setting attribute {} to {}".format(attr, node_data[attr]))
                    if attr in wob.orm.keys():
                        try:
                            if attr in wob.default_value.keys():
                                if isinstance(wob.default_value[attr], str):
                                    setattr(wob, attr, node_data[attr])
                                elif isinstance(wob.default_value[attr], int):
                                    setattr(wob, attr, int(node_data[attr]))
                                elif isinstance(wob.default_value[attr], float):
                                    setattr(wob, attr, float(node_data[attr]))
                            else:
                                setattr(wob, attr, node_data[attr])
                        except Exception as e:
                            print(f"  Error setting attribute {attr} : {e}")
                    else:
                        wob.set_attribute(attr, node_data[attr])
                wobs[node] = wob
                wob.update(ko.sctx)
        except Exception as e:
            print(f"  Error: {e}")
            print(traceback.format_exc())

    for edge in G.edges(data=True):
        print(f"Creating edge {edge[0]} -> {edge[1]}, data= {edge[2]}")
        try:
            src = wobs[edge[0]]
            dst = wobs[edge[1]]
            print(
                f"Translated to wob mids: edge {src.metadata_id} -> {dst.metadata_id}"
            )
            attr: dict = json.loads(edge[2]["attributes"])
            for dest_socket in attr.keys():
                src_socket = attr[dest_socket]
                print(f"  Linking {src_socket} to {dest_socket}")
                try:
                    if miranda.link(
                        ko.sctx,
                        src,
                        dst,
                        src_socket[0],
                        dest_socket,
                        datatype=src_socket[1],
                    ):
                        print("Error while linking")
                except Exception as e:
                    print(f"  Error: {e}")
                    print(traceback.format_exc())

        except Exception as e:
            print(f"  Error: {e}")
            print(traceback.format_exc())


def export_graph(graph):
    root = etree.Element("graphml")
    root.set("xmlns", "http://graphml.graphdrawing.org/xmlns")
    root.set("__xmlnsxsi__", "http://www.w3.org/2001/XMLSchema-instance")
    root.set("__xmlnsy__", "http://www.yworks.com/xml/graphml")
    root.set(
        "__xsischemaLocation__",
        "http://graphml.graphdrawing.org/xmlns http://www.yworks.com/xml/schema/graphml/1.1/ygraphml.xsd",
    )
    root.append(etree.Comment(f"Generated by Miranda on {datetime.datetime.now()}"))
    gr = etree.SubElement(root, "graph")
    gr.set("edgedefault", "directed")
    key_map = []
    node_els = []
    for id, node in graph.nodes(data=True):
        if "wob" not in node:
            continue
        node_el = etree.SubElement(gr, "node")
        node_el.set("id", f"n{id}")
        node_el.set("attr.name", node["wob"]["name"])
        for key in node["wob"]:
            key_i = key_map.index(key) if key in key_map else None
            if key_i is None:
                key_i = len(key_map)
                key_map.append(key)
                key_el = etree.SubElement(root, "key")
                key_el.set("id", f"d{key_i}")
                key_el.set("for", "node")
                key_el.set("attr.name", key)
                if isinstance(node["wob"][key], int):
                    key_el.set("attr.type", "int")
                elif isinstance(node["wob"][key], float):
                    key_el.set("attr.type", "double")
                else:
                    key_el.set("attr.type", "string")
            data_el = etree.SubElement(node_el, "data")
            data_el.set("key", f"d{key_i}")
            data_el.text = str(node["wob"][key]).replace("\r\n", "\n")
        node_els.append(node_el)
        gr.append(node_el)

    key_map.append("nodegraphics")
    ngkeyi = len(key_map) - 1
    ngkey_el = etree.SubElement(root, "key")
    ngkey_el.set("id", f"d{ngkeyi}")
    ngkey_el.set("for", "node")
    ngkey_el.set("yfiles.type", "nodegraphics")

    for node in node_els:
        data_el = etree.SubElement(node, "data")
        data_el.set("key", f"d{ngkeyi}")
        shape_el = etree.SubElement(data_el, "__yShapeNode__")
        label_el = etree.SubElement(shape_el, "__yNodeLabel__")
        label_el.text = node.get("attr.name")
        geometry_el = etree.SubElement(shape_el, "__yGeometry__")
        geometry_el.set("height", "30.0")
        geometry_el.set("width", f"{len(label_el.text) * 6.0 + 10.0}")

    id = 0
    for src, dest, edge in graph.edges(data=True):
        edge_el = etree.SubElement(gr, "edge")
        edge_el.set("source", f"n{src}")
        edge_el.set("target", f"n{dest}")
        data = etree.SubElement(edge_el, "data")
        data.set("key", "e{}".format(id))
        data.text = json.dumps(edge["attributes"])
        key = etree.SubElement(root, "key")
        key.set("id", f"e{id}")
        key.set("for", "edge")
        key.set("attr.name", "attributes")
        key.set("attr.type", "string")
        id += 1
        gr.append(edge_el)
        root.append(key)

    treestr = etree.tostring(root, pretty_print=True)
    # stupid hack to get past the lxml limitation on not using : in tag and attribute names
    treestr = treestr.replace(b"__yShapeNode__", b"y:ShapeNode")
    treestr = treestr.replace(b"__yNodeLabel__", b"y:NodeLabel")
    treestr = treestr.replace(b"__yGeometry__", b"y:Geometry")
    treestr = treestr.replace(b"__xmlnsxsi__", b"xmlns:xsi")
    treestr = treestr.replace(b"__xmlnsy__", b"xmlns:y")
    treestr = treestr.replace(b"__xsischemaLocation__", b"xsi:schemaLocation")
    return treestr
