"""Question enrichment (phase A4, E-SQL arXiv:2409.16751).

E-SQL reports ~+5 EA on challenging BIRD from a single extra LLM call that
rewrites the question into an explicit specification — conditions, steps and
schema names spelled out. The enriched text goes into the generate prompt
*in addition to* the original question (never instead: the original stays
authoritative, so an enrichment mistake cannot override it).

Default OFF — wired only when ``PipelineConfig.enrich_question`` is True and
``enrichment_provider`` is set. The call rides the ordinary provider stack,
so it is cached alongside generation calls and replays on reruns.
"""

from __future__ import annotations

import re

from nl_sql.agent.prompts import load_prompt
from nl_sql.llm.providers.base import GenerateRequest, LLMProvider

_ENRICH_MAX_TOKENS = 1024
_MAX_ENRICHED_CHARS = 2000
_FENCE_RE = re.compile(r"```\w*\s*|\s*```")


def clean_enriched_text(text: str) -> str:
    """Normalise the model reply into a short plain-text restatement.

    Strips code fences and whitespace; truncates runaway replies (the block
    is auxiliary — a wall of text would drown the actual question). Returns
    ``""`` for empty/whitespace replies so the caller can skip the block.
    """
    cleaned = _FENCE_RE.sub("", text or "").strip()
    if len(cleaned) > _MAX_ENRICHED_CHARS:
        cleaned = cleaned[:_MAX_ENRICHED_CHARS].rsplit("\n", 1)[0].strip()
    return cleaned


def enrich_question(
    provider: LLMProvider,
    *,
    question: str,
    schema_text: str,
) -> str:
    """One LLM call → explicit restatement of the question ("" on empty)."""
    prompt = load_prompt(
        "enrich_question",
        schema_block=schema_text,
        question=question,
    )
    response = provider.generate(
        GenerateRequest(prompt=prompt, max_tokens=_ENRICH_MAX_TOKENS, temperature=0.0)
    )
    return clean_enriched_text(response.text)