import logging
import os
import time
from typing import List
from uuid import uuid4

from pandas import DataFrame

from graph.connections.db_connection import DbConnection

logger = logging.getLogger(__name__)


# Export to a CSV first because kuzu does not handle disjointed dataframe for the moment.
# See https://github.com/kuzudb/kuzu/discussions/4568
def batch_insert_nodes(temp_dir: str, connection: DbConnection, table: str, df: DataFrame) -> None:
    start_time = time.time()
    filename = str(uuid4())
    filepath = f"{temp_dir}/nodes_{filename}.csv"
    try:
        df.to_csv(filepath, index=False, sep=",")

        file_write_time = time.time()

        query = f"""
            COPY `{table}`({",".join(f"`{col}`" for col in df.columns)}) FROM "{filepath}" (IGNORE_ERRORS=true, HEADER=true, DELIMITER=",")
        """
        logger.debug(f"Inserting nodes with query {query}.")
        connection.execute(query)

        insert_time = time.time()
        logger.debug(
            f"Batch insertion (node): file write {round(file_write_time - start_time, 5)}, insert {round(insert_time - file_write_time, 5)}"
        )
    finally:
        os.unlink(filepath)


def batch_insert_edges(
    temp_dir: str, connection: DbConnection, table: str, properties: List[str], df: DataFrame
) -> None:
    start_time = time.time()
    filename = str(uuid4())
    filepath = f"{temp_dir}/edges_{filename}.csv"
    try:
        df.to_csv(filepath, index=False, sep=",")

        file_write_time = time.time()

        # This is necessary that dataframe columns and properties are in the same order.
        ordered_properties = [col for col in df.columns if col in properties]
        query = f"""
            COPY `{table}`({",".join(f"`{p}`" for p in ordered_properties)}) FROM "{filepath}" (IGNORE_ERRORS=true, HEADER=true, DELIMITER=",")
        """
        logger.debug(f"Inserting edges with query {query}.")
        connection.execute(query)

        insert_time = time.time()
        logger.debug(
            f"Batch insertion (edge): file write {round(file_write_time - start_time, 5)}, insert {round(insert_time - file_write_time, 5)}"
        )
    finally:
        os.unlink(filepath)
