"""Generic Databend I/O for the Python transform engine.

Reads a relation into a pandas DataFrame and writes a DataFrame back as a CREATE OR
REPLACE table, over the Rust ``databend_driver`` (the same client ``load_seeds.py`` uses,
which handles INSERTs against the TiDB Cloud Lake gateway where the dbt adapter's HTTP
client does not). No product facts live here — callers pass the SQL and the column
coercion sets (config-driven-clean-structure).

Coercion: the driver returns DECIMAL/DATE/etc. as Python ``Decimal``/``date`` and VARCHAR
as ``str`` — including numeric-looking VARCHARs (e.g. seed amounts loaded all-VARCHAR).
``read_df`` auto-numifies columns the driver reports as Decimal/Float and auto-dates Date/
Timestamp columns, then applies the caller's explicit ``numeric``/``date``/``boolean``
overrides (needed for VARCHAR columns that hold numbers, while leaving e.g. JSON-string
columns as strings).
"""

from __future__ import annotations

import datetime as _dt
import math
import os
from decimal import Decimal
from typing import Iterable

import pandas as pd
from databend_driver import BlockingDatabendClient


def dsn() -> str:
    """Build a databend:// DSN with the warehouse param the gateway requires."""
    d = os.environ["DATABEND_DSN"]
    if d.startswith("lake://"):
        d = "databend://" + d[len("lake://") :]
    if "warehouse=" not in d:
        d += ("&" if "?" in d else "?") + f"warehouse={os.environ.get('DATABEND_WAREHOUSE', 'default')}"
    return d


def connect():
    """Return a blocking cursor against the lakehouse warehouse."""
    return BlockingDatabendClient(dsn()).cursor()


def _rows(cur) -> list[list]:
    return [list(r.values()) if hasattr(r, "values") else list(r) for r in cur.fetchall()]


def _to_date(v):
    if v is None or (isinstance(v, float) and math.isnan(v)):
        return None
    if isinstance(v, _dt.datetime):
        return v.date()
    if isinstance(v, _dt.date):
        return v
    ts = pd.to_datetime(v, errors="coerce")
    return None if pd.isna(ts) else ts.date()


def read_df(
    cur,
    sql: str,
    *,
    numeric: Iterable[str] = (),
    date: Iterable[str] = (),
    boolean: Iterable[str] = (),
) -> pd.DataFrame:
    """Run ``sql`` and return a typed DataFrame. ``cur.description`` carries the driver
    types, which we use to auto-coerce; ``numeric``/``date``/``boolean`` force columns the
    driver reports as String (the all-VARCHAR seed amounts, TINYINT flags, etc.)."""
    cur.execute(sql)
    desc = cur.description or []
    cols = [d[0] for d in desc]
    types = {d[0]: str(d[1]) for d in desc}
    df = pd.DataFrame(_rows(cur), columns=cols)

    numeric, date, boolean = set(numeric), set(date), set(boolean)
    for c in cols:
        t = types.get(c, "")
        if c in numeric or "Decimal" in t or "Float" in t:
            df[c] = pd.to_numeric(df[c], errors="coerce")
        elif c in date or t.startswith("Date") or "Timestamp" in t:
            df[c] = df[c].map(_to_date)
        elif c in boolean:
            df[c] = df[c].map(lambda v: None if v is None else bool(int(v)) if not isinstance(v, bool) else v)
    return df


def _ddl_type(series: pd.Series) -> str:
    """Pick a Databend column type from a pandas Series (first non-null value wins)."""
    dt = series.dtype
    if pd.api.types.is_bool_dtype(dt):
        return "BOOLEAN"
    if pd.api.types.is_integer_dtype(dt):
        return "BIGINT"
    if pd.api.types.is_float_dtype(dt):
        return "DOUBLE"
    if pd.api.types.is_datetime64_any_dtype(dt):
        return "TIMESTAMP"
    for v in series:
        if v is None or (isinstance(v, float) and math.isnan(v)) or v is pd.NaT:
            continue
        if isinstance(v, bool):
            return "BOOLEAN"
        if isinstance(v, _dt.datetime):
            return "TIMESTAMP"
        if isinstance(v, _dt.date):
            return "DATE"
        if isinstance(v, (int, Decimal)):
            return "DOUBLE" if isinstance(v, Decimal) else "BIGINT"
        if isinstance(v, float):
            return "DOUBLE"
        break
    return "VARCHAR"


def _literal(v) -> str:
    """Render a Python value as a Databend SQL literal."""
    if v is None or v is pd.NaT:
        return "NULL"
    if hasattr(v, "item") and not isinstance(v, (str, bytes)):  # numpy scalar
        v = v.item()
    if isinstance(v, float):
        return "NULL" if math.isnan(v) else repr(v)
    if isinstance(v, bool):
        return "TRUE" if v else "FALSE"
    if isinstance(v, Decimal):
        return str(v)
    if isinstance(v, int):
        return str(v)
    if isinstance(v, _dt.datetime):
        return "'" + v.strftime("%Y-%m-%d %H:%M:%S") + "'"
    if isinstance(v, _dt.date):
        return "'" + v.isoformat() + "'"
    try:
        if pd.isna(v):
            return "NULL"
    except (TypeError, ValueError):
        pass
    return "'" + str(v).replace("'", "''") + "'"


def replace_table(cur, fq_name: str, df: pd.DataFrame, *, chunk: int = 500) -> int:
    """CREATE OR REPLACE ``fq_name`` from ``df`` (write-truncate semantics) and INSERT its
    rows. Returns the row count written."""
    if df.shape[1] == 0:
        raise ValueError(f"refusing to create {fq_name} from a DataFrame with no columns")
    cols_ddl = ", ".join(f'"{c}" {_ddl_type(df[c])}' for c in df.columns)
    cur.execute(f"CREATE OR REPLACE TABLE {fq_name} ({cols_ddl})")
    records = df.to_dict("records")
    collist = ", ".join(f'"{c}"' for c in df.columns)
    for i in range(0, len(records), chunk):
        block = records[i : i + chunk]
        values = ", ".join("(" + ", ".join(_literal(r[c]) for c in df.columns) + ")" for r in block)
        cur.execute(f"INSERT INTO {fq_name} ({collist}) VALUES {values}")
    return len(records)