"""Execute read-only SQL against the financial warehouse.

Defence in depth:
- Engine-level `search_path` and `statement_timeout` via libpq options in
  `api.tools.db.get_engine` — no SQL-side SETs.
- `execution_options(postgresql_readonly=True, postgresql_deferrable=True)`
  — SQLAlchemy emits `SET TRANSACTION READ ONLY DEFERRABLE` internally; we
  don't compose a SQL string for it.
- Hard refusal of any statement that doesn't begin with SELECT or WITH.
- Row cap on the returned result set.
- No commit — read-only transactions roll back on context exit.

The SQL passed in IS a string (it's the LLM's output); nothing on the Python
side concatenates SQL fragments around it.
"""

from __future__ import annotations

import logging

from sqlalchemy import text

from api import langfuse_client as lf
from api.tools.db import get_engine
from api.tools.types import QueryResult

logger = logging.getLogger("nvi.tools.execute_sql")

MAX_ROWS = 1000


def execute_sql(sql: str) -> QueryResult:
    sql_clean = sql.strip().rstrip(";").strip()
    head = sql_clean.split(None, 1)[0].upper() if sql_clean else ""
    if head not in {"SELECT", "WITH"}:
        raise ValueError(f"refusing non-SELECT statement (starts with {head!r})")

    engine = get_engine()
    with lf.span("execute_sql", input={"sql": sql_clean}) as span:
        ro_engine = engine.execution_options(
            postgresql_readonly=True,
            postgresql_deferrable=True,
        )
        with ro_engine.connect() as conn:
            res = conn.execute(text(sql_clean))
            columns = list(res.keys())
            fetched = res.fetchmany(MAX_ROWS + 1)

        truncated = len(fetched) > MAX_ROWS
        rows = [list(r) for r in fetched[:MAX_ROWS]]
        result = QueryResult(
            columns=columns, rows=rows, row_count=len(rows), truncated=truncated
        )
        span.update(output={"row_count": result.row_count, "truncated": truncated})
        return result