from datetime import datetime
from typing import Iterator, Tuple, override

from sqlalchemy import Column, MetaData, Table, UniqueConstraint

from sandwich import SANDWICH_VERSION
from sandwich.dialects.base import DialectHandler
from sandwich.modeling import modeling_metadata, StgInfo, Dv2SystemInfo

from .base import HubBasedValidator, HubSchemaGenerator, ValidationResult


class Scd2DimValidator(HubBasedValidator):
    def __init__(self, template: str):
        super().__init__(template)
        self._on_validate_staging = self._validate_staging

    @staticmethod
    def _validate_staging(stg_info: StgInfo, _: Dv2SystemInfo) -> None:

        # -----------------
        # BKs
        # -----------------
        # You don't need a hub or/and a dim tables for a non-business entity.
        # So you have to have at least one business key, and you can have more.
        # Naming convention is to just add a `bk_` prefix to the original key name
        # because we want to keep information of the original names
        if len(stg_info.bk_keys) == 0:
            raise Exception("bk column(s) are required for `scd2dim` validation")


        system_column_names = stg_info.sys_columns.keys()

        if modeling_metadata.hashdiff not in system_column_names:
            raise Exception(f"{modeling_metadata.hashdiff} column is required for scd2dim validation")
        if modeling_metadata.is_available not in system_column_names:
            raise Exception(f"{modeling_metadata.is_available} column is required for scd2dim validation")

class Scd2DimSchemaGenerator(HubSchemaGenerator):

    def __init__(self, dialect_handler: DialectHandler, validation_result: ValidationResult, entity_registration_date: datetime):
        super().__init__(dialect_handler, validation_result, entity_registration_date)
        self._on_make_procedures.extend([self._make_stg_proc, self._make_sat_proc, self._make_dim_proc])

    @override
    def make_tables(self) -> dict[str, Table]:
        entity_name = self._validation_result.entity_name
        bk_keys = self._validation_result.bk_keys
        hk_key = self._validation_result.hk_keys[0]
        business_column_types = self._validation_result.business_column_types
        system_column_types = self._validation_result.system_column_types

        # Helper functions for creating columns
        def get_bk_columns() -> Iterator[Column]:
            return (Column(bk_key[0], bk_key[1], nullable=False) for bk_key in bk_keys)

        def get_bk_pk_columns() -> Iterator[Column]:
            return (Column(bk_key[0], bk_key[1], primary_key=True) for bk_key in bk_keys)

        def get_hk_pk_column() -> Column:
            return Column(hk_key[0], hk_key[1], primary_key=True)

        # def get_loaddate_column() -> Column:
        #     _load_date = modeling_metadata.loaddate
        #     _load_date_type = system_column_types[_load_date]
        #     return Column(_load_date, _load_date_type, nullable=False)

        def get_loaddate_pk_column() -> Column:
            _load_date = modeling_metadata.loaddate
            _load_date_type = system_column_types[_load_date]
            return Column(_load_date, _load_date_type, primary_key=True)

        def get_datefrom_pk_column() -> Column:
            _load_date = modeling_metadata.loaddate
            _load_date_type = system_column_types[_load_date]
            return Column("DateFrom", _load_date_type, primary_key=True)

        def get_dateto_column() -> Column:
            _load_date = modeling_metadata.loaddate
            _load_date_type = system_column_types[_load_date]
            return Column("DateTo", _load_date_type, nullable=True)

        def get_recordsource_column() -> Column:
            _record_source = modeling_metadata.recordsource
            _record_source_type = system_column_types[_record_source]
            return Column(_record_source, _record_source_type, nullable=False)

        def get_business_columns() -> Iterator[Column]:
            return (Column(col_name, col_type, nullable=True) for (col_name, col_type) in business_column_types.items())

        def get_is_available_column() -> Column:
            _is_available = modeling_metadata.is_available
            _is_available_type = system_column_types[_is_available]
            return Column(_is_available, _is_available_type, nullable=False)

        def get_hashdiff_column() -> Column:
            _hashdiff = modeling_metadata.hashdiff
            _hashdiff_type = system_column_types[_hashdiff]
            return Column(_hashdiff, _hashdiff_type, nullable=False)

        # Create hub table
        # hub_table = Table(entity_name, MetaData(), schema="hub")
        # for bk_col in get_bk_columns():
        #     hub_table.append_column(bk_col)
        # hub_table.append_column(get_hk_pk_column())
        # hub_table.append_column(get_loaddate_column())
        # hub_table.append_column(get_recordsource_column())
        # hub_table.append_constraint(UniqueConstraint(*[bk[0] for bk in bk_keys]))

        # Create sat table
        sat_table = Table(entity_name, MetaData(), schema="sat")
        for bk_col in get_bk_columns():
            sat_table.append_column(bk_col)
        sat_table.append_column(get_hk_pk_column())
        sat_table.append_column(get_loaddate_pk_column())
        sat_table.append_column(get_recordsource_column())
        sat_table.append_column(get_hashdiff_column())
        for business_col in get_business_columns():
            sat_table.append_column(business_col)
        sat_table.append_column(get_is_available_column())

        # Create dim table
        dim_table = Table(entity_name, MetaData(), schema="dim")
        for bk_col in get_bk_pk_columns():
            dim_table.append_column(bk_col)
        for business_col in get_business_columns():
            dim_table.append_column(business_col)
        dim_table.append_column(get_is_available_column())
        dim_table.append_column(Column("IsCurrent", self.dialect_handler.get_boolean_type(), nullable=False))
        dim_table.append_column(get_datefrom_pk_column())
        dim_table.append_column(get_dateto_column())

        result = super().make_tables() # hub is already there
        result["sat"] = sat_table
        result["dim"] = dim_table
        return result

    def _make_sat_proc(self, procedures: dict[str, Tuple[str, str, str]], tables: dict[str, Table]) -> None:
        sat_proc_code, sat_proc_name, sat_call_stmt = self.dialect_handler.make_scd2_sat_proc(
            sat_table=tables["sat"],
            hk_name=self._validation_result.hk_keys[0][0],
            hashdiff_col=modeling_metadata.hashdiff,
            is_available_col=modeling_metadata.is_available,
            loaddate_col=modeling_metadata.loaddate,
            stg_schema=self._validation_result.stg_schema,
            header=self.header
        )
        procedures["sat"] = (sat_proc_code, sat_proc_name, sat_call_stmt)

    def _make_dim_proc(self, procedures: dict[str, Tuple[str, str, str]], tables: dict[str, Table]) -> None:
        dim_proc_code, dim_proc_name, dim_call_stmt = self.dialect_handler.make_scd2_dim_proc(
            dim_table=tables["dim"],
            bk_keys=self._validation_result.bk_keys,
            header=self.header
        )
        procedures["dim"] = (dim_proc_code, dim_proc_name, dim_call_stmt)

    def _make_stg_proc(self, procedures: dict[str, Tuple[str, str, str]], tables: dict[str, Table]) -> None:
        if self._validation_result.stg_schema == "proxy":
            stg_proc_code, stg_proc_name, stg_call_stmt = self.dialect_handler.make_stg_materialization_proc(
                entity_name=self._validation_result.entity_name,
                header=self.header
            )
            procedures["stg"] = (stg_proc_code, stg_proc_name, stg_call_stmt)

