import logging
from typing import Dict

from graph.connections.db_connection import DbConnection
from graph.ingestion.dataiku.data_types import DataikuDataType
from graph.ingestion.kuzu.data_types import TYPE_MAPPING_DATAIKU_KUZU
from solutions.graph.models import PropertyTypeNotSupportedError

logger = logging.getLogger(__name__)


def __validate_property_types_supported(
    property_to_type: Dict[str, DataikuDataType],
) -> None:
    for prop, type in property_to_type.items():
        if type not in TYPE_MAPPING_DATAIKU_KUZU:
            raise PropertyTypeNotSupportedError(property_name=prop, dataiku_type=type, db_type="kuzu")


def create_node_table(
    table_name: str, primary_key: str, property_to_type: Dict[str, DataikuDataType], connection: DbConnection
) -> None:
    ordered_properties = sorted(property_to_type.items(), key=lambda item: item[0])
    __validate_property_types_supported(property_to_type)

    logger.debug(
        f"Creating query to create node table %s with PK %s and properties %s.",
        table_name,
        primary_key,
        ordered_properties,
    )
    query = f"""
            CREATE NODE TABLE `{table_name}` (
                {",".join([f"`{prop}` {TYPE_MAPPING_DATAIKU_KUZU[type]}" for prop, type in ordered_properties])},
                PRIMARY KEY (`{primary_key}`)
            )
        """

    logger.debug(f"Creating node table {table_name} with query '{query}'.")

    connection.execute(query).close()

    logger.info(f"Done creating node table {table_name}.")


def create_edge_table(
    table_name: str,
    from_node: str,
    to_node: str,
    property_to_type: Dict[str, DataikuDataType],
    connection: DbConnection,
) -> None:
    ordered_properties = sorted(property_to_type.items(), key=lambda item: item[0])
    __validate_property_types_supported(property_to_type)

    logger.debug(
        f"Creating query to create edge table %s from %s to %s with properties %s.",
        table_name,
        from_node,
        to_node,
        ordered_properties,
    )
    query = f"""
        CREATE REL TABLE `{table_name}` (
            FROM `{from_node}` TO `{to_node}`{"," if ordered_properties else ""}
            {",".join([f"`{prop}` {TYPE_MAPPING_DATAIKU_KUZU[type]}" for prop, type in ordered_properties])}
        )
    """

    logger.debug(f"Creating edge table {table_name} with query '{query}'.")

    connection.execute(query).close()

    logger.info(f"Done creating edge table {table_name}.")
