"""`run_sql` — execute SQL against the active dataset's DuckDB.

Identical behavior across all three data sources (bundled PKDD, uploaded
file, S3 connector views) — the dataset layer normalises everything to a
DuckDB substrate.

Results are returned as a pandas DataFrame in `payload["df"]`. The
`summary` field gives the LLM a compact preview (shape + head) so it
doesn't have to ingest the whole result set.
"""

from __future__ import annotations

import logging

import pandas as pd
from pydantic import BaseModel, Field

from lexsi_ds.agent.context import AgentContext
from lexsi_ds.agent.tools.base import Tool, ToolResult, ToolSpec

log = logging.getLogger(__name__)

# Preview cap for the LLM observation. Full DataFrame is in payload.
_PREVIEW_ROWS = 20
_PREVIEW_CHARS = 4000


class RunSqlArgs(BaseModel):
    sql: str = Field(..., description="DuckDB-compatible SQL to execute.")
    label: str | None = Field(
        default=None,
        description=(
            "Optional human label for this query (e.g. 'context_df' or "
            "'top10_risky_loans'). Used by the UI/trace; the planner can "
            "pass it through so later tools can reference the result by name."
        ),
    )


def _run(args: RunSqlArgs, ctx: AgentContext) -> ToolResult:
    sql = (args.sql or "").strip().rstrip(";")
    if not sql:
        return ToolResult(ok=False, summary="run_sql: empty SQL", error="empty_sql")

    con = ctx.duck()
    try:
        df = con.execute(sql).df()
    except Exception as e:  # noqa: BLE001 — surface the SQL error to the LLM
        return ToolResult(
            ok=False,
            summary=f"run_sql: execution failed: {type(e).__name__}: {str(e)[:300]}",
            error=str(e),
        )
    finally:
        con.close()

    # Stash on context so explainability tools can find the row set by label.
    if args.label:
        ctx.cache[f"sql_result:{args.label}"] = df
    ctx.cache["last_sql_result"] = df
    ctx.cache["last_sql"] = sql

    preview = _preview(df)
    summary = (
        f"run_sql: {len(df):,} rows × {len(df.columns)} cols"
        + (f"  (label: {args.label})" if args.label else "")
        + f"\ncolumns: {list(df.columns)}\n"
        + f"head:\n{preview}"
    )

    return ToolResult(
        ok=True,
        summary=summary,
        payload={"df": df, "n_rows": len(df), "columns": list(df.columns),
                 "label": args.label},
    )


def _preview(df: pd.DataFrame) -> str:
    if df.empty:
        return "(empty)"
    head = df.head(_PREVIEW_ROWS).to_string(index=False, max_cols=12)
    if len(head) > _PREVIEW_CHARS:
        head = head[:_PREVIEW_CHARS] + "\n…[truncated]"
    return head


class _RunSqlTool:
    spec = ToolSpec(
        name="run_sql",
        description=(
            "Execute DuckDB-compatible SQL against the active dataset and "
            "return the result as a DataFrame. "
            "PREFER `text_to_sql` for any SQL involving JOINs, conditional "
            "aggregates, or string-value filters — that tool sees the full "
            "KG and is far less likely to pick the wrong join path. Use "
            "`run_sql` inline ONLY for trivial single-table queries (e.g. "
            "`SELECT AVG(amount) FROM fin_loan`) OR to execute SQL that "
            "`text_to_sql` returned in a previous step. "
            "Pass `label` to give the result a name later tools can "
            "reference (e.g. label='context_df' before `train_tabular_model`)."
        ),
        args_schema=RunSqlArgs,
        returns=(
            "summary: '<N> rows × <M> cols', columns list, and the head of "
            "the result. payload['df'] holds the full pandas DataFrame. "
            "0-row results on a string filter mean the value is wrong — "
            "call `sample_values` and retry."
        ),
    )

    def run(self, args: RunSqlArgs, ctx: AgentContext) -> ToolResult:
        return _run(args, ctx)


TOOL: Tool = _RunSqlTool()