"""Shared Pydantic response schemas for all Chanakya agent REST endpoints.

Widget-based response contract
-------------------------------
Every successful ``POST /query`` response returns a list of *widgets* that
the frontend renders in order.

``text``   -- a plain-English explanation of the data.
``chart``  -- a chart matching the frontend ``ChartWidget`` TypeScript interface.
``table``  -- raw tabular data (columns + rows) for a grid view.

Success shape
-------------
{
    "status": "success",
    "data": [
        { "type": "text",  "content": "Here are the top 5 policy types ..." },
        {
            "type": "chart",
            "content": {
                "chartType":  "bar",
                "title":      "Policies by Type",
                "dataset":    [ {"policy_type": "Health", "count": 150} ],
                "labelField": "policy_type",
                "valueField": "count",
                "labels":  ["Health", "Life"],
                "values":  [150, 120],
                "insight": "Health insurance dominates with 150 policies."
            }
        }
    ],
    "error": null
}

Failure shape
-------------
{
    "status": "failure",
    "data": null,
    "error": {
        "message": "<human-readable explanation>",
        "code":    "PIPELINE_ERROR | DB_ERROR | SAFETY_VIOLATION | NO_DATA"
    }
}
"""
from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel


# ---------------------------------------------------------------------------
# Widget -- text
# ---------------------------------------------------------------------------
class TextWidget(BaseModel):
    """A plain-English explanation rendered as prose by the frontend."""
    type: Literal["text"] = "text"
    content: str


# ---------------------------------------------------------------------------
# Widget -- chart  (matches the frontend ChartWidget TypeScript interface)
# ---------------------------------------------------------------------------
class SeriesConfig(BaseModel):
    """A single data series for multi-series charts (stacked_bar, grouped_bar, line)."""
    name: str
    values: list[float] | None = None
    field: str | None = None
    color: str | None = None


class ChartWidgetContent(BaseModel):
    """Payload inside a chart widget -- mirrors the frontend ChartWidget interface."""
    chartType: str
    title: str | None = None
    labels: list[str] | None = None
    values: list[float] | None = None
    series: list[SeriesConfig] | None = None
    unit: str | None = None
    dataset: list[dict[str, Any]] | None = None
    labelField: str | None = None
    valueField: str | None = None
    insight: str | None = None


class ChartWidget(BaseModel):
    type: Literal["chart"] = "chart"
    content: ChartWidgetContent


# ---------------------------------------------------------------------------
# Widget -- table  (raw grid; used when no chart type fits)
# ---------------------------------------------------------------------------
class TableWidgetContent(BaseModel):
    columns: list[str]
    rows: list[dict[str, Any]]


class TableWidget(BaseModel):
    type: Literal["table"] = "table"
    content: TableWidgetContent


# ---------------------------------------------------------------------------
# Error descriptor
# ---------------------------------------------------------------------------
class ErrorInfo(BaseModel):
    """Error payload returned when an agent query fails."""
    message: str
    code: str


# ---------------------------------------------------------------------------
# Top-level request / response
# ---------------------------------------------------------------------------
class TatasteelAgentOutput(BaseModel):
    """Structured output from the Tatasteel NL-to-SQL agent."""
    explanation: str
    what_to_act_on: list[str] = []
    urgency: str = ""
    cost_of_inaction: str = ""


class AgentQueryRequest(BaseModel):
    """Inbound body for a ``POST /query`` endpoint on any Chanakya agent."""
    question: str
    user_email: str | None = None
    conversation_id: int | None = None  # For conversation history continuity
    history: list[dict[str, Any]] | None = None  # Pre-serialized [{role, content}] from chanakya hub


# ---------------------------------------------------------------------------
# UI agent render request / response
# ---------------------------------------------------------------------------
class UIRenderRequest(BaseModel):
    """Inbound body for ``POST /render`` on the UI agent.

    Contains the original user question, the plain-text explanation from the
    data agent, and the raw query rows to visualise.
    """
    question: str
    explanation: str
    rows: list[dict[str, Any]]
    user_email: str | None = None
    conversation_id: int | None = None
    progress_steps: list[dict[str, Any]] = []
    # ISO 4217 currency code for display (detected from user question or DISPLAY_CURRENCY env)
    currency_code: str = "GBP"
    # Market and storyline enrichment (populated by Chanakya after web_search)
    market_context: str | None = None
    storyline_hints: list[str] | None = None
    # Additional decision-support datasets collected by Chanakya after the primary answer query.
    decision_support_context: list[dict[str, Any]] | None = None
    # The exact SQL executed by the NL-to-SQL agent (for display in the UI)
    generated_sql: str | None = None
    # Plain-English explanation of what the SQL does and any translation assumptions
    sql_explanation: str | None = None
    # Decision Context mapping explanation — which DC was matched and why each
    # follow-up question was derived from that DC's decision options / data signals.
    dc_explanation: str | None = None
    # Pre-loaded conversation history passed from rt_agent so ui_agent skips its
    # own DB round-trip. Wire format: [{role: "user"|"assistant", content: str}].
    history: list[dict[str, str]] | None = None
    # Structured decision context — populated when the query is a decision-support question
    what_to_act_on: list[str] | None = None
    urgency: str | None = None
    cost_of_inaction: str | None = None


class AgentQueryResponse(BaseModel):
    """Top-level JSON envelope for all agent query responses.

    ``data`` is a list of widgets ordered for sequential rendering.
    Either ``data`` or ``error`` will be non-null, never both.
    """
    status: str
    data: list[dict[str, Any]] | None = None
    error: ErrorInfo | None = None

    @classmethod
    def success(cls, *, widgets: list[dict[str, Any]]) -> "AgentQueryResponse":
        """Convenience constructor for a successful widget response."""
        return cls(status="success", data=widgets, error=None)

    @classmethod
    def failure(cls, *, message: str, code: str) -> "AgentQueryResponse":
        """Convenience constructor for a failure response."""
        return cls(
            status="failure",
            data=None,
            error=ErrorInfo(message=message, code=code),
        )