"""NL2SQLDeps — per-request context dataclass for NL-to-SQL agents.

Extends ``RequestContext`` with a side-channel that tools write into so the
``/query`` REST endpoint can return structured rows without re-parsing the
agent's prose output.
"""
from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import Any

import asyncpg

from backend.chanakya.a2a_context import RequestContext


@dataclass
class NL2SQLDeps(RequestContext):
    """Generic per-request deps for any NL-to-SQL agent.

    The ``run_query`` tool writes results here during execution.
    The ``/query`` REST handler reads them after ``agent.run()`` completes.

    Attributes:
        query_result:  Raw rows returned by the most recent ``run_query`` call.
                       Each element is a ``{column: value}`` dict.
        generated_sql: The exact SQL that was executed (after LIMIT injection).
        truncated:     ``True`` when the result was capped by the row limit.
        widget_mode:   ``True`` when the caller is the REST ``/query`` endpoint
                       and expects structured widget JSON.  ``False`` (default)
                       means the agent is running inside a pylogue chat session
                       and should respond in plain prose.
        formatted_widgets: Populated by the ``format_response`` tool only when
                       ``widget_mode=True``.
    """
    query_result: list[dict[str, Any]] | None = field(default=None)
    generated_sql: str | None = field(default=None)
    sql_explanation: str | None = field(default=None)
    query_log: list[dict[str, str]] = field(default_factory=list)
    truncated: bool = field(default=False)
    widget_mode: bool = field(default=False)
    formatted_widgets: list[dict[str, Any]] = field(default_factory=list)
    # Populated by WS handlers only; tools push progress events into this queue
    progress_queue: asyncio.Queue | None = field(default=None)
    pool: asyncpg.Pool | None = field(default=None)