import logging
import time
from typing import Any

from pandas import DataFrame

from graph.connections.db_connection import DbConnection

logger = logging.getLogger(__name__)


def _edge_rows_without_endpoint_properties(df: DataFrame, source_col: str, target_col: str) -> list[dict[str, Any]]:
    endpoint_columns = {source_col, target_col}
    rows = []
    for row in df.to_dict("records"):
        rows.append(
            {
                "source": row[source_col],
                "target": row[target_col],
                "properties": {key: value for key, value in row.items() if key not in endpoint_columns},
            }
        )
    return rows


def batch_insert_nodes(connection: DbConnection, table: str, primary_key_prop_name: str, df: DataFrame) -> None:
    start_time = time.time()

    rows = df.to_dict("records")

    # Cypher query
    # UNWIND expands the list of rows so we can process them one by one in Cypher
    # WHERE filters out rows with null/NaN primary keys (Neo4j rejects these in MERGE)
    #   NaN check: `x <> x` is true only for NaN (IEEE 754 spec: NaN != NaN)
    #   See: https://en.wikipedia.org/wiki/NaN#Comparison_with_NaN
    # MERGE ensures we don't create duplicate nodes if the ID exists
    # SET n += row adds/updates all other properties from the dataframe
    query = f"""
    UNWIND $rows AS row
    WITH row WHERE row.`{primary_key_prop_name}` IS NOT NULL AND NOT (row.`{primary_key_prop_name}` <> row.`{primary_key_prop_name}`)
    MERGE (n:`{table}` {{ `{primary_key_prop_name}`: row.`{primary_key_prop_name}` }})
    SET n += row
    """

    logger.debug("Inserting %s nodes with label %s.", len(rows), table)

    connection.execute(query, parameters={"rows": rows})

    insert_time = time.time()
    logger.debug("Batch insertion (node): processed %s rows in %s s", len(rows), round(insert_time - start_time, 5))


def batch_insert_edges(
    connection: DbConnection,
    rel_type: str,
    source_label: str,
    target_label: str,
    source_col: str,
    target_col: str,
    source_pk_prop: str,
    target_pk_prop: str,
    df: DataFrame,
) -> None:
    start_time = time.time()

    rows = _edge_rows_without_endpoint_properties(df, source_col, target_col)

    # Cypher query logic:
    # 1. Filter out rows with null/NaN source or target keys.
    # 2. MATCH the source node by its primary key property.
    # 3. MATCH the target node by its primary key property.
    # 4. If either MATCH fails (returns null), the query stops for that specific row.
    # 5. CREATE the relationship (not MERGE, to preserve duplicate edges like Kuzu).

    query = f"""
    UNWIND $rows AS row
    WITH row WHERE row.source IS NOT NULL AND NOT (row.source <> row.source)
                   AND row.target IS NOT NULL AND NOT (row.target <> row.target)
    MATCH (source:`{source_label}` {{ `{source_pk_prop}`: row.source }})
    MATCH (target:`{target_label}` {{ `{target_pk_prop}`: row.target }})
    CREATE (source)-[r:`{rel_type}`]->(target)
    SET r += row.properties
    """

    logger.debug("Inserting %s edges of type %s.", len(rows), rel_type)

    connection.execute(query, parameters={"rows": rows})

    insert_time = time.time()
    logger.debug("Batch insertion (edge): processed %s rows in %s s", len(rows), round(insert_time - start_time, 5))
