"""Unified warehouse adapter with query-tag cost-attribution support.

This module provides a single ``query(sql, *, tag, ...)`` entry-point that
works against both local DuckDB and cloud warehouses (Snowflake / BigQuery).

Cost attribution
----------------
Every query can carry an optional ``tag`` string that is attached to the
query in the warehouse's audit / history metadata.  Convention:

    "team=<name>;job=<name>;run_id=<uuid>"

Warehouse behaviour per backend:

    Snowflake  — executes ``ALTER SESSION SET QUERY_TAG = '<tag>'`` before
                 running the SQL.  The tag appears in QUERY_HISTORY.QUERY_TAG
                 and can be grouped to produce per-team cost reports.

    BigQuery   — (not yet implemented) would use ``job.labels`` on the
                 QueryJobConfig.

    DuckDB     — no native query-cost metadata.  The adapter records
                 ``(tag, elapsed_s, ts)`` in a local ``_query_costs`` table
                 inside the same DuckDB connection so the cost-attribution
                 fallback in ``governance/cost_attribution.py`` can read it.
"""

from __future__ import annotations

import os
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

try:
    import duckdb
    import pandas as pd
    _DUCKDB_AVAILABLE = True
except ImportError:
    _DUCKDB_AVAILABLE = False


# ---------------------------------------------------------------------------
# Backend detection
# ---------------------------------------------------------------------------

def _backend() -> str:
    """Return ``"snowflake"``, ``"bigquery"``, or ``"duckdb"``."""
    if os.environ.get("SNOWFLAKE_ACCOUNT"):
        return "snowflake"
    if os.environ.get("BIGQUERY_PROJECT"):
        return "bigquery"
    return "duckdb"


# ---------------------------------------------------------------------------
# DuckDB connection + _query_costs table bootstrap
# ---------------------------------------------------------------------------

_LOCAL_DB_PATH = (Path(__file__).resolve().parent.parent
                  / "warehouse" / "warehouse.duckdb").as_posix()

_duckdb_con: Any = None  # module-level cached connection


def _get_duckdb_con() -> Any:
    global _duckdb_con
    if _duckdb_con is None:
        if not _DUCKDB_AVAILABLE:
            raise RuntimeError(
                "duckdb is not installed; pip install duckdb pandas"
            )
        _duckdb_con = duckdb.connect(_LOCAL_DB_PATH)
        _duckdb_con.execute("""
            CREATE TABLE IF NOT EXISTS _query_costs (
                ts        TIMESTAMPTZ NOT NULL,
                tag       TEXT,
                elapsed_s DOUBLE NOT NULL
            )
        """)
    return _duckdb_con


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def query(sql: str, *, tag: str | None = None, **params: Any) -> "pd.DataFrame":
    """Run *sql* against the configured backend and return a DataFrame.

    Args:
        sql:    SQL query to execute.  Named parameters (``$name`` or
                ``:name``) are substituted from *params*.
        tag:    Arbitrary string attached to the query for cost attribution.
                Convention: ``"team=<name>;job=<name>;run_id=<uuid>"``.
                On Snowflake the tag is propagated via
                ``ALTER SESSION SET QUERY_TAG``.
                On DuckDB it is recorded in the ``_query_costs`` local table
                (wall-clock seconds; no credit-cost available on DuckDB).
        **params: Bind-parameter values for parameterised queries.

    Returns:
        A ``pandas.DataFrame`` with the query results.

    Raises:
        RuntimeError: If a required cloud SDK is not installed or env vars
                      are missing for the selected backend.
    """
    backend = _backend()

    if backend == "snowflake":
        return _query_snowflake(sql, tag=tag, **params)
    if backend == "bigquery":
        return _query_bigquery(sql, tag=tag, **params)
    return _query_duckdb(sql, tag=tag, **params)


# ---------------------------------------------------------------------------
# Backend implementations
# ---------------------------------------------------------------------------

def _query_snowflake(sql: str, *, tag: str | None = None,
                     **params: Any) -> "pd.DataFrame":
    """Execute *sql* on Snowflake, optionally setting QUERY_TAG first."""
    try:
        import snowflake.connector  # type: ignore[import]
    except ImportError as exc:
        raise RuntimeError(
            "snowflake-connector-python is not installed; "
            "pip install 'snowflake-connector-python[pandas]'"
        ) from exc

    conn = snowflake.connector.connect(
        account=os.environ["SNOWFLAKE_ACCOUNT"],
        user=os.environ["SNOWFLAKE_USER"],
        password=os.environ.get("SNOWFLAKE_PASSWORD", ""),
        database=os.environ.get("SNOWFLAKE_DATABASE", ""),
        schema=os.environ.get("SNOWFLAKE_SCHEMA", "PUBLIC"),
        warehouse=os.environ.get("SNOWFLAKE_WAREHOUSE", ""),
    )
    try:
        cur = conn.cursor()
        if tag:
            # Attach tag so it appears in QUERY_HISTORY.QUERY_TAG
            safe_tag = tag.replace("'", "''")
            cur.execute(f"ALTER SESSION SET QUERY_TAG = '{safe_tag}'")
        cur.execute(sql, params or None)
        rows = cur.fetchall()
        cols = [d[0] for d in cur.description]
        return pd.DataFrame(rows, columns=cols)
    finally:
        conn.close()


def _query_bigquery(sql: str, *, tag: str | None = None,
                    **params: Any) -> "pd.DataFrame":
    """Execute *sql* on BigQuery.

    Note: BigQuery label-based cost attribution (INFORMATION_SCHEMA.JOBS.labels)
    is not yet implemented; this stub runs the query without a label.
    """
    try:
        from google.cloud import bigquery  # type: ignore[import]
    except ImportError as exc:
        raise RuntimeError(
            "google-cloud-bigquery is not installed; "
            "pip install google-cloud-bigquery"
        ) from exc

    project = os.environ["BIGQUERY_PROJECT"]
    client = bigquery.Client(project=project)
    # TODO(L11): pass job_config.labels = parse_tag(tag) when tag is set
    result = client.query(sql).result()
    return result.to_dataframe()


def _query_duckdb(sql: str, *, tag: str | None = None,
                  **params: Any) -> "pd.DataFrame":
    """Execute *sql* locally via DuckDB and record cost in ``_query_costs``."""
    con = _get_duckdb_con()
    t0 = time.monotonic()
    rel = con.execute(sql, list(params.values()) if params else [])
    elapsed = time.monotonic() - t0

    if tag:
        con.execute(
            "INSERT INTO _query_costs (ts, tag, elapsed_s) VALUES (?, ?, ?)",
            [datetime.now(timezone.utc), tag, elapsed],
        )

    return rel.df()