import pandas as pd
from neo4j import GraphDatabase
from typing import List

class GraphBuilder(object):

    def __init__(self,
                 uri: str = "bolt://localhost:7687",
                 user: str = "neo4j",
                 password: str = "neo4j1"):
        """
        Args:
            uri: uri to access the Bolt DirectDriver.
            user: neo4j user.
            password: neo4j password.
        """
        self._driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self) -> None:
        self._driver.close()

    def create_nodes(self,
                     df : pd.DataFrame,
                     node_props: List[str],
                     label: str) -> None:
        """Creates nodes with a fixed set properties based on the dataframe provided.

        Args:
            df: Dataframe in which each column and each row corresponds to a different node.
            node_props: names of the properties of each node.
            label: neo4j label for each node created.

        """
        cypher_statement = f"CREATE (new_node: {label})"
        for prop in node_props:
            cypher_statement += f"\n SET new_node.{prop} = ${prop}"

        with self._driver.session() as session:
            props_dict = {}
            for index, row in df.iterrows():
                for prop in node_props:
                    props_dict[prop] = row[prop]
                session.run(
                    cypher_statement,
                    props_dict)

    def create_edges(self,
                     df: pd.DataFrame,
                     predecessor_prop: str,
                     predecessor_label: str,
                     successor_prop: str,
                     successor_label: str,
                     edge_props: List[str],
                     edge_label: str) -> None:
        """Creates edges with a fixed set properties based on the dataframe provided.

        Args:
            df: Dataframe in which each column and each row corresponds to a different node.
            predecessor_prop: name of the matching property in the source node of the edge.
            predecessor_label: label of the source node of the edge.
            sucessor_prop: name of the matching property in the target node of the edge.
            sucessor_label: label of the target node of the edge.
            edge_props: names of the properties of each edge.
            edge_label: neo4j label for each edge created.

        """
        cypher_statement = f"""
            MATCH
                (predecessor_node: {predecessor_label} {{ id: ${predecessor_prop} }}),
                (successor_node: {successor_label} {{ id: ${successor_prop} }})
            CREATE (predecessor_node)-[new_edge: {edge_label}]->(successor_node)"""
        for prop in edge_props:
            cypher_statement += f"\n SET new_edge.{prop} = ${prop}"

        with self._driver.session() as session:
            props_dict = {}
            for index, row in df.iterrows():
                for prop in edge_props + [successor_prop, predecessor_prop]:
                    props_dict[prop] = row[prop]
                session.run(
                    cypher_statement,
                    props_dict)