"""
VSS SQL Agent
Generates and executes read-only SQL against ClickHouse using a LangGraph pipeline.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional, Sequence, Type, TypedDict, Union

from config.settings import get_settings
from orchestration.state import VSSState, add_trace
from langchain_community.utilities import SQLDatabase
from langchain_core.callbacks import AsyncCallbackManagerForToolRun, CallbackManagerForToolRun
from langchain_core.language_models import BaseLanguageModel
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import BaseTool
from langchain_community.tools.sql_database.prompt import QUERY_CHECKER
from langchain_ollama import ChatOllama
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, ConfigDict, Field, model_validator
from sqlalchemy.engine import Result

logger = logging.getLogger("vss.sql_agent")

# ─────────────────────────────────────────────────────────────────────────────
# Configuration
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class AgentConfig:
    """
    Runtime configuration for the SQL agent.

    No field carries a default value — every value must be explicitly
    provided by the constructor so configuration provenance is always clear.
    Use from_settings() (canonical) to build an instance; the legacy
    from_env() is kept only for backward-compatibility.
    """

    database_url: str
    model_name: str
    ollama_base_url: str
    sql_dialect: str
    temperature: float
    max_retries: int
    top_k_tables: int
    sample_rows_per_table: int
    schema_loading_mode: Literal["selected_tables", "all_tables"]
    enforce_read_only_sql: bool

    @classmethod
    def from_settings(cls) -> "AgentConfig":
        """
        Build config from VSSSettings — the single source of truth.

        Reads the singleton from config.settings.get_settings(), which loads
        vss/.env automatically via pydantic-settings.  Every field maps to a
        VSS_-prefixed env var:

            VSS_CLICKHOUSE__*             → database_url (built as sqlalchemy_url)
            VSS_LLM__MODEL_NAME           → model_name
            VSS_LLM__ENDPOINT             → ollama_base_url
            VSS_LLM__TEMPERATURE          → temperature
            VSS_SQL_DIALECT               → sql_dialect
            VSS_SQL_MAX_RETRIES           → max_retries
            VSS_SQL_TOP_K_TABLES          → top_k_tables
            VSS_SQL_SAMPLE_ROWS_PER_TABLE → sample_rows_per_table
            VSS_SQL_SCHEMA_LOADING_MODE   → schema_loading_mode
            VSS_SQL_ENFORCE_READ_ONLY     → enforce_read_only_sql
        """
        s = get_settings()
        return cls(
            database_url=s.clickhouse.sqlalchemy_url,
            model_name=s.llm.model_name,
            ollama_base_url=s.llm.endpoint,
            sql_dialect=s.sql_dialect,
            temperature=s.llm.temperature,
            max_retries=s.sql_max_retries,
            top_k_tables=s.sql_top_k_tables,
            sample_rows_per_table=s.sql_sample_rows_per_table,
            schema_loading_mode=s.sql_schema_loading_mode,  # type: ignore[arg-type]
            enforce_read_only_sql=s.sql_enforce_read_only,
        )

    @classmethod
    def from_env(cls) -> "AgentConfig":
        """
        Build config from raw, unprefixed environment variables.

        Deprecated — use from_settings() instead.
        Kept only for backward-compatibility with standalone scripts that set
        the legacy env vars directly.  When a legacy var is absent, the value
        falls back to the VSSSettings singleton (which reads vss/.env) rather
        than a hardcoded constant, so the two paths stay in sync.
        """
        import warnings
        warnings.warn(
            "AgentConfig.from_env() is deprecated. Use AgentConfig.from_settings() instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        s = get_settings()
        return cls(
            database_url=os.environ.get("CLICKHOUSE_URL") or s.clickhouse.sqlalchemy_url,
            model_name=os.environ.get("OLLAMA_MODEL") or s.llm.model_name,
            ollama_base_url=os.environ.get("OLLAMA_BASE_URL") or s.llm.endpoint,
            sql_dialect=os.environ.get("SQL_DIALECT") or s.sql_dialect,
            temperature=float(os.environ.get("LLM_TEMPERATURE", s.llm.temperature)),
            max_retries=int(os.environ.get("MAX_RETRIES", s.sql_max_retries)),
            top_k_tables=int(os.environ.get("TOP_K_TABLES", s.sql_top_k_tables)),
            sample_rows_per_table=int(os.environ.get("SAMPLE_ROWS_PER_TABLE", s.sql_sample_rows_per_table)),
            schema_loading_mode=os.environ.get("SCHEMA_LOADING_MODE") or s.sql_schema_loading_mode,  # type: ignore[arg-type]
            enforce_read_only_sql=os.environ.get("ENFORCE_READ_ONLY_SQL", str(s.sql_enforce_read_only)).lower() == "true",
        )


# ─────────────────────────────────────────────────────────────────────────────
# Table metadata
# ─────────────────────────────────────────────────────────────────────────────

TABLE_DESCRIPTIONS: Dict[str, str] = {
    "alert_instances": (
        "Stores individual alert events generated by different applications for each camera. "
        "Use this table when the user asks about alerts, alert counts, alert history, or alert details."
    ),
    "incidents": (
        "Stores detected incidents with their severity levels across different applications and cameras. "
        "Use this table when the user asks about incidents, incident severity, incident history, or incident statistics."
    ),
    "aggregated_analytics_totals": (
        "Stores aggregated analytics summaries (counts and totals) for different object categories "
        "across applications and cameras over time. Use this table when the user asks for totals, "
        "summaries, trends, or analytics reports."
    ),
}

INCLUDED_TABLES = list(TABLE_DESCRIPTIONS.keys())

SEMANTIC_COLUMN_MAP: Dict[str, str] = {
    "aggregated_analytics_totals": "applicationName",
    "alert_instances": "applicationName",
    "incidents": "incidentType",
}


# ─────────────────────────────────────────────────────────────────────────────
# Pydantic schemas (LLM structured outputs)
# ─────────────────────────────────────────────────────────────────────────────

class IntentDecision(BaseModel):
    use_sql_agent: bool = Field(..., description="True if question requires database lookup")
    reason: str = Field(..., description="Why the route decision was made")


class TableSelection(BaseModel):
    selected_tables: List[str] = Field(default_factory=list)
    reason: str = Field(..., description="Why these tables were selected")


# ─────────────────────────────────────────────────────────────────────────────
# LangGraph state
# ─────────────────────────────────────────────────────────────────────────────

class AgentState(TypedDict, total=False):
    user_query: str
    route_to_sql: bool
    intent_reason: str
    candidate_tables: List[str]
    selected_tables: List[str]
    schema_context: str
    sql_attempt: int
    sql_query: str
    checked_sql_query: str
    execution_result: str
    execution_error: str
    should_retry: bool
    retry_reason: str
    final_answer: str
    trace: List[str]


# ─────────────────────────────────────────────────────────────────────────────
# SQL safety helpers
# ─────────────────────────────────────────────────────────────────────────────

_READ_ONLY_RE = re.compile(r"^\s*(select|with)\b", re.IGNORECASE)
_FORBIDDEN_RE = re.compile(
    r"\b(insert|update|delete|drop|create|alter|truncate|grant|revoke)\b",
    re.IGNORECASE,
)


def sanitize_sql(sql: str) -> str:
    """Strip markdown fences, backticks, and enforce read-only policy."""
    cleaned = sql.strip()
    
    # Extract SQL from a ```sql ... ``` block anywhere in the text
    match = re.search(r"```sql\s*(.*?)\s*```", cleaned, re.IGNORECASE | re.DOTALL)
    if match:
        cleaned = match.group(1).strip()
    else:
        # Fallback: strip fences if they're at start/end
        cleaned = re.sub(r"^```sql\s*", "", cleaned, flags=re.IGNORECASE)
        cleaned = re.sub(r"^```\s*", "", cleaned)
        cleaned = re.sub(r"\s*```$", "", cleaned)
        cleaned = cleaned.strip()
    
    cleaned = cleaned.replace("`", "")

    if not _READ_ONLY_RE.search(cleaned):
        raise ValueError("Only SELECT/WITH queries are permitted.")
    if _FORBIDDEN_RE.search(cleaned):
        raise ValueError("Query contains a forbidden keyword.")
    return cleaned


# ─────────────────────────────────────────────────────────────────────────────
# LLM helpers
# ─────────────────────────────────────────────────────────────────────────────

def build_llm(config: AgentConfig) -> ChatOllama:
    return ChatOllama(
        model=config.model_name,
        base_url=config.ollama_base_url,
        temperature=config.temperature,
    )


def _extract_json_object(text: str) -> Dict[str, Any]:
    """Extract the first JSON object from a model text response."""
    cleaned = text.strip()
    cleaned = re.sub(r"^```json\s*", "", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"^```\s*", "", cleaned)
    cleaned = re.sub(r"\s*```$", "", cleaned)

    try:
        parsed = json.loads(cleaned)
        if isinstance(parsed, dict):
            return parsed
    except Exception:
        pass

    match = re.search(r"\{.*\}", cleaned, flags=re.DOTALL)
    if not match:
        raise ValueError(f"No JSON object found in model output: {text!r}")
    parsed = json.loads(match.group(0))
    if not isinstance(parsed, dict):
        raise ValueError(f"Expected JSON object, got {type(parsed)}")
    return parsed


def invoke_structured(llm: Any, schema: Type[BaseModel], task_prompt: str) -> BaseModel:
    """Invoke an LLM and parse its response into *schema*."""
    prompt = (
        "Return ONLY a valid JSON object. No markdown, no explanations.\n"
        f"JSON must satisfy this schema:\n{json.dumps(schema.model_json_schema(), ensure_ascii=True)}\n\n"
        f"Task:\n{task_prompt}"
    )
    response = llm.invoke(prompt)
    content = response.content if hasattr(response, "content") else str(response)
    payload = _extract_json_object(str(content))
    return schema.model_validate(payload)


# ─────────────────────────────────────────────────────────────────────────────
# Result reducer
# ─────────────────────────────────────────────────────────────────────────────

def _row_hash(row: Any) -> str:
    return hashlib.md5(str(row).encode()).hexdigest()


def reduce_sql_result(rows: Any, *, max_rows: int = 20, deduplicate: bool = True) -> Any:
    """Trim and optionally deduplicate SQL results before sending to LLM."""
    if rows is None or not isinstance(rows, list):
        return rows[:4000] if isinstance(rows, str) else rows

    if deduplicate:
        seen: set = set()
        unique: List[Any] = []
        for r in rows:
            h = _row_hash(r)
            if h not in seen:
                seen.add(h)
                unique.append(r)
        rows = unique

    if len(rows) <= max_rows:
        return rows

    head = max_rows // 2
    return rows[:head] + rows[-(max_rows - head):]


# ─────────────────────────────────────────────────────────────────────────────
# Semantic context helper
# ─────────────────────────────────────────────────────────────────────────────

def get_semantic_table_context(db: SQLDatabase, table_name: str) -> str:
    """Return a sample of distinct key-column values for grounding the LLM."""
    column = SEMANTIC_COLUMN_MAP.get(table_name)
    if not column:
        return ""
    query = (
        f"SELECT * FROM default.{table_name} "
        f"WHERE {column} IS NOT NULL LIMIT 1 BY applicationName LIMIT 50"
    )
    try:
        result = db.run_no_throw(query)
        return f"\nDistinct {column} values:\n{result}\n"
    except Exception as exc:
        logger.warning("Failed to fetch semantic context for %s: %s", table_name, exc)
        return ""


# ─────────────────────────────────────────────────────────────────────────────
# Custom SQL tools
# ─────────────────────────────────────────────────────────────────────────────

class _SQLBaseTool(BaseModel):
    db: SQLDatabase = Field(exclude=True)
    model_config = ConfigDict(arbitrary_types_allowed=True)


class _QueryInput(BaseModel):
    query: str = Field(..., description="A correct, read-only SQL query.")


class CustomQueryTool(_SQLBaseTool, BaseTool):
    name: str = "sql_db_query"
    description: str = (
        "Execute a SQL query against the database and return the result. "
        "Returns an error message if the query fails."
    )
    args_schema: Type[BaseModel] = _QueryInput
    enforce_read_only_sql: bool = True

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Union[str, Sequence[Dict[str, Any]], Result]:
        if self.enforce_read_only_sql and not _READ_ONLY_RE.match(query.strip()):
            return "Error: Non read-only SQL is blocked by policy."
        return self.db.run_no_throw(query)


class _SchemaInput(BaseModel):
    table_names: str = Field(..., description="Comma-separated table names.")


class CustomSchemaTool(_SQLBaseTool, BaseTool):
    name: str = "sql_db_schema"
    description: str = "Get schema and sample rows for specified SQL tables."
    args_schema: Type[BaseModel] = _SchemaInput
    table_descriptions: Dict[str, str] = Field(default_factory=dict)

    def _run(
        self,
        table_names: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        table_list = [t.strip() for t in table_names.split(",") if t.strip()]
        base_info = self.db.get_table_info_no_throw(table_list)
        if not table_list:
            return base_info
        desc_block = "\n".join(
            f"- {t}: {self.table_descriptions.get(t, 'No description provided.')}"
            for t in table_list
        )
        return f"Table descriptions:\n{desc_block}\n\n{base_info}"


class _ListInput(BaseModel):
    tool_input: str = Field("", description="Empty string.")


class CustomListTablesTool(_SQLBaseTool, BaseTool):
    name: str = "sql_db_list_tables"
    description: str = "Returns a comma-separated list of available tables."
    args_schema: Type[BaseModel] = _ListInput

    def _run(
        self,
        tool_input: str = "",
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return ", ".join(self.db.get_usable_table_names())


class _CheckerInput(BaseModel):
    query: str = Field(..., description="SQL query to check and improve.")


class CustomQueryCheckerTool(_SQLBaseTool, BaseTool):
    name: str = "sql_db_query_checker"
    description: str = "Use LLM to validate SQL before execution."
    args_schema: Type[BaseModel] = _CheckerInput
    llm: BaseLanguageModel
    sql_dialect: str = "clickhouse"
    llm_chain: Any = Field(init=False)

    @model_validator(mode="before")
    @classmethod
    def _init_chain(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        if "llm_chain" not in values:
            from langchain_classic.chains.llm import LLMChain

            values["llm_chain"] = LLMChain(
                llm=values["llm"],
                prompt=PromptTemplate(
                    template=QUERY_CHECKER,
                    input_variables=["dialect", "query"],
                ),
            )
        if values["llm_chain"].prompt.input_variables != ["dialect", "query"]:
            raise ValueError("LLM chain must use input vars ['dialect', 'query'].")
        return values

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return self.llm_chain.predict(
            query=query,
            dialect=self.sql_dialect,
            callbacks=run_manager.get_child() if run_manager else None,
        )

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return await self.llm_chain.apredict(
            query=query,
            dialect=self.sql_dialect,
            callbacks=run_manager.get_child() if run_manager else None,
        )


# ─────────────────────────────────────────────────────────────────────────────
# Toolkit
# ─────────────────────────────────────────────────────────────────────────────

class SQLAgentToolkit(BaseModel):
    db: SQLDatabase = Field(exclude=True)
    llm: BaseLanguageModel = Field(exclude=True)
    table_descriptions: Dict[str, str] = Field(default_factory=dict)
    sql_dialect: str = "clickhouse"
    enforce_read_only_sql: bool = True

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def get_tools(self) -> List[BaseTool]:
        return [
            CustomQueryTool(db=self.db, enforce_read_only_sql=self.enforce_read_only_sql),
            CustomSchemaTool(db=self.db, table_descriptions=self.table_descriptions),
            CustomListTablesTool(db=self.db),
            CustomQueryCheckerTool(db=self.db, llm=self.llm, sql_dialect=self.sql_dialect),
        ]


# ─────────────────────────────────────────────────────────────────────────────
# LangGraph agent
# ─────────────────────────────────────────────────────────────────────────────

class SQLAgent:
    """LangGraph-based SQL agent that uses the custom toolkit."""

    def __init__(
        self,
        config: AgentConfig,
        llm: Any,
        toolkit: SQLAgentToolkit,
        table_descriptions: Dict[str, str],
    ) -> None:
        self.config = config
        self.llm = llm
        self.table_descriptions = table_descriptions

        tools = {tool.name: tool for tool in toolkit.get_tools()}
        self._query_tool: CustomQueryTool = tools["sql_db_query"]  # type: ignore[assignment]
        self._schema_tool: CustomSchemaTool = tools["sql_db_schema"]  # type: ignore[assignment]
        self._list_tool: CustomListTablesTool = tools["sql_db_list_tables"]  # type: ignore[assignment]
        self._checker_tool: CustomQueryCheckerTool = tools["sql_db_query_checker"]  # type: ignore[assignment]

        self._graph = self._build_graph()

    # ── Graph construction ────────────────────────────────────────────────

    def _build_graph(self):
        builder = StateGraph(AgentState)

        builder.add_node("route_query", self._route_query)
        builder.add_node("non_sql_response", self._non_sql_response)
        builder.add_node("discover_tables", self._discover_tables)
        builder.add_node("select_tables", self._select_tables)
        builder.add_node("load_schema", self._load_schema)
        builder.add_node("generate_sql", self._generate_sql)
        builder.add_node("check_sql", self._check_sql)
        builder.add_node("execute_sql", self._execute_sql)
        builder.add_node("assess_execution", self._assess_execution)
        builder.add_node("finalize_answer", self._finalize_answer)

        builder.add_edge(START, "route_query")
        builder.add_conditional_edges(
            "route_query",
            lambda s: "sql" if s.get("route_to_sql") else "non_sql",
            {"sql": "discover_tables", "non_sql": "non_sql_response"},
        )
        builder.add_edge("non_sql_response", END)
        builder.add_edge("discover_tables", "select_tables")
        builder.add_edge("select_tables", "load_schema")
        builder.add_edge("load_schema", "generate_sql")
        builder.add_edge("generate_sql", "check_sql")
        builder.add_edge("check_sql", "execute_sql")
        builder.add_edge("execute_sql", "assess_execution")
        builder.add_conditional_edges(
            "assess_execution",
            lambda s: "retry" if s.get("should_retry") else "final",
            {"retry": "generate_sql", "final": "finalize_answer"},
        )
        builder.add_edge("finalize_answer", END)
        return builder.compile()

    # ── Graph nodes ───────────────────────────────────────────────────────

    def _route_query(self, state: AgentState) -> AgentState:
        add_trace(state, "SQLAgent: Starting")
        add_trace(state, "Route: Classifying intent")
        prompt = (
            "You are an intent classifier for a SQL analytics assistant.\n"
            "Return true when database querying is needed.\n"
            "Return false for greetings, chitchat, or general capability questions.\n\n"
            f"User question: {state['user_query']}"
        )
        decision = invoke_structured(self.llm, IntentDecision, prompt)
        add_trace(state, f"Route: {'SQL' if decision.use_sql_agent else 'non-SQL'} — {decision.reason}")
        return {"route_to_sql": decision.use_sql_agent, "intent_reason": decision.reason, "sql_attempt": 0}

    def _non_sql_response(self, state: AgentState) -> AgentState:
        add_trace(state, "Route: Non-SQL response")
        return {
            "final_answer": (
                "I can query your database and answer questions related to "
                "alerts, incidents, trends, and operations data. How can I help you?"
            )
        }

    def _discover_tables(self, state: AgentState) -> AgentState:
        add_trace(state, "Discovering available tables")
        tables_csv = self._list_tool.run("")
        candidates = [t.strip() for t in tables_csv.split(",") if t.strip()]
        add_trace(state, f"Tables discovered: {candidates}")
        return {"candidate_tables": candidates}

    def _select_tables(self, state: AgentState) -> AgentState:
        add_trace(state, "Selecting relevant tables")
        candidates = state.get("candidate_tables", [])
        if not candidates:
            add_trace(state, "Table selection error: no tables available")
            return {
                "selected_tables": [],
                "execution_error": "No tables available in database.",
                "should_retry": False,
            }

        if self.config.schema_loading_mode == "all_tables":
            add_trace(state, f"Tables selected (all): {candidates}")
            return {"selected_tables": candidates}

        table_context = "\n".join(
            f"- {t}: {self.table_descriptions.get(t, 'No description provided.')}"
            for t in candidates
        )
        prompt = (
            f"Select the minimal set of SQL tables needed for the question.\n"
            f"Max tables: {self.config.top_k_tables}\n\n"
            f"Question: {state['user_query']}\n\nTables:\n{table_context}"
        )
        selection = invoke_structured(self.llm, TableSelection, prompt)
        selected = [t for t in selection.selected_tables if t in candidates]
        chosen = selected or candidates[: self.config.top_k_tables]
        add_trace(state, f"Tables selected: {chosen} — {selection.reason}")
        return {"selected_tables": chosen}

    def _load_schema(self, state: AgentState) -> AgentState:
        selected = state.get("selected_tables", [])
        add_trace(state, f"Loading schema for: {selected}")
        if not selected:
            add_trace(state, "Schema load error: no tables selected")
            return {"schema_context": "", "execution_error": "No tables selected for schema loading."}

        base_schema = self._schema_tool.run(", ".join(selected))
        semantic_blocks = [
            get_semantic_table_context(self._query_tool.db, t) for t in selected
        ]
        add_trace(state, f"Schema loaded for: {selected}")
        return {"schema_context": base_schema + "".join(semantic_blocks)}

    def _generate_sql(self, state: AgentState) -> AgentState:
        attempt = int(state.get("sql_attempt", 0)) + 1
        add_trace(state, f"Generating SQL (attempt {attempt}/{self.config.max_retries})")
        if state.get("retry_reason"):
            add_trace(state, f"Retry reason: {state['retry_reason']}")
        prompt = (
            "Generate one read-only SQL query for this question.\n"
            "RULES:\n"
            "- MUST start with SELECT\n"
            "- DO NOT use backticks\n"
            "- DO NOT use INSERT/UPDATE/DELETE/CREATE\n"
            "- Use plain table names only\n\n"
            "CLICKHOUSE RULES:\n"
            "- Use yesterday(), today(), now() for dates\n"
            "- DO NOT use toDateTime('yesterday') or unsupported INTERVAL syntax\n\n"
            f"SQL Dialect: {self.config.sql_dialect}\n"
            f"Attempt: {attempt}/{self.config.max_retries}\n"
            f"Question: {state['user_query']}\n"
            f"Retry reason: {state.get('retry_reason', '')}\n\n"
            f"Schema context:\n{state.get('schema_context', '')}"
        )
        try:
            sql_query = sanitize_sql(str(self.llm.invoke(prompt).content))
            add_trace(state, f"SQL generated: {sql_query[:100]}...")
        except ValueError as e:
            add_trace(state, f"SQL generation error: {e}")
            raise
        return {
            "sql_attempt": attempt,
            "sql_query": sql_query,
            "execution_error": "",
            "should_retry": False,
            "retry_reason": "",
        }

    def _check_sql(self, state: AgentState) -> AgentState:
        add_trace(state, "Checking SQL validity")
        query = state.get("sql_query", "")
        if not query:
            add_trace(state, "SQL check error: generated SQL is empty")
            return {"checked_sql_query": "", "execution_error": "Generated SQL is empty."}

        checked_raw = self._checker_tool.run(query)
        try:
            checked = sanitize_sql(checked_raw or query)
        except ValueError:
            checked = query  # fall back to original if checker output is invalid

        if not checked.lower().startswith(("select", "with")):
            checked = query
        add_trace(state, f"SQL checked: {checked[:100]}...")
        return {"checked_sql_query": checked}

    def _execute_sql(self, state: AgentState) -> AgentState:
        if state.get("execution_error"):
            add_trace(state, f"SQL execution skipped: prior error — {state['execution_error']}")
            return state

        sql = state.get("checked_sql_query") or state.get("sql_query", "")
        if not sql:
            add_trace(state, "SQL execution error: no SQL available")
            return {"execution_error": "No SQL available to execute.", "execution_result": ""}

        add_trace(state, f"Executing SQL: {sql[:100]}...")
        logger.info("Executing SQL: %s", sql)
        result = self._query_tool.run(sql)
        logger.info("SQL result: %s", result)

        if isinstance(result, str) and result.strip().lower().startswith("error"):
            add_trace(state, f"SQL execution error: {result[:200]}")
            return {"execution_error": result, "execution_result": ""}

        add_trace(state, f"SQL executed: {str(result)[:100]}...")
        return {"execution_result": str(result), "execution_error": ""}

    def _assess_execution(self, state: AgentState) -> AgentState:
        attempt = int(state.get("sql_attempt", 0))
        can_retry = attempt < self.config.max_retries

        if state.get("execution_error"):
            reason = "Execution error; regenerate SQL."
            if can_retry:
                add_trace(state, f"SQL retry {attempt}: {state['execution_error']}")
            else:
                add_trace(state, f"SQL failed after {attempt} attempt(s): {state['execution_error']}")
            return {"should_retry": can_retry, "retry_reason": reason}

        if str(state.get("execution_result", "")).strip().lower() in {"", "[]", "none"}:
            reason = "Empty result; adjust query strategy."
            if can_retry:
                add_trace(state, f"SQL retry {attempt}: empty result returned")
            else:
                add_trace(state, f"SQL failed after {attempt} attempt(s): empty result")
            return {"should_retry": can_retry, "retry_reason": reason}

        add_trace(state, f"Execution assessed: success after {attempt} attempt(s)")
        return {"should_retry": False, "retry_reason": ""}

    def _finalize_answer(self, state: AgentState) -> AgentState:
        add_trace(state, "Finalizing answer")
        if state.get("execution_error"):
            add_trace(state, f"Answer: failed after {state.get('sql_attempt', 0)} attempt(s)")
            return {
                "final_answer": (
                    f"I could not complete your request after "
                    f"{state.get('sql_attempt', 0)} attempt(s). Please try again later."
                )
            }

        reduced = reduce_sql_result(state.get("execution_result"), max_rows=20, deduplicate=True)
        prompt = (
            "Convert the SQL result into clear, concise natural language.\n"
            "Rules:\n"
            "- Answer in English only\n"
            "- Be professional and user-friendly\n"
            "- Include key numbers\n"
            "- Avoid technical jargon or code\n\n"
            f"Question: {state['user_query']}\n"
            f"SQL: {state.get('checked_sql_query') or state.get('sql_query', '')}\n"
            f"Result: {reduced}"
        )
        answer = self.llm.invoke(prompt).content
        add_trace(state, f"Answer finalized: {str(answer)[:100]}...")
        return {"final_answer": str(answer)}

    # ── Public API ────────────────────────────────────────────────────────

    def run(self, user_query: str) -> AgentState:
        return self._graph.invoke({"user_query": user_query, "trace": []})


# ─────────────────────────────────────────────────────────────────────────────
# Lazy initialisation — runs only on first call, not at import time
# ─────────────────────────────────────────────────────────────────────────────

_agent: Optional[SQLAgent] = None


def _get_agent() -> SQLAgent:
    """Return the SQL agent singleton, initialising it on first use."""
    global _agent
    if _agent is not None:
        return _agent

    _config = AgentConfig.from_settings()
    _llm = build_llm(_config)
    _db = SQLDatabase.from_uri(
        _config.database_url,
        schema="default",
        sample_rows_in_table_info=0,
        include_tables=INCLUDED_TABLES,
        view_support=True,
        lazy_table_reflection=True,
    )
    _toolkit = SQLAgentToolkit(
        db=_db,
        llm=_llm,
        table_descriptions=TABLE_DESCRIPTIONS,
        sql_dialect=_config.sql_dialect,
        enforce_read_only_sql=_config.enforce_read_only_sql,
    )
    _agent = SQLAgent(
        config=_config,
        llm=_llm,
        toolkit=_toolkit,
        table_descriptions=TABLE_DESCRIPTIONS,
    )
    logger.info("VSS SQL agent initialised.")
    return _agent


# ─────────────────────────────────────────────────────────────────────────────
# Public helper
# ─────────────────────────────────────────────────────────────────────────────
def ask_sql_agent(question: str) -> Dict[str, Any]:
    """Run the SQL agent and return a structured result dict."""
    state = _get_agent().run(user_query=question)
    return {
        "answer": state.get("final_answer", ""),
        "sql_generated": state.get("sql_query", ""),
        "sql_checked": state.get("checked_sql_query", ""),
        "attempts": state.get("sql_attempt", 0),
        "selected_tables": state.get("selected_tables", []),
        "execution_error": state.get("execution_error", ""),
        "execution_result": state.get("execution_result", ""),
        "intent_reason": state.get("intent_reason", ""),
        "trace": state.get("trace", []),
    }


# ─────────────────────────────────────────────────────────────────────────────
# Orchestration entry point  (called by graph.py → VSSOrchestrator._sql node)
# ─────────────────────────────────────────────────────────────────────────────

async def process_sql(state: VSSState) -> VSSState:
    """
    LangGraph node: bridge between VSSState (orchestrator) and the internal
    AgentState-based SQL pipeline.

    Reads:  state["user_query"]
    Writes: state["sql_query"], state["sql_result"], state["sql_result_formatted"],
            state["sql_error"], state["escalated"], state["escalation_reason"],
            state["trace"]
    """
    add_trace(state, "SQLAgent: process_sql called")

    try:
        sql_agent = _get_agent()
        user_query: str = state.get("user_query", "")

        # The internal SQL pipeline is fully synchronous (LangChain + LangGraph
        # sync invocations via ChatOllama).  Run it in a thread-pool executor so
        # it does not block the asyncio event loop that serves FastAPI requests.
        result: AgentState = await asyncio.to_thread(
            sql_agent.run, user_query=user_query
        )

        # ── SQL query ────────────────────────────────────────────────────────
        state["sql_query"] = (
            result.get("checked_sql_query") or result.get("sql_query", "")
        )

        # ── Errors ───────────────────────────────────────────────────────────
        execution_error = result.get("execution_error", "")
        state["sql_error"] = execution_error if execution_error else None

        # ── Raw execution result string ───────────────────────────────────────
        raw = result.get("execution_result", "")
        has_data = bool(raw and raw.strip() not in ("", "[]", "None", "none"))

        # ── Rows: try to parse to list-of-dicts for structured formatting ─────
        # SQLDatabase.run_no_throw() returns a string like "[(val,), ...]"
        # (tuples, not dicts). We attempt to parse but only keep actual dicts.
        rows: List[Any] = []
        if has_data:
            try:
                import ast
                parsed = ast.literal_eval(raw)
                if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict):
                    rows = parsed
            except Exception:
                pass  # leave rows as [] — formatted fallback will be used
        state["sql_result"] = rows

        # ── Formatted result: use the agent's final_answer (LLM-generated) ──
        final_answer = result.get("final_answer", "")
        state["sql_result_formatted"] = final_answer or raw or "No results."

        # ── Escalation: based on raw execution result, not parsed rows ────────
        # This correctly signals "no data" without depending on parse success.
        if not has_data and not execution_error:
            state["escalated"] = True
            state["escalation_reason"] = "SQL returned no results"
            add_trace(state, "SQLAgent: no results — escalating")
        else:
            state["escalated"] = False

        # ── Merge internal trace into VSSState trace ──────────────────────
        for entry in result.get("trace", []):
            state["trace"].append(entry)

        add_trace(
            state,
            f"SQLAgent: completed ({result.get('sql_attempt', 0)} attempt(s))"
        )

    except Exception as exc:
        logger.error("process_sql failed: %s", exc, exc_info=True)
        state["sql_error"] = str(exc)
        add_trace(state, f"SQLAgent: error — {exc}")

    return state

