"""
Contract tests for the LLM query-plan JSON schema.

`app.planning.models.RawQueryPlan` is the single source of truth for the
LLM query-plan contract. `planning.prompts.query_plan_prompt.QUERY_PLAN_JSON_SCHEMA`
is generated from it and shared verbatim by both contract surfaces:

  * the Ollama structured-output `format` (`nlp.client`), and
  * the planner system prompt (`planning.prompts.query_plan_prompt`).

These tests pin that relationship so a hand-maintained schema can never
silently drift from the Pydantic model again: if the model gains or loses a
field, the prompt schema and the Ollama format change with it automatically.
"""
from __future__ import annotations

import json

from app.planning.prompts.query_plan_prompt import QUERY_PLAN_JSON_SCHEMA, SYSTEM_PROMPT
from app.planning.models import (
    AdditionalAggregateSpec,
    FilterSpec,
    GroupSpec,
    RawQueryPlan,
    SortSpec,
)
from app.planning.operators import FilterOperator, Operation


def test_schema_is_generated_from_raw_query_plan() -> None:
    assert QUERY_PLAN_JSON_SCHEMA == RawQueryPlan.model_json_schema()
    assert QUERY_PLAN_JSON_SCHEMA["title"] == "RawQueryPlan"


def test_system_prompt_embeds_the_exact_generated_schema() -> None:
    # The prompt must carry the serialized *generated* schema verbatim - not a
    # hand-maintained copy - so the two can never drift apart.
    assert json.dumps(QUERY_PLAN_JSON_SCHEMA, indent=2) in SYSTEM_PROMPT


def test_schema_covers_every_raw_query_plan_field() -> None:
    properties = QUERY_PLAN_JSON_SCHEMA["properties"]
    assert set(properties) == set(RawQueryPlan.model_fields)
    # The multi-intent / additional-aggregate contract must be represented.
    assert "additional_aggregates" in properties
    # Only the fields without defaults are mandatory - anything else may be
    # omitted and the model's default applies.
    assert set(QUERY_PLAN_JSON_SCHEMA["required"]) == {"operation", "entity"}


def test_schema_covers_every_nested_spec_field() -> None:
    defs = QUERY_PLAN_JSON_SCHEMA["$defs"]
    assert set(defs["SortSpec"]["properties"]) == set(SortSpec.model_fields)
    assert set(defs["FilterSpec"]["properties"]) == set(FilterSpec.model_fields)
    assert set(defs["GroupSpec"]["properties"]) == set(GroupSpec.model_fields)
    assert set(defs["AdditionalAggregateSpec"]["properties"]) == set(AdditionalAggregateSpec.model_fields)
    assert set(defs["AdditionalAggregateSpec"]["required"]) == {"operation"}


def test_full_featured_plan_round_trips_through_schema() -> None:
    plan = RawQueryPlan(
        operation=Operation.SUM,
        entity="invoices",
        fields=["vendor_id"],
        select_all=False,
        record_positions=[2, 5],
        filters=[
            FilterSpec(
                field="date",
                operator=FilterOperator.BETWEEN,
                value="2024-01-01",
                value2="2024-03-31",
            ),
            FilterSpec(field="status", operator=FilterOperator.IN, value=["paid", "open"]),
        ],
        having=[FilterSpec(field="result", operator=FilterOperator.GREATER_THAN, value=0)],
        sort=SortSpec(field="total", direction="desc"),
        group_by=GroupSpec(fields=["vendor_id"]),
        aggregate_field="total",
        additional_aggregates=[AdditionalAggregateSpec(operation=Operation.AVG, field="total")],
        distinct=False,
        limit=5,
    )
    payload = plan.model_dump()
    # Every emitted key is declared by the schema...
    assert set(payload) <= set(QUERY_PLAN_JSON_SCHEMA["properties"])
    # ...and the serialized payload survives a JSON round trip back through the
    # model the schema was generated from (all multi-intent fields included).
    assert RawQueryPlan.model_validate(json.loads(json.dumps(payload))) == plan


def test_plan_with_additional_aggregates_is_accepted_end_to_end() -> None:
    # The previously-undocumented multi-aggregate contract: parse a raw LLM
    # payload that uses additional_aggregates exactly as the schema describes.
    payload = {
        "operation": "SUM",
        "entity": "invoices",
        "aggregate_field": "total",
        "additional_aggregates": [
            {"operation": "AVG", "field": "total"},
            {"operation": "COUNT", "field": None},
        ],
    }
    plan = RawQueryPlan.model_validate(payload)
    assert plan.operation == Operation.SUM
    assert [(s.operation, s.field) for s in plan.additional_aggregates] == [
        (Operation.AVG, "total"),
        (Operation.COUNT, None),
    ]