from typing import List

from pandas import DataFrame

from graph.connections.db_connection import DbConnection


def edge_properties_without_endpoints(properties: List[str], source_col: str, target_col: str) -> List[str]:
    endpoint_columns = {source_col, target_col}
    return [prop for prop in properties if prop not in endpoint_columns]


def batch_insert_nodes(
    temp_dir: str, connection: DbConnection, table: str, primary_key_prop_name: str, df: DataFrame
) -> None:
    from graph.connections.kuzu.db_connection import KuzuConnection
    from graph.connections.neo4j.db_connection import Neo4jConnection
    from graph.ingestion.kuzu.batch_insert import (
        batch_insert_nodes as batch_insert_nodes_kuzu,
    )
    from graph.ingestion.neo4j.batch_insert import (
        batch_insert_nodes as batch_insert_nodes_neo4j,
    )

    if isinstance(connection, KuzuConnection):
        batch_insert_nodes_kuzu(temp_dir, connection, table, df)
    elif isinstance(connection, Neo4jConnection):
        batch_insert_nodes_neo4j(connection, table, primary_key_prop_name, df)
    else:
        raise TypeError(f"Unsupported connection type: {type(connection)}")


def batch_insert_edges(
    temp_dir: str,
    connection: DbConnection,
    table: str,
    properties: List[str],
    source_table: str,
    target_table: str,
    source_col: str,
    target_col: str,
    source_pk_prop: str,
    target_pk_prop: str,
    df: DataFrame,
) -> None:
    from graph.connections.kuzu.db_connection import KuzuConnection
    from graph.connections.neo4j.db_connection import Neo4jConnection
    from graph.ingestion.kuzu.batch_insert import (
        batch_insert_edges as batch_insert_edges_kuzu,
    )
    from graph.ingestion.neo4j.batch_insert import (
        batch_insert_edges as batch_insert_edges_neo4j,
    )

    if isinstance(connection, KuzuConnection):
        kuzu_properties = edge_properties_without_endpoints(properties, source_col, target_col)
        batch_insert_edges_kuzu(temp_dir, connection, table, kuzu_properties, df)
    elif isinstance(connection, Neo4jConnection):
        batch_insert_edges_neo4j(
            connection, table, source_table, target_table, source_col, target_col, source_pk_prop, target_pk_prop, df
        )
    else:
        raise TypeError(f"Unsupported connection type: {type(connection)}")
