from __future__ import annotations

import logging
from typing import Dict, List, Tuple

from graph.connections.db_connection import DbConnection
from graph.ingestion.create_table import create_edge_table, create_node_table
from graph.ingestion.dataiku.data_types import DataikuDataType

logger = logging.getLogger(__name__)


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

    def with_property(self, property: str, type: DataikuDataType) -> 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, DataikuDataType]]) -> NodeTableBuilder:
        for prop, type in properties:
            self.with_property(prop, type)

        return self

    def with_primary_key(self, primary_key: str, type: DataikuDataType) -> 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.")

        create_node_table(self.__table_name, self.__primary_key, self.__property_to_type, self.__connection)


class EdgeTableBuilder:
    def __init__(self, connection: DbConnection) -> 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, DataikuDataType] = {}

    def with_property(self, property: str, type: DataikuDataType) -> 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, DataikuDataType]]) -> 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

        create_edge_table(self.__table_name, self.__from, self.__to, self.__property_to_type, self.__connection)
