"""
SQLAlchemy dialect registrations and patches.
Ensures that custom dialects like ClickHouse are available globally.
"""

import logging

logger = logging.getLogger("sibi_flux.dialects")

def register_dialects():
    """Registers custom dialects and applies necessary monkey patches."""
    
    # 1. ClickHouse Dialect Registration
    try:
        import clickhouse_connect
        from clickhouse_connect.cc_sqlalchemy.dialect import ClickHouseDialect
        from clickhouse_connect.driver.binding import quote_identifier
        from sqlalchemy import text
        from sqlalchemy.dialects import registry

        from clickhouse_connect.cc_sqlalchemy.datatypes.base import sqla_type_from_name

        # Monkey patch for SQLA 2.0 compatibility in Clickhouse-Connect
        def patched_get_table_names(self, connection, schema=None, **kw):
            cmd = "SHOW TABLES"
            if schema:
                cmd += " FROM " + quote_identifier(schema)
            return [row.name for row in connection.execute(text(cmd))]

        def patched_get_columns(self, connection, table_name, schema=None, **kw):
            """Robust column reflection via DESCRIBE TABLE."""
            full_table = quote_identifier(table_name)
            if schema:
                full_table = f"{quote_identifier(schema)}.{full_table}"
            
            # Use DESCRIBE TABLE as it's more standard in ClickHouse for metadata
            result = connection.execute(text(f"DESCRIBE TABLE {full_table}"))
            
            columns = []
            for row in result:
                # row is a Row object, accessing by name is safer if the driver supports it
                # ClickHouse DESCRIBE columns: name, type, default_type, default_expression, comment, codec_expression, ttl_expression
                name = getattr(row, 'name', row[0])
                type_str = getattr(row, 'type', row[1])
                default_expr = getattr(row, 'default_expression', row[3] if len(row) > 3 else None)
                comment = getattr(row, 'comment', row[4] if len(row) > 4 else None)

                sqla_type = sqla_type_from_name(type_str.replace('\n', ''))
                
                columns.append({
                    "name": name,
                    "type": sqla_type,
                    "nullable": getattr(sqla_type, 'nullable', 'Nullable' in type_str),
                    "default": default_expr if default_expr else None,
                    "autoincrement": False,
                    "comment": comment if comment else None
                })
            return columns

        def patched_get_pk_constraint(self, connection, table_name, schema=None, **kw):
            """ClickHouse doesn't have traditional PKs in the SQLA sense for reflection."""
            return {"constrained_columns": [], "name": None}

        ClickHouseDialect.get_table_names = patched_get_table_names
        ClickHouseDialect.get_columns = patched_get_columns
        ClickHouseDialect.get_pk_constraint = patched_get_pk_constraint

        # 2. Patch ChInspector helper get_engine to handle schema=None
        import clickhouse_connect.cc_sqlalchemy.inspector as ch_inspector
        from sqlalchemy.orm.exc import NoResultFound

        def patched_get_engine(connection, table_name, schema=None):
            """Ensures schema is resolved for ClickHouse engine lookup."""
            if schema is None:
                try:
                    # In SQLA 2.0 connection.execute(text(...)) is the standard
                    schema = connection.execute(text("SELECT database()")).scalar()
                except Exception:
                    # Fallback to dialect default if query fails
                    schema = getattr(connection.dialect, 'default_schema_name', 'default')
            
            # Clickhouse-Connect's original get_engine query
            result_set = connection.execute(text(
                f"SELECT engine_full FROM system.tables WHERE database = '{schema}' and name = '{table_name}'"))
            row = next(result_set, None)
            if not row:
                raise NoResultFound(f'Table {schema}.{table_name} does not exist')
            
            from clickhouse_connect.cc_sqlalchemy.ddl.tableengine import build_engine
            return build_engine(row.engine_full)

        ch_inspector.get_engine = patched_get_engine

        # Register Dialects
        registry.register(
            "clickhouse.connect",
            "clickhouse_connect.cc_sqlalchemy.dialect",
            "ClickHouseDialect",
        )
        registry.register(
            "clickhouse", 
            "clickhouse_connect.cc_sqlalchemy.dialect", 
            "ClickHouseDialect"
        )
        
        logger.debug("Successfully registered ClickHouse dialects.")

    except ImportError:
        logger.debug("clickhouse-connect not found; skipping ClickHouse dialect registration.")
    except Exception as e:
        logger.warning(f"Failed to register ClickHouse dialects: {e}")

# Trigger registration on import
register_dialects()
