"""drill_down helpers — one function per concern.

Kept separate from `analysis.py` so the DrillDown class stays readable as
high-level orchestration: decide → slice → loop → interpret.
"""

from __future__ import annotations

import json
import logging
import re
from typing import Any

from api import langfuse_client as lf
from api.analyses.drill_down.types import Slice
from api.composer import Pick, compose
from api.llm import chat
from api.prompts import load, render
from api.recon import load_recon
from api.runtime import events
from api.tools.execute_sql import execute_sql

logger = logging.getLogger("nvi.analyses.drill_down.helpers")


# ── Decision step ──

async def decide_next(question: str, metric: str, dimensions: list[str],
                      slices: list[Slice], budget: int) -> dict[str, Any]:
    """One LLM call to pick the next dimension or stop.

    If the LLM picks something not in the candidate list, drop the choice
    and stop — the planner's fallback (if any) takes over. We don't try to
    coerce the LLM into a valid dimension because it might be hallucinating
    a name (e.g. a metric name) that has no obvious mapping.
    """
    system = load("drill_down.next.system")
    user = render(
        "drill_down.next.user",
        question=question,
        metric=metric,
        dimensions=", ".join(dimensions),
        history=format_history(slices),
        budget=budget,
    )
    with lf.span("drill_down.next", input={"budget": budget}) as span:
        await events.publish_current(events.llm_call(
            "drill_down.next", system_len=len(system), user_len=len(user),
        ))
        decision = _parse_json(chat(
            system=system, user=user, max_tokens=256, span_name="drill_down.next.gen",
        ))

        # Hard guard: the chosen dimension MUST be in the candidate list.
        if decision.get("action") == "drill":
            dim = decision.get("dimension")
            if dim not in dimensions:
                logger.warning(
                    "drill_down.next picked %r (not in candidates %s); stopping",
                    dim, dimensions,
                )
                decision = {
                    "action": "stop",
                    "reason": f"LLM picked {dim!r}, which isn't in the candidate dimensions",
                }
        span.update(output=decision)
        return decision


# ── Slice execution ──

async def execute_slice(question: str, metric: str, dim: str, reason: str) -> Slice:
    """Compose SQL for one slice deterministically (metric + dim are already
    chosen by `decide_next`), execute it, emit tool-call events around both,
    and return the Slice. Zero LLM calls — recon authors the SQL."""
    recon = load_recon()
    pick = Pick(kind="aggregate", metric=metric, group_by=[dim], limit=10)

    await events.publish_current(events.tool_call_start(
        "compose_sql", input={"pick": pick.to_dict()},
    ))
    composed = compose(pick, recon)
    await events.publish_current(events.tool_call_end(
        "compose_sql", output={"sql": composed.sql, "tables": composed.used_tables},
    ))

    await events.publish_current(events.tool_call_start(
        "execute_sql", input={"sql": composed.sql, "dimension": dim},
    ))
    result = execute_sql(composed.sql)
    await events.publish_current(events.tool_call_end(
        "execute_sql",
        output={"dimension": dim, "row_count": result.row_count,
                "preview": result.as_dicts()[:5]},
    ))

    return Slice(
        dimension=dim,
        sql=composed.sql,
        rows=result.as_dicts()[:20],
        reason=reason,
    )


# ── Final interpretation ──

async def interpret(question: str, metric: str, slices: list[Slice]) -> str:
    system = load("drill_down.interpret.system")
    user = render(
        "drill_down.interpret.user",
        question=question,
        metric=metric,
        slices_block=format_slices_block(slices),
    )
    await events.publish_current(events.llm_call(
        "drill_down.interpret", system_len=len(system), user_len=len(user),
    ))
    return chat(system=system, user=user, max_tokens=512, span_name="drill_down.interpret")


# ── Prompt-context formatters ──

def format_history(slices: list[Slice]) -> str:
    """Render the 'already tried' block fed to drill_down.next."""
    if not slices:
        return "(none yet)"
    return "\n".join(f"- {s.dimension}: {repr(s.rows[:5])}" for s in slices)


def format_slices_block(slices: list[Slice]) -> str:
    """Render every slice's rows for drill_down.interpret."""
    return "\n\n".join(
        f"dimension: {s.dimension}\nrows: {repr(s.rows)}"
        for s in slices
    )


# ── JSON extraction ──

def _parse_json(text: str) -> dict[str, Any]:
    m = re.search(r"```(?:json)?\s*(\{.*\})\s*```", text, re.DOTALL)
    raw = m.group(1) if m else text
    start, end = raw.find("{"), raw.rfind("}")
    if start < 0 or end <= start:
        raise ValueError(f"drill_down.next returned no JSON: {text[:200]!r}")
    return json.loads(raw[start:end + 1])