"""
Minimal Databricks engine shim to mimic SQLAlchemy for metadata + queries.

The native Databricks SQL driver is not a SQLAlchemy dialect, so we wrap it
to provide:
- `connection.execute(...)` that behaves like a Result (returns_rows, keys, fetchall)
- `engine.dialect.get_schema_names/get_table_names/get_columns` for reflection-like use
"""
from dataclasses import dataclass
from typing import Any, Iterable, List, Optional

from databricks import sql
from sqlalchemy.sql import ClauseElement


@dataclass
class _ResultAdapter:
    """Lightweight adapter to mimic SQLAlchemy Result for mnemon usage."""
    cursor: Any
    returns_rows: bool = True

    def keys(self) -> List[str]:
        if not getattr(self.cursor, "description", None):
            return []
        return [col[0] for col in self.cursor.description]

    def fetchall(self) -> Iterable:
        return self.cursor.fetchall()


class _DatabricksDialect:
    """
    Tiny dialect shim providing inspector-style helpers for tables/columns.
    Uses SHOW/DESCRIBE to avoid relying on information_schema availability.
    """

    def get_inspector(self, engine: "DatabricksEngine") -> "DatabricksInspector":
        return DatabricksInspector(engine)

    def get_schema_names(self, connection: "DatabricksConnection") -> List[str]:
        cur = connection.execute("SHOW DATABASES")
        return [row[0] for row in cur.fetchall()]

    def get_table_names(
        self,
        connection: "DatabricksConnection",
        schema: Optional[str] = None,
    ) -> List[str]:
        stmt = f"SHOW TABLES IN `{schema}`" if schema else "SHOW TABLES"
        cur = connection.execute(stmt)
        rows = cur.fetchall()
        # SHOW TABLES columns: database, tableName, isTemporary
        return [row[1] if len(row) > 1 else row[0] for row in rows]

    def get_columns(
        self,
        connection: "DatabricksConnection",
        table_name: str,
        schema: Optional[str] = None,
    ) -> List[dict]:
        qualified = f"`{schema}`.`{table_name}`" if schema else f"`{table_name}`"
        cur = connection.execute(f"DESCRIBE TABLE {qualified}")
        cols: List[dict] = []
        for row in cur.fetchall():
            # DESCRIBE TABLE returns: col_name, data_type, comment, ...
            if not row or len(row) < 2:
                continue
            name, dtype = row[0], row[1]
            if name is None or str(name).startswith("#"):
                continue  # skip footer rows
            cols.append(
                {
                    "name": name,
                    "type": dtype,
                    "nullable": True,  # DESCRIBE does not expose nullability
                    "default": None,
                }
            )
        return cols


class DatabricksInspector:
    """Inspector-style helper to mirror SQLAlchemy's reflection surface."""

    def __init__(self, engine: "DatabricksEngine"):
        self.engine = engine

    def get_schema_names(self) -> List[str]:
        with self.engine.connect() as conn:
            return self.engine.dialect.get_schema_names(conn)

    def get_table_names(self, schema: Optional[str] = None) -> List[str]:
        with self.engine.connect() as conn:
            return self.engine.dialect.get_table_names(conn, schema=schema)

    def get_columns(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
        with self.engine.connect() as conn:
            return self.engine.dialect.get_columns(conn, table_name, schema=schema)


class DatabricksEngine:
    def __init__(self, conn):
        self._conn = conn
        self.dialect = _DatabricksDialect()

    def connect(self):
        return DatabricksConnection(self._conn, self.dialect, engine=self)

    def dispose(self):
        try:
            self._conn.close()
        except Exception:
            pass


class DatabricksConnection:
    def __init__(self, conn, dialect: _DatabricksDialect, engine: Optional["DatabricksEngine"] = None):
        self._conn = conn
        self.dialect = dialect
        self.engine = engine

    def execute(self, statement, parameters=None):
        # Accept either raw SQL string or a compiled ClauseElement
        if isinstance(statement, ClauseElement):
            compiled = statement.compile(compile_kwargs={"literal_binds": True})
            sql_stmt = str(compiled)
        else:
            sql_stmt = str(statement)

        cursor = self._conn.cursor()
        cursor.execute(sql_stmt, parameters or {})
        return _ResultAdapter(cursor)

    # Alias to mirror SQLAlchemy Connection API
    def exec_driver_sql(self, statement: str, parameters=None):
        return self.execute(statement, parameters)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        self.close()

    def close(self):
        pass


def build_databricks_connection(auth_method: str, config: dict):
    common = {
        "server_hostname": config["host"],
        "http_path": config["http_path"],
    }

    if auth_method == "oauth_m2m":
        return sql.connect(
            **common,
            auth_type="oauth-m2m",
            client_id=config["client_id"],
            client_secret=config["client_secret"],
        )

    return sql.connect(
        **common,
        access_token=config["access_token"],
    )
