"""Regression coverage for error-guided NL2SQL plan recovery."""
from __future__ import annotations

import json

import pytest

from app.core.exceptions import UnsupportedRequestError
from app.nlp.client import LLMClient
from app.planning.planner import build_query_plan
from app.schema.registry import ColumnDescriptor, TableDescriptor


BASE = {
    "fields": [], "select_all": False, "record_positions": [], "filters": [],
    "having": [], "sort": None, "group_by": None, "aggregate_field": None,
    "distinct": False, "limit": None, "needs_clarification": False,
    "clarification_question": None, "unsupported": False, "unsupported_reason": None,
}


def plan(**changes):
    result = dict(BASE)
    result.update(changes)
    return result


class RecoveryLLM(LLMClient):
    def __init__(self, responses):
        self.responses = responses
        self.prompts: list[str] = []

    async def generate_json(self, system_prompt, user_prompt):
        self.prompts.append(user_prompt)
        return json.dumps(self.responses[len(self.prompts) - 1])


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("question", "first", "second", "expected_sql"),
    [
        # Invented aggregate: the correction must not turn an unsupported
        # report request into a made-up COUNT query.
        (
            "show me executive summary data",
            plan(operation="COUNT", entity="invoices"),
            plan(operation="LIST", entity="invoices", unsupported=True,
                 unsupported_reason="Executive summary data is not available in this schema."),
            None,
        ),
        (
            "show invoices with total above 50000",
            plan(operation="LIST", entity="invoices", fields=["id", "total"]),
            plan(operation="LIST", entity="invoices", fields=["id", "total"],
                 filters=[{"field": "total", "operator": "GREATER_THAN", "value": 50000}]),
            "invoices.total >",
        ),
        (
            "show spaceships",
            plan(operation="LIST", entity="spaceships"),
            plan(operation="LIST", entity="spaceships", unsupported=True,
                 unsupported_reason="Spaceship records are not available."),
            None,
        ),
        (
            "show invoices from last month",
            plan(operation="LIST", entity="invoices", fields=["id", "total"]),
            plan(operation="LIST", entity="invoices", fields=["id", "total"],
                 filters=[{"field": "date", "operator": "EQUALS", "value": "last month"}]),
            "invoices.date BETWEEN",
        ),
        (
            "show expense records",
            plan(operation="LIST", entity="zb_expenses", fields=["amount"]),
            plan(operation="LIST", entity="expense_entries", fields=["amount"]),
            "FROM expense_entries",
        ),
        (
            "count invoices by customer",
            plan(operation="GROUP", entity="invoices"),
            plan(operation="GROUP", entity="invoices", group_by={"fields": ["customer_id"]}),
            "GROUP BY invoices.customer_id",
        ),
    ],
    ids=["invented-aggregate", "dropped-comparison", "invented-entity", "missing-date", "provider-table-without-platform", "missing-group-by"],
)
async def test_invalid_plan_gets_one_structured_recovery_and_revalidates(
    fake_snapshot, monkeypatch, question, first, second, expected_sql
):
    # Disable local reconciliation for this focused recovery-loop test: each
    # first response must reach the semantic/schema validator unchanged.
    monkeypatch.setattr("app.planning.planner.reconcile_plan", lambda _q, raw, *_args, **_kw: raw)
    monkeypatch.setattr("app.planning.planner._build_schema_backed_simple_plan", lambda *_args, **_kw: None)
    fake_snapshot.tables["zb_expenses"] = TableDescriptor(
        table="zb_expenses", description="Provider expense records", columns={
            "id": ColumnDescriptor("id", "integer", False, True, "id"),
            "amount": ColumnDescriptor("amount", "decimal", True, False, "amount"),
        },
    )
    llm = RecoveryLLM([first, second])

    if expected_sql is None:
        with pytest.raises(UnsupportedRequestError):
            await build_query_plan(question, fake_snapshot, llm)
    else:
        result = await build_query_plan(question, fake_snapshot, llm)
        assert expected_sql in result.data_query.sql

    assert len(llm.prompts) == 2  # initial plan + exactly one correction
    recovery_prompt = llm.prompts[1]
    assert question in recovery_prompt  # original wording is unchanged
    assert "VALIDATION_ERRORS" in recovery_prompt
    assert "QUERY_PLAN_VALIDATION_FAILED" in recovery_prompt
    assert "corrects EVERY" in recovery_prompt