"""PostgreSQL execution operations used by the direct DWH adapter."""

from sqlalchemy import Engine

from tht.execute import ExecResult, PlanSummary, explain as _explain, run_controlled

DEFAULT_LIMIT = 10
DEFAULT_TIMEOUT_MS = 30_000


def run_query(engine: Engine, sql: str, *, limit: int | None = None) -> ExecResult:
    return run_controlled(
        engine,
        sql,
        limit=DEFAULT_LIMIT if limit is None else limit,
        timeout_ms=DEFAULT_TIMEOUT_MS,
    )


def explain(engine: Engine, sql: str) -> PlanSummary:
    return _explain(engine, sql, timeout_ms=DEFAULT_TIMEOUT_MS)


def sample_column(
    engine: Engine, schema: str, table: str, column: str, *, limit: int
) -> list[object]:
    quote = engine.dialect.identifier_preparer.quote
    sql = (
        f"SELECT {quote(column)} FROM {quote(schema)}.{quote(table)} "
        f"WHERE {quote(column)} IS NOT NULL"
    )
    return [row[0] for row in run_query(engine, sql, limit=limit).rows]


def distinct_values(engine: Engine, schema: str, table: str, column: str) -> list[object]:
    quote = engine.dialect.identifier_preparer.quote
    sql = (
        f"SELECT DISTINCT {quote(column)} FROM {quote(schema)}.{quote(table)} "
        f"WHERE {quote(column)} IS NOT NULL ORDER BY {quote(column)}"
    )
    result = run_controlled(engine, sql, limit=1000, timeout_ms=DEFAULT_TIMEOUT_MS)
    return [row[0] for row in result.rows]