"""Few-shot examples for the planner prompt.

Two illustrative (question -> TaskList) pairs that teach the OUTPUT SHAPE:
stages, dependency edges, ordered tool-call chains, inline QueryIR,
"${t<id>}" placeholders, and the assumed data-flow convention — `retrieve_data`
pulls rows, then a composite `analyze_*` tool consumes them via a `data` placeholder
referencing the upstream result's column aliases (Pattern A; the tool team may
instead pick self-fetch by `source_id`, in which case these examples are reshaped
to match — see registry.py). They reference a hypothetical sales catalog
(`src_sales` / `t_orders`); these ids are part of the illustration and are not
validated against the user's real catalog. v1 is descriptive/diagnostic — no
modeling tasks.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3 (Examples A and B).
"""

from __future__ import annotations

from .schemas import Task, TaskList, ToolCall

# --------------------------------------------------------------------------- #
# Example A — exploratory, no modeling.
# "Which product categories drove last quarter's revenue?"
# Shows: retrieve_data pulls rows -> analyze_aggregate sums revenue per
# category in one call (no manual per-category queries).
# --------------------------------------------------------------------------- #

_EXAMPLE_A = TaskList(
    plan_id="example_a",
    goal_restated="Identify which product categories contributed most to last quarter's revenue.",
    assumptions=["'last quarter' = 2026-01-01 to 2026-03-31."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category, revenue, and order date.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; the 3 needed columns are present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull last quarter's order-level category and revenue rows.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_order_date",
                                    "op": "between",
                                    "value": ["2026-01-01", "2026-03-31"],
                                    "value_type": "date",
                                }
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="quarter_rows",
            success_criteria="Produced last quarter's order rows with category and revenue.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Sum revenue per category for the quarter.",
            tool_calls=[
                ToolCall(
                    tool="analyze_aggregate",
                    args={
                        "data": "${t2}",
                        "aggregations": {"revenue": ["sum"]},
                        "group_by": ["category"],
                    },
                )
            ],
            expected_output="category_revenue",
            success_criteria="Produced total revenue per category, one row each.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example B — descriptive / trend.
# "How has monthly revenue trended by region this year, and what's unusual?"
# --------------------------------------------------------------------------- #

_EXAMPLE_B = TaskList(
    plan_id="example_b",
    goal_restated="Describe this year's monthly revenue trend and flag unusual months.",
    assumptions=["'this year' starts 2026-01-01."],
    open_questions=["'Unusual' is interpreted as months far from the typical monthly revenue."],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order date, revenue, and region.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; the needed columns are present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull this year's order dates, revenue, and region.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {
                                    "kind": "column",
                                    "column_id": "c_order_date",
                                    "alias": "order_date",
                                },
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                                {"kind": "column", "column_id": "c_region", "alias": "region"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_order_date",
                                    "op": ">=",
                                    "value": "2026-01-01",
                                    "value_type": "date",
                                }
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="ytd_rows",
            success_criteria="Produced this year's order-level rows with date, revenue, region.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Bucket revenue into months and summarize the trend and movement.",
            tool_calls=[
                ToolCall(
                    tool="analyze_trend",
                    args={
                        "data": "${t2}",
                        "date_column": "order_date",
                        "value_column": "revenue",
                        "freq": "month",
                        "agg": "sum",
                    },
                )
            ],
            expected_output="monthly_trend",
            success_criteria=(
                "Produced a per-month revenue series with direction and change rate to "
                "flag months above/below the typical level."
            ),
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example C — mixed structured + unstructured.
# "Revenue dipped in Q1 — what happened?"
# Shows: a structured branch (query -> analyze_trend) runs alongside an
# INDEPENDENT retrieve_knowledge branch that pulls qualitative context. Note
# retrieve_knowledge takes a natural-language `query` (NOT a `${t<id>}` data
# placeholder — it is a source, not a consumer) and can run in parallel; the
# Assembler folds the document context into the explanation.
# --------------------------------------------------------------------------- #

_EXAMPLE_C = TaskList(
    plan_id="example_c",
    goal_restated="Explain Q1's revenue dip using both the numbers and the qualitative record.",
    assumptions=["'Q1' = 2026-01-01 to 2026-03-31."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order date and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; date and revenue columns present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull Q1 order dates and revenue.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {
                                    "kind": "column",
                                    "column_id": "c_order_date",
                                    "alias": "order_date",
                                },
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_order_date",
                                    "op": "between",
                                    "value": ["2026-01-01", "2026-03-31"],
                                    "value_type": "date",
                                }
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="q1_rows",
            success_criteria="Produced Q1 order rows with date and revenue.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Summarize the Q1 monthly revenue trend to locate the dip.",
            tool_calls=[
                ToolCall(
                    tool="analyze_trend",
                    args={
                        "data": "${t2}",
                        "date_column": "order_date",
                        "value_column": "revenue",
                        "freq": "month",
                        "agg": "sum",
                    },
                )
            ],
            expected_output="q1_trend",
            success_criteria="Produced a per-month revenue series showing where revenue fell.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
        Task(
            id="t4",
            stage="data_understanding",
            objective="Retrieve qualitative context on Q1 operational events behind a dip.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_knowledge",
                    args={
                        "query": "operational issues, outages, or notable events in Q1 2026",
                        "top_k": 5,
                    },
                )
            ],
            expected_output="q1_context_chunks",
            success_criteria="Produced relevant document chunks about Q1 operations.",
            depends_on=[],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example D — group-by aggregation (analyze_aggregate arg shape).
# "What is the average and total order value per region?"
# Shows the EXACT analyze_aggregate args: `aggregations` is an OBJECT mapping each
# column to a LIST of functions ({"revenue": ["mean", "sum"]}), and `group_by` is a
# SEPARATE array — NOT a nested list of metric specs. Supported funcs: sum, mean,
# count, min, max, median, nunique.
# --------------------------------------------------------------------------- #

_EXAMPLE_D = TaskList(
    plan_id="example_d",
    goal_restated="Report the average and total order value for each region.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes region and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; region and revenue present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull order-level region and revenue.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_region", "alias": "region"},
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="region_rows",
            success_criteria="Produced order rows with region and revenue.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Aggregate mean and total revenue per region.",
            tool_calls=[
                ToolCall(
                    tool="analyze_aggregate",
                    args={
                        "data": "${t2}",
                        "aggregations": {"revenue": ["mean", "sum"]},
                        "group_by": ["region"],
                    },
                )
            ],
            expected_output="region_aggregates",
            success_criteria="Produced one row per region with mean and total revenue.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example E — non-date filters (value_type is the ELEMENT type, never a container).
# "Total revenue for the East and West regions, counting orders of at least 100."
# Shows: an `in` filter over a list of strings uses value_type "string" (NOT
# "list"); a numeric comparison uses "decimal"; and every filter carries a
# value_type copied from the column's catalog [data_type].
# --------------------------------------------------------------------------- #

_EXAMPLE_E = TaskList(
    plan_id="example_e",
    goal_restated="Total revenue for the East and West regions, orders of at least 100.",
    assumptions=["'at least 100' filters order revenue >= 100."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes region and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; region and revenue present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull East/West order rows of at least 100 with region and revenue.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_region", "alias": "region"},
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_region",
                                    "op": "in",
                                    "value": ["East", "West"],
                                    "value_type": "string",
                                },
                                {
                                    "column_id": "c_revenue",
                                    "op": ">=",
                                    "value": 100,
                                    "value_type": "decimal",
                                },
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="filtered_rows",
            success_criteria="Produced East/West order rows above the revenue threshold.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Sum revenue per region.",
            tool_calls=[
                ToolCall(
                    tool="analyze_aggregate",
                    args={
                        "data": "${t2}",
                        "aggregations": {"revenue": ["sum"]},
                        "group_by": ["region"],
                    },
                )
            ],
            expected_output="region_revenue",
            success_criteria="Produced total revenue per region, one row each.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example F — descriptive statistics (analyze_descriptive).
# "Give me the summary statistics for order revenue and quantity."
# Shows: retrieve_data pulls the numeric columns -> analyze_descriptive summarizes
# them. The `data` comes from the retrieve_data task (t2), NEVER from the check_data
# inspection step (t1) — check_data returns column metadata, not data rows.
# --------------------------------------------------------------------------- #

_EXAMPLE_F = TaskList(
    plan_id="example_f",
    goal_restated="Summarize the distribution of order revenue and quantity.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order revenue and quantity.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; the needed columns are present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull the order revenue and quantity rows.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                                {"kind": "column", "column_id": "c_quantity", "alias": "quantity"},
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="order_rows",
            success_criteria="Produced order-level rows with revenue and quantity.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Summarize the distribution of revenue and quantity.",
            tool_calls=[
                ToolCall(
                    tool="analyze_descriptive",
                    args={
                        # `data` references t2 (retrieve_data rows), NOT t1 (check_data).
                        "data": "${t2}",
                        # column refs are the retrieve_data output aliases.
                        "column_ids": ["revenue", "quantity"],
                    },
                )
            ],
            expected_output="summary_stats",
            success_criteria=(
                "Produced mean/median/std/quartiles for revenue and quantity, above/below "
                "the typical range."
            ),
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example G — top-N ranking.
# "Top 3 product categories by total revenue."
# Shows: top-N is ONE retrieve_data query — group by the entity, aggregate the
# measure with an alias, order by that alias, limit N. NEVER a bare
# order-by-measure + limit (that ranks raw rows, so the same entity can appear
# twice — observed in production: "top 3 models" returned one model twice).
# --------------------------------------------------------------------------- #

_EXAMPLE_G = TaskList(
    plan_id="example_g",
    goal_restated="Rank product categories by total revenue and return the top 3.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria=(
                "Produced the orders table schema; category and revenue columns "
                "are present."
            ),
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Aggregate revenue per category, rank descending, keep the top 3.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "sum",
                                    "column_id": "c_revenue",
                                    "alias": "total_revenue",
                                },
                            ],
                            "group_by": ["c_category"],
                            "order_by": [{"column_id": "total_revenue", "dir": "desc"}],
                            "limit": 3,
                        }
                    },
                )
            ],
            expected_output="top3_categories",
            success_criteria=(
                "Produced at most 3 rows, one distinct category each, ranked by "
                "total revenue."
            ),
            depends_on=["t1"],
            estimated_cost="low",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example H — infeasible question (see planner.md "When the catalog cannot
# answer"). "What is our customer churn rate?" against a sales catalog with no
# subscription/churn data: no task list is forced onto unrelated columns;
# instead `infeasible_reason` states the gap + the nearest available data.
# --------------------------------------------------------------------------- #

_EXAMPLE_H = TaskList(
    plan_id="example_h",
    goal_restated="Measure the customer churn rate.",
    assumptions=[],
    open_questions=[],
    tasks=[],
    infeasible_reason=(
        "The connected source has no churn or subscription-status data — the "
        "orders table only carries order-level category, revenue, quantity, and "
        "dates. Nearest available analyses: repeat-purchase behaviour or revenue "
        "per customer over time."
    ),
)


# --------------------------------------------------------------------------- #
# Example I — combine two measures per entity (KM-703).
# "Which category has both the highest revenue and the highest average order
# quantity?" Shows: each measure is computed in its OWN grouped retrieve_data
# task (a "${t<id>}" placeholder resolves to a task's LAST output, so the two
# retrievals must be separate tasks), then analyze_merge aligns them on the
# shared entity alias. The merged table answers "both A and B" questions that
# a single query cannot express.
# --------------------------------------------------------------------------- #

_EXAMPLE_I = TaskList(
    plan_id="example_i",
    goal_restated=(
        "Identify the product category with both the highest total revenue and the "
        "highest average order quantity."
    ),
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category, revenue, and quantity.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria=(
                "Produced the orders table schema; category, revenue, and quantity "
                "columns are present."
            ),
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Total revenue per category.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "sum",
                                    "column_id": "c_revenue",
                                    "alias": "total_revenue",
                                },
                            ],
                            "group_by": ["c_category"],
                        }
                    },
                )
            ],
            expected_output="revenue_per_category",
            success_criteria="Produced one total-revenue row per category.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
        Task(
            id="t3",
            stage="data_preparation",
            objective="Average order quantity per category.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "avg",
                                    "column_id": "c_quantity",
                                    "alias": "avg_quantity",
                                },
                            ],
                            "group_by": ["c_category"],
                        }
                    },
                )
            ],
            expected_output="quantity_per_category",
            success_criteria="Produced one average-quantity row per category.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
        Task(
            id="t4",
            stage="evaluation",
            objective="Align both measures per category to find the category leading on both.",
            tool_calls=[
                ToolCall(
                    tool="analyze_merge",
                    args={
                        "data": "${t2}",
                        "data_right": "${t3}",
                        "on": ["category"],
                    },
                )
            ],
            expected_output="combined_measures",
            success_criteria=(
                "Produced one row per category carrying both total_revenue and "
                "avg_quantity."
            ),
            depends_on=["t2", "t3"],
            estimated_cost="low",
        ),
    ],
)


EXAMPLES: list[tuple[str, TaskList]] = [
    ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
    ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
    ("Revenue dipped in Q1 — what happened?", _EXAMPLE_C),
    ("What is the average and total order value per region?", _EXAMPLE_D),
    ("Total revenue for the East and West regions, counting orders of at least 100.", _EXAMPLE_E),
    ("Give me the summary statistics for order revenue and quantity.", _EXAMPLE_F),
    ("Which 3 product categories have the best revenue performance?", _EXAMPLE_G),
    ("What is our customer churn rate?", _EXAMPLE_H),
    (
        "Which product category has both the highest revenue and the highest average "
        "order quantity?",
        _EXAMPLE_I,
    ),
]


def render_examples() -> str:
    """Render the few-shots as text for the planner prompt."""
    blocks: list[str] = []
    for i, (question, plan) in enumerate(EXAMPLES, start=1):
        blocks.append(
            f"## Example {i}\n\n"
            f"Question:\n{question}\n\n"
            f"TaskList:\n{plan.model_dump_json(indent=2)}"
        )
    return "\n\n".join(blocks)