from __future__ import annotations

import logging
from typing import Dict, List, Tuple

import kuzu

from .data_types import KuzuDataType

logger = logging.getLogger(__name__)


class NodeTableBuilder:
    def __init__(self, connection: kuzu.Connection) -> None:
        self.__connection = connection
        self.__table_name: str | None = None
        self.__primary_key: str | None = None
        self.__property_to_type: Dict[str, KuzuDataType] = {}

    def with_property(self, property: str, type: KuzuDataType) -> NodeTableBuilder:
        if property in self.__property_to_type:
            if type != self.__property_to_type[property]:
                raise Exception(
                    f"""Trying to set two different types for the same property {property}: {self.__property_to_type[property]} to {type}."""
                )
        else:
            self.__property_to_type[property] = type

        return self

    def with_properties(self, properties: List[Tuple[str, KuzuDataType]]) -> NodeTableBuilder:
        for prop, type in properties:
            self.with_property(prop, type)

        return self

    def with_primary_key(self, primary_key: str, type: KuzuDataType) -> NodeTableBuilder:
        self.__primary_key = primary_key
        self.with_property(primary_key, type)
        return self

    def with_table_name(self, name: str) -> NodeTableBuilder:
        self.__table_name = name
        return self

    def create(self) -> None:
        if not self.__table_name:
            raise Exception("A table name is mandatory to create a table.")
        if not self.__primary_key:
            raise Exception("Primary key is mandatory to create a table.")

        ordered_properties = sorted(self.__property_to_type.items(), key=lambda item: item[0])
        query = f"""
            CREATE NODE TABLE `{self.__table_name}` (
                {",".join([f"`{prop}` {type}" for prop, type in ordered_properties])},
                PRIMARY KEY (`{self.__primary_key}`)
            )
        """

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

        self.__connection.execute(query)

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


class EdgeTableBuilder:
    def __init__(self, connection: kuzu.Connection) -> None:
        self.__connection = connection
        self.__table_name: str | None = None
        self.__from: str | None = None
        self.__to: str | None = None
        self.__property_to_type: Dict[str, KuzuDataType] = {}

    def with_property(self, property: str, type: KuzuDataType) -> EdgeTableBuilder:
        if property in self.__property_to_type:
            if type != self.__property_to_type[property]:
                raise Exception(
                    f"""Trying to set two different types for the same property {property}: {self.__property_to_type[property]} to {type}."""
                )
        else:
            self.__property_to_type[property] = type

        return self

    def with_properties(self, properties: List[Tuple[str, KuzuDataType]]) -> EdgeTableBuilder:
        for prop, type in properties:
            self.with_property(prop, type)

        return self

    def from_node(self, from_node: str) -> EdgeTableBuilder:
        self.__from = from_node
        return self

    def to_node(self, to_node: str) -> EdgeTableBuilder:
        self.__to = to_node
        return self

    def with_table_name(self, name: str) -> EdgeTableBuilder:
        self.__table_name = name
        return self

    def create(self) -> None:
        if not self.__table_name:
            raise Exception("A table name is mandatory to create a table.")

        assert self.__from
        assert self.__to

        ordered_properties = sorted(self.__property_to_type.items(), key=lambda item: item[0])
        query = f"""
            CREATE REL TABLE `{self.__table_name}` (
                FROM `{self.__from}` TO `{self.__to}`{"," if ordered_properties else ""}
                {",".join([f"`{prop}` {type}" for prop, type in ordered_properties])}
            )
        """

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

        self.__connection.execute(query)

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