"""Dynamic, data-shape-driven dashboard synthesizer.

Given one or more SQL result sets, this module *profiles* each row collection
and emits a variable number of UI components — KPIs, bars, lines, pies,
treemaps, histograms, scatter, heatmaps — chosen by the column dtypes,
cardinality, and temporal hints. Nothing is hardcoded per question; the same
synthesizer reacts purely to the data.

Used by both the offline planner and the LLM-backed planner so every answer
becomes a true mini-dashboard that adapts as the underlying database changes.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from statistics import mean, median
from typing import Any, Iterable

from ..config import get_settings
from ..retrieval.duck import run_sql
from ..ui_plan.schema import (
    Citation,
    Component,
    Encoding,
    GridCell,
    Layout,
    Persona,
    UIPlan,
)

ACCENT = "#F47C2C"
PALETTE = ["#F47C2C", "#FFC72C", "#FB9A2E", "#D6611B", "#FFB84D", "#A8470F", "#E63946", "#FFD27A"]


# --------------------------------------------------------------------------- #
# Schema introspection                                                        #
# --------------------------------------------------------------------------- #

def _schema() -> dict[str, dict]:
    p = Path(get_settings().schema_path)
    if not p.exists():
        return {}
    try:
        return json.loads(p.read_text())
    except Exception:  # noqa: BLE001
        return {}


def _table_cols(table: str) -> list[dict[str, str]]:
    sch = _schema().get(table) or {}
    return list(sch.get("columns") or [])


def _col_names(table: str) -> list[str]:
    return [c.get("name", "") for c in _table_cols(table)]


def _has_col(table: str, name: str) -> bool:
    return name in _col_names(table)


# --------------------------------------------------------------------------- #
# Column profiling                                                            #
# --------------------------------------------------------------------------- #

_DATE_RX = re.compile(r"^\d{4}-\d{2}(-\d{2})?")
_TIMEY_NAMES = {"month", "period", "date", "year", "fy", "quarter", "week"}


@dataclass
class ColProfile:
    name: str
    kind: str = "text"           # numeric|categorical|temporal|id|text
    n_unique: int = 0
    n_null: int = 0
    sample: list[Any] = field(default_factory=list)
    total: float = 0.0
    vmin: float | None = None
    vmax: float | None = None


def _is_number(v: Any) -> bool:
    return isinstance(v, (int, float)) and not isinstance(v, bool)


def _looks_temporal(name: str, vals: list[Any]) -> bool:
    n = name.lower()
    if any(t in n for t in _TIMEY_NAMES):
        return True
    if vals and isinstance(vals[0], str) and _DATE_RX.match(vals[0]):
        return True
    return False


def profile(rows: list[dict[str, Any]]) -> dict[str, ColProfile]:
    if not rows:
        return {}
    cols = list(rows[0].keys())
    out: dict[str, ColProfile] = {c: ColProfile(name=c) for c in cols}
    for c in cols:
        vals = [r.get(c) for r in rows]
        nonnull = [v for v in vals if v is not None]
        out[c].n_null = len(vals) - len(nonnull)
        out[c].sample = nonnull[:8]
        out[c].n_unique = len({str(v) for v in nonnull})
        nums = [float(v) for v in nonnull if _is_number(v)]
        if nums and len(nums) >= max(1, int(len(nonnull) * 0.7)):
            out[c].kind = "numeric"
            out[c].total = sum(nums)
            out[c].vmin = min(nums)
            out[c].vmax = max(nums)
        elif _looks_temporal(c, nonnull):
            out[c].kind = "temporal"
        elif "id" == c.lower() or c.lower().endswith("_id"):
            out[c].kind = "id"
        else:
            out[c].kind = "categorical" if out[c].n_unique <= max(40, len(rows) // 2) else "text"
    return out


# --------------------------------------------------------------------------- #
# Result-set → components                                                     #
# --------------------------------------------------------------------------- #

def _kpi_card(idx: int, label: str, value: Any, sub: str = "") -> Component:
    fmt = "currency_inr" if isinstance(value, (int, float)) else "text"
    return Component(
        id=f"k{idx}",
        type="kpi",
        title=label,
        subtitle=sub,
        data={"value": value, "format": fmt, "label": label},
        why=None,
    )


def _bar(idx: int, title: str, rows: list[dict], x: str, y: str, why: str) -> Component:
    return Component(
        id=f"c{idx}",
        type="bar",
        title=title,
        data=rows,
        encoding=Encoding(x=x, y=y, color=ACCENT),
        why=why,
    )


def _stacked_bar(idx: int, title: str, rows: list[dict], x: str,
                 y_cols: list[str], why: str) -> Component:
    """Stacked bar when multiple numeric measures share one categorical axis."""
    return Component(
        id=f"sb{idx}", type="stacked-bar", title=title, data=rows[:60],
        encoding=Encoding(x=x, y=y_cols, color=ACCENT),
        why=why,
    )


def _sankey(idx: int, title: str, rows: list[dict], source: str,
            target: str, value_col: str, why: str) -> Component | None:
    """Sankey for flow between two categorical columns weighted by a numeric."""
    links: list[dict[str, Any]] = []
    for r in rows[:100]:
        s = str(r.get(source, ""))[:24] or "(blank)"
        t = str(r.get(target, ""))[:24] or "(blank)"
        v = float(r.get(value_col) or 0)
        if v > 0 and s != t:
            links.append({"source": s, "target": t, "value": round(v, 2)})
    if len(links) < 2:
        return None
    nodes = list({l["source"] for l in links} | {l["target"] for l in links})
    return Component(
        id=f"sk{idx}", type="sankey", title=title,
        data={"nodes": [{"name": n} for n in nodes], "links": links},
        why=why,
    )


def _line(idx: int, title: str, rows: list[dict], x: str, y: str, why: str) -> Component:
    return Component(
        id=f"c{idx}",
        type="line",
        title=title,
        data=rows,
        encoding=Encoding(x=x, y=y, color=ACCENT),
        why=why,
    )


def _treemap(idx: int, title: str, rows: list[dict], name_col: str, value_col: str, why: str) -> Component:
    series = [
        {"name": str(r.get(name_col, "—"))[:32] or "—",
         "value": float(r.get(value_col) or 0)}
        for r in rows[:30]
    ]
    return Component(
        id=f"c{idx}",
        type="treemap",
        title=title,
        data=series,
        why=why,
    )


def _histogram(idx: int, title: str, rows: list[dict], col: str) -> Component | None:
    nums = [float(r.get(col) or 0) for r in rows if _is_number(r.get(col))]
    if len(nums) < 8:
        return None
    nums.sort()
    lo, hi = nums[0], nums[-1]
    if hi <= lo:
        return None
    nbins = min(20, max(5, int(len(nums) ** 0.5)))
    width = (hi - lo) / nbins
    buckets = [0] * nbins
    for v in nums:
        b = min(nbins - 1, int((v - lo) / width))
        buckets[b] += 1
    data = [
        {"bucket": f"{lo + i * width:.0f}–{lo + (i + 1) * width:.0f}", "count": n}
        for i, n in enumerate(buckets)
    ]
    return Component(
        id=f"h{idx}",
        type="bar",
        title=title,
        data=data,
        encoding=Encoding(x="bucket", y="count", color=ACCENT),
        why=f"Distribution of {col} across {len(nums)} rows in {nbins} buckets.",
    )


def _heatmap_corr(idx: int, title: str, rows: list[dict], num_cols: list[str]) -> Component | None:
    if len(num_cols) < 2 or len(rows) < 5:
        return None
    cols = num_cols[:6]
    series = [[float(r.get(c) or 0) for c in cols] for r in rows[:300]]
    means = [mean(col) for col in zip(*series)]
    stds = [
        (sum((v - m) ** 2 for v in col) / max(1, len(col) - 1)) ** 0.5
        for col, m in zip(zip(*series), means)
    ]
    data: list[list[float]] = []
    x_labels, y_labels = list(cols), list(cols)
    for i, ci in enumerate(cols):
        for j, cj in enumerate(cols):
            if stds[i] == 0 or stds[j] == 0:
                corr = 0.0
            else:
                vi = [r[i] for r in series]
                vj = [r[j] for r in series]
                cov = sum((a - means[i]) * (b - means[j]) for a, b in zip(vi, vj)) / max(1, len(vi) - 1)
                corr = cov / (stds[i] * stds[j])
            data.append([float(j), float(i), round(corr, 3)])
    return Component(
        id=f"hm{idx}",
        type="heatmap",
        title=title,
        data=data,  # type: ignore[arg-type]
        options={"x": x_labels, "y": y_labels, "min": -1, "max": 1},
        why="Pearson correlation across the numeric columns in this result.",
    )


def _table(idx: int, title: str, rows: list[dict]) -> Component:
    return Component(id=f"t{idx}", type="data-table", title=title, data=rows[:200], why=None)


def _raw_sql(idx: int, sql: str) -> Component:
    return Component(id=f"sql{idx}", type="raw-sql", title="SQL", data={"sql": sql}, why=None)


def _donut(idx: int, title: str, rows: list[dict], name_col: str,
           value_col: str, why: str) -> Component:
    """Render as sunburst (frontend supports sunburst, renders like a donut for flat series)."""
    data = [
        {"name": str(r.get(name_col, "—"))[:28] or "—",
         "value": float(r.get(value_col) or 0)}
        for r in rows if float(r.get(value_col) or 0) > 0
    ][:8]
    return Component(
        id=f"dn{idx}", type="sunburst", title=title, data=data, why=why,
    )


def _sunburst_hier(idx: int, title: str, rows: list[dict], outer: str, inner: str,
                   value_col: str) -> Component | None:
    """2-level sunburst: outer category contains inner category → value."""
    tree: dict[str, dict[str, float]] = {}
    for r in rows:
        o = str(r.get(outer, "—") or "—")[:24]
        i = str(r.get(inner, "—") or "—")[:24]
        v = float(r.get(value_col) or 0)
        if v <= 0:
            continue
        tree.setdefault(o, {}).setdefault(i, 0.0)
        tree[o][i] += v
    if not tree:
        return None
    data = [
        {"name": o, "children": [{"name": k, "value": round(v, 2)}
                                  for k, v in sorted(kids.items(),
                                                     key=lambda x: x[1], reverse=True)[:8]]}
        for o, kids in list(sorted(tree.items(),
                                    key=lambda x: sum(x[1].values()), reverse=True))[:10]
    ]
    return Component(
        id=f"sb{idx}", type="sunburst", title=title, data=data,
        why=f"Hierarchical share: {outer} → {inner} on {value_col}.",
    )


def _funnel(idx: int, title: str, rows: list[dict], name_col: str,
            value_col: str, why: str) -> Component:
    data = sorted(
        [
            {"name": str(r.get(name_col, "—"))[:28] or "—",
             "value": float(r.get(value_col) or 0)}
            for r in rows if float(r.get(value_col) or 0) > 0
        ],
        key=lambda d: d["value"],
        reverse=True,
    )[:8]
    return Component(id=f"fn{idx}", type="funnel", title=title, data=data, why=why)


def _radar(idx: int, title: str, rows: list[dict], name_col: str,
           value_cols: list[str], why: str) -> Component | None:
    if not rows or not value_cols:
        return None
    top = rows[:5]
    maxes = {c: max((float(r.get(c) or 0) for r in rows), default=1.0) or 1.0 for c in value_cols}
    indicators = [{"name": c, "max": round(maxes[c] * 1.15, 2)} for c in value_cols]
    series = [
        {"name": str(r.get(name_col, f"row {i+1}"))[:24],
         "value": [round(float(r.get(c) or 0), 2) for c in value_cols]}
        for i, r in enumerate(top)
    ]
    return Component(
        id=f"rd{idx}", type="radar", title=title, data=series,
        options={"indicators": indicators}, why=why,
    )


def _waterfall(idx: int, title: str, rows: list[dict], name_col: str,
               budget_col: str, actual_col: str) -> Component | None:
    """Variance waterfall: budget − actual per category (as bar with signed deltas)."""
    data = []
    for r in rows[:12]:
        b = float(r.get(budget_col) or 0)
        a = float(r.get(actual_col) or 0)
        data.append({
            "name": str(r.get(name_col, "—"))[:24] or "—",
            "variance": round(b - a, 2),
            "budget": round(b, 2),
            "actual": round(a, 2),
        })
    if not data:
        return None
    return Component(
        id=f"wf{idx}", type="bar", title=title, data=data,
        encoding=Encoding(x="name", y=["budget", "actual"], color=ACCENT),
        why="Budget vs actual side-by-side reveals under/overspend per category.",
    )


def _scatter(idx: int, title: str, rows: list[dict], x_col: str, y_col: str,
             cat_col: str | None, why: str) -> Component:
    """Scatter plot for two numeric dimensions, optionally colored by a categorical."""
    data = []
    for r in rows[:300]:
        xv = r.get(x_col)
        yv = r.get(y_col)
        if xv is None or yv is None:
            continue
        pt: dict[str, Any] = {"x": float(xv), "y": float(yv)}
        if cat_col:
            pt["category"] = str(r.get(cat_col, ""))[:24]
        data.append(pt)
    enc = Encoding(x=x_col, y=y_col, color=cat_col or ACCENT)
    return Component(id=f"sc{idx}", type="scatter", title=title, data=data,
                     encoding=enc, why=why)


def _pie(idx: int, title: str, rows: list[dict], name_col: str,
         value_col: str, why: str) -> Component:
    """True pie chart for low-cardinality share-of-total."""
    data = [
        {"name": str(r.get(name_col, "—"))[:24] or "—",
         "value": float(r.get(value_col) or 0)}
        for r in rows if float(r.get(value_col) or 0) > 0
    ][:8]
    return Component(id=f"pi{idx}", type="pie", title=title, data=data, why=why)


def _area(idx: int, title: str, rows: list[dict], x: str, y: str, why: str) -> Component:
    """Area chart — temporal continuous fill."""
    return Component(
        id=f"ar{idx}", type="area", title=title, data=rows[:200],
        encoding=Encoding(x=x, y=y, color=ACCENT), why=why,
    )


def _gauge(idx: int, title: str, value: float, max_val: float, why: str) -> Component:
    """Gauge for utilization / completion rates."""
    return Component(
        id=f"ga{idx}", type="gauge", title=title,
        data={"value": round(value, 2), "max": round(max_val, 2), "label": title},
        why=why,
    )


def _pick_kpis(seq: int, rows: list[dict], prof: dict[str, ColProfile],
               num_cols: list[str], cat_cols: list[str],
               time_cols: list[str] | None = None) -> list[Component]:
    """Data-tier-aware KPI selection — 1 to 4 KPIs with optional sparklines."""
    out: list[Component] = []
    if not num_cols:
        return out
    primary = max(num_cols, key=lambda c: prof[c].total)
    p = prof[primary]
    nrows = len(rows)

    # Build sparkline trend from time-ordered values (up to 12 points)
    sparkline: list[float] | None = None
    if time_cols and nrows >= 3:
        sparkline = [float(r.get(primary) or 0) for r in rows[:12]]

    kpi1 = _kpi_card(seq * 10 + 1, f"Σ {primary}", round(p.total, 2),
                     f"{nrows} rows")
    if sparkline:
        kpi1.options = {"sparkline": sparkline}
    out.append(kpi1)

    # median / average tells a different story than sum
    if nrows >= 4:
        nums = sorted(float(r.get(primary) or 0) for r in rows)
        med = nums[len(nums) // 2]
        kpi2 = _kpi_card(seq * 10 + 2, f"median {primary}", round(med, 2),
                         f"typical value")
        out.append(kpi2)
    # top category + its share of total — but only if the value is meaningful
    if cat_cols and p.total > 0:
        # Filter out junk category values (NULL, empty, "grand total", "(unspecified)" etc.)
        _JUNK = {"", "—", "(unspecified)", "grand total", "grand total:",
                 "total", "totals", "subtotal", "n/a", "na", "null", "none",
                 "all", "overall"}
        clean_rows = [r for r in rows
                      if str(r.get(cat_cols[0]) or "").strip().lower() not in _JUNK
                      and r.get(cat_cols[0]) is not None]
        if clean_rows:
            top = max(clean_rows, key=lambda r: float(r.get(primary) or 0))
            tv = float(top.get(primary) or 0)
            if tv > 0:
                share_pct = (tv / p.total) * 100 if p.total > 0 else 0
                out.append(_kpi_card(
                    seq * 10 + 3,
                    f"top {cat_cols[0]}",
                    str(top.get(cat_cols[0]) or "—")[:28],
                    f"₹{tv:,.0f} · {share_pct:.0f}% of total",
                ))
    # spread (range) — only for large/diverse sets
    if nrows >= 10 and p.vmin is not None and p.vmax is not None:
        spread = p.vmax - p.vmin
        if spread > 0:
            out.append(_kpi_card(seq * 10 + 4, f"range {primary}", round(spread, 2),
                                 f"min ₹{p.vmin:,.0f} → max ₹{p.vmax:,.0f}"))
    return out


def _components_for_result(seq: int, title: str, sql: str, rows: list[dict],
                           emit_table: bool = True, lean: bool = False) -> list[Component]:
    """Emit a data-shape-driven component list.

    `lean` trims secondary result-sets so the full dashboard stays digestible.
    Component count varies by row count, cardinality, numeric/categorical mix
    and temporal hints — nothing is emitted just to fill space.
    """
    out: list[Component] = []
    if not rows:
        out.append(Component(id=f"n{seq}", type="callout",
                             title=title,
                             data={"text": "No rows returned for this slice."},
                             why=None))
        out.append(_raw_sql(seq, sql))
        return out

    # Check for all-null/all-zero numeric data — don't render charts for it
    if _all_zero_numeric(rows):
        out.append(Component(id=f"n{seq}", type="callout",
                             title=title,
                             data={"text": "All numeric values are zero or empty — "
                                   "this part of the query could not be answered "
                                   "from the available data."},
                             why=None))
        out.append(_raw_sql(seq, sql))
        return out

    prof = profile(rows)
    num_cols = [c for c, p in prof.items() if p.kind == "numeric"]
    cat_cols = [c for c, p in prof.items() if p.kind == "categorical"]
    time_cols = [c for c, p in prof.items() if p.kind == "temporal"]
    nrows = len(rows)

    # --- Tier (controls how rich the dashboard gets) -------------------------
    tier = "tiny" if nrows <= 3 else "small" if nrows <= 12 else \
           "medium" if nrows <= 60 else "large"

    # Detect budget/actual pair (waterfall candidate)
    b_col = next((c for c in num_cols
                  if "budget" in c.lower() and "total" in c.lower()), None)
    a_col = next((c for c in num_cols
                  if ("actual" in c.lower() and "total" in c.lower())
                  or c.lower() == "actuals"), None)
    has_variance = bool(b_col and a_col and cat_cols)

    # 1) KPI strip — only for primary result (lean=False); secondaries
    # skip KPIs to avoid duplicate Σ/median/top tiles cluttering the dashboard.
    if not lean:
        out.extend(_pick_kpis(seq, rows, prof, num_cols, cat_cols, time_cols))

    # 2) Temporal trace → line (always if time exists)
    if time_cols and num_cols:
        tc, nc = time_cols[0], num_cols[0]
        out.append(_line(seq * 10 + 5, f"{nc} over {tc}",
                         rows[:200], x=tc, y=nc,
                         why=f"Trace of {nc} along {tc}."))

    # 3) Primary categorical view — pick a shape from data characteristics
    if cat_cols and num_cols:
        pc, pn = cat_cols[0], num_cols[0]
        uniq = len({r.get(pc) for r in rows})
        ranked = sorted(rows, key=lambda r: float(r.get(pn) or 0), reverse=True)
        if tier == "tiny" or uniq <= 5:
            # tiny / low-cardinality → pie (clear share-of-total)
            out.append(_pie(seq * 10 + 6, f"{pc} share of {pn}",
                            ranked, name_col=pc, value_col=pn,
                            why="Pie reveals share-of-total with few categories."))
        elif uniq >= 20:
            out.append(_bar(seq * 10 + 6,
                            f"Top 15 {pc} by {pn}",
                            ranked[:15], x=pc, y=pn,
                            why=f"Ranks the 15 largest {pc} by {pn}."))
            if not lean:
                out.append(_treemap(seq * 10 + 7,
                                    f"{pc} share of {pn}",
                                    ranked[:30], name_col=pc, value_col=pn,
                                    why="Treemap reveals share-of-total across many categories."))
        else:
            out.append(_bar(seq * 10 + 6,
                            f"{pc} by {pn}",
                            ranked[:15], x=pc, y=pn,
                            why=f"Bar compares {pn} across {pc}."))

    # 4) Budget-vs-actual waterfall (only when both present)
    if has_variance and not lean:
        wf = _waterfall(seq * 10 + 8, f"Budget vs Actual by {cat_cols[0]}",
                        rows, name_col=cat_cols[0],
                        budget_col=b_col, actual_col=a_col)  # type: ignore[arg-type]
        if wf is not None:
            out.append(wf)

    # 5) Hierarchical sunburst when 2 cats + 1 num
    if len(cat_cols) >= 2 and num_cols and tier in {"medium", "large"}:
        sb = _sunburst_hier(seq * 10 + 9,
                            f"{cat_cols[0]} → {cat_cols[1]} share of {num_cols[0]}",
                            rows, outer=cat_cols[0], inner=cat_cols[1],
                            value_col=num_cols[0])
        if sb is not None:
            out.append(sb)

    # 6) Funnel for status/stage-like categorical
    status_cols = [c for c in cat_cols if any(k in c.lower()
                  for k in ("status", "stage", "state", "phase", "approval"))]
    if status_cols and num_cols:
        out.append(_funnel(seq * 10 + 10, f"{status_cols[0]} funnel",
                           rows, name_col=status_cols[0], value_col=num_cols[0],
                           why=f"Funnel ranks {status_cols[0]} by {num_cols[0]}."))

    # 7) Radar when 3+ numeric metrics across a few categories
    if len(num_cols) >= 3 and cat_cols and tier in {"small", "medium"} and not lean:
        rd = _radar(seq * 10 + 11, f"{cat_cols[0]} multi-metric profile",
                    rows, name_col=cat_cols[0], value_cols=num_cols[:5],
                    why="Radar compares multiple metrics for the top rows.")
        if rd is not None:
            out.append(rd)

    # 8) Distribution histogram — only for medium/large single-numeric slices
    if num_cols and tier in {"medium", "large"} and not lean:
        h = _histogram(seq * 10 + 12, f"Distribution of {num_cols[0]}", rows, num_cols[0])
        if h is not None:
            out.append(h)

    # 9) Correlation heatmap when genuine multi-numeric dataset
    if len(num_cols) >= 3 and tier == "large" and not lean:
        hm = _heatmap_corr(seq * 10 + 13, "Numeric correlations", rows, num_cols)
        if hm is not None:
            out.append(hm)

    # 10) Scatter when 2+ numeric cols (explore relationships)
    if len(num_cols) >= 2 and tier in {"medium", "large"} and not lean:
        x_sc, y_sc = num_cols[0], num_cols[1]
        cat_sc = cat_cols[0] if cat_cols else None
        out.append(_scatter(seq * 10 + 15, f"{x_sc} vs {y_sc}",
                            rows, x_col=x_sc, y_col=y_sc, cat_col=cat_sc,
                            why=f"Scatter reveals relationship between {x_sc} and {y_sc}."))

    # 11) Area chart as alternate temporal view (stacked feel)
    if time_cols and len(num_cols) >= 2 and not lean:
        tc, nc2 = time_cols[0], num_cols[1] if len(num_cols) > 1 else num_cols[0]
        out.append(_area(seq * 10 + 16, f"{nc2} trend",
                         rows[:200], x=tc, y=nc2,
                         why=f"Area fill highlights the trend of {nc2} over time."))

    # 12) Gauge for utilization / completion when budget+actual exist
    if has_variance and not lean:
        total_b = sum(float(r.get(b_col) or 0) for r in rows)  # type: ignore[arg-type]
        total_a = sum(float(r.get(a_col) or 0) for r in rows)  # type: ignore[arg-type]
        if total_b > 0:
            out.append(_gauge(seq * 10 + 17,
                              "Budget Utilisation",
                              total_a, total_b,
                              why=f"Utilisation gauge: ₹{total_a:,.0f} spent of ₹{total_b:,.0f} budget."))

    # 13) Stacked bar when 2+ numeric measures share one categorical axis
    if len(num_cols) >= 2 and cat_cols and tier in {"small", "medium"} and not lean:
        pc = cat_cols[0]
        y_names = num_cols[:3]  # up to 3 stacked measures
        ranked = sorted(rows, key=lambda r: sum(float(r.get(c) or 0) for c in y_names), reverse=True)
        out.append(_stacked_bar(
            seq * 10 + 18,
            f"{pc} — stacked {', '.join(y_names)}",
            ranked[:15], x=pc, y_cols=y_names,
            why=f"Stacked bar compares {', '.join(y_names)} side-by-side per {pc}.",
        ))

    # 14) Sankey flow when 2 categoricals + numeric (flow semantics)
    flow_hints = ("from", "to", "source", "target", "sender", "receiver",
                  "origin", "destination", "funder", "region")
    if len(cat_cols) >= 2 and num_cols and tier in {"medium", "large"} and not lean:
        c1, c2 = cat_cols[0], cat_cols[1]
        # Only emit sankey when columns suggest a flow relationship
        c1_flow = any(h in c1.lower() for h in flow_hints)
        c2_flow = any(h in c2.lower() for h in flow_hints)
        if c1_flow or c2_flow or len(cat_cols) >= 2:
            sk = _sankey(seq * 10 + 19,
                         f"{c1} → {c2} flow by {num_cols[0]}",
                         rows, source=c1, target=c2, value_col=num_cols[0],
                         why=f"Sankey traces fund flow from {c1} to {c2}.")
            if sk is not None:
                out.append(sk)

    # 15) Data-table preview + raw SQL (always, for transparency)
    if emit_table:
        out.append(_table(seq * 10 + 14, f"{title} — rows", rows))
    out.append(_raw_sql(seq, sql))
    return out
    return out


# --------------------------------------------------------------------------- #
# Result-quality helper (used by offline plan & component rendering)          #
# --------------------------------------------------------------------------- #

def _all_zero_numeric(rows: list[dict]) -> bool:
    """True when every numeric value across all rows is 0 or None."""
    has_num = False
    for r in rows:
        for v in r.values():
            if v is None:
                continue
            if isinstance(v, (int, float)):
                has_num = True
                if v != 0:
                    return False
            elif isinstance(v, str):
                try:
                    fv = float(v)
                    has_num = True
                    if fv != 0:
                        return False
                except (ValueError, OverflowError):
                    pass
    return has_num


# --------------------------------------------------------------------------- #
# Heuristic SQL generation (offline path)                                     #
# --------------------------------------------------------------------------- #

# Tables ranked by usefulness for keyword fallback
_PREFERRED_TABLES = ["bva", "extended_bva_v5", "expense_mb", "item_mb",
                     "open_po", "opening_budget_fy_25_26", "opening_underspents"]


def _pick_table(question: str) -> str:
    q = question.lower()
    # New analytical views (computed budget rollups)
    if ("ytd" in q or "year to date" in q or "balance" in q
            or ("summary" in q and ("budget" in q or "actual" in q))):
        return "SummaryBVA"
    if "finalbva" in q or ("rolled" in q and "budget" in q):
        return "finalbva5"
    if "open" in q and "po" in q or "purchase" in q or "vendor" in q:
        return "open_po"
    if "expense" in q:
        return "expense_mb"
    if "item" in q and "mb" in q:
        return "item_mb"
    if "underspent" in q or "underspend" in q:
        return "opening_underspents"
    if "opening" in q and "budget" in q:
        return "opening_budget_fy_25_26"
    if "extended" in q:
        return "extended_bva_v5"
    return "bva"


def _pick_dim(table: str, question: str) -> str | None:
    q = question.lower()
    cols = _col_names(table)
    cmap = {c.lower(): c for c in cols}
    priority = []
    for kw, cand in [
        ("funder", "funder_name"),
        ("region", "region"),
        ("region", "regions"),
        ("project", "project_name"),
        ("category", "category_name"),
        ("vendor", "vendor_name"),
        ("state", "state"),
        ("status", "approval_status"),
        ("status", "purchase_order_status"),
        ("vertical", "vertical"),
        ("program", "program_type"),
        ("district", "district"),
        ("department", "department"),
    ]:
        if kw in q and cand in cmap:
            priority.append(cmap[cand])
    if priority:
        return priority[0]
    for cand in ["funder_name", "region", "regions", "category_name",
                 "project_name", "vendor_name", "approval_status",
                 "vertical", "program_type", "state", "district"]:
        if cand in cmap:
            return cmap[cand]
    for c in _table_cols(table):
        n = c.get("name", "")
        if not n.lower().endswith("_id") and not _is_numeric_type(c.get("type")):
            return n
    return None


_NUMERIC_TYPES = {"BIGINT", "DOUBLE", "INTEGER", "DECIMAL", "FLOAT", "REAL"}
_AMOUNT_HINTS = ("amount", "total", "value", "price", "cost", "qty", "quantity", "actual", "budget")


def _is_numeric_type(t: str | None) -> bool:
    u = (t or "").upper()
    return any(k in u for k in _NUMERIC_TYPES)


def _measure_col(table: str) -> str:
    cols = _table_cols(table)
    cmap = {c.get("name", "").lower(): c.get("name", "") for c in cols}
    for cand in ["actual_total", "budget_total", "amount", "total_amount",
                 "total_cost_amount", "value", "qty", "quantity"]:
        if cand in cmap:
            return cmap[cand]
    for c in cols:
        if _is_numeric_type(c.get("type")):
            return c["name"]
    for c in cols:
        n = (c.get("name") or "").lower()
        if any(h in n for h in _AMOUNT_HINTS) and not n.startswith("actual_0"):
            return c["name"]
    return ""


def _has_monthly(table: str) -> bool:
    return any(c.startswith("actual_0") for c in _col_names(table))


def heuristic_primary_sql(question: str, table: str) -> tuple[str, str]:
    dim = _pick_dim(table, question)
    measure = _measure_col(table)
    if dim and measure:
        sql = (
            f"SELECT COALESCE(CAST({dim} AS VARCHAR), '(unspecified)') AS {dim}, "
            f"ROUND(SUM(TRY_CAST({measure} AS DOUBLE))::DOUBLE, 2) AS {measure} "
            f"FROM {table} "
            f"WHERE {dim} IS NOT NULL "
            f"GROUP BY 1 "
            f"ORDER BY 2 DESC NULLS LAST "
            f"LIMIT 25"
        )
        title = f"{dim} by {measure}"
    elif dim:
        sql = (
            f"SELECT COALESCE(CAST({dim} AS VARCHAR), '(unspecified)') AS {dim}, "
            f"COUNT(*) AS row_count "
            f"FROM {table} "
            f"GROUP BY 1 "
            f"ORDER BY 2 DESC NULLS LAST "
            f"LIMIT 25"
        )
        title = f"{dim} row counts"
    else:
        sql = f"SELECT * FROM {table} LIMIT 200"
        title = f"{table} preview"
    return title, sql


def heuristic_secondaries(question: str, table: str) -> list[tuple[str, str]]:
    """Build 1–3 enrichment queries based on what columns the table actually has."""
    out: list[tuple[str, str]] = []
    cols = _col_names(table)
    cmap = {c.lower(): c for c in cols}
    measure = _measure_col(table)
    region_col = cmap.get("region") or cmap.get("regions") or ""

    # Monthly trend (FY26) when monthly columns exist and we have a measure
    if measure and _has_monthly(table):
        months = [
            ("Apr-25", "actual_01042025"), ("May-25", "actual_01052025"),
            ("Jun-25", "actual_01062025"), ("Jul-25", "actual_01072025"),
            ("Aug-25", "actual_01082025"), ("Sep-25", "actual_01092025"),
            ("Oct-25", "actual_01102025"), ("Nov-25", "actual_01112025"),
            ("Dec-25", "actual_01122025"), ("Jan-26", "actual_01012026"),
            ("Feb-26", "actual_01022026"), ("Mar-26", "actual_01032026"),
        ]
        present = [(m, c) for m, c in months if c in cols]
        if present:
            union = " UNION ALL ".join(
                f"SELECT '{m}' AS month, {i:02d} AS ord, "
                f"ROUND(SUM(TRY_CAST({c} AS DOUBLE))::DOUBLE, 2) AS actual_inr FROM {table}"
                for i, (m, c) in enumerate(present, 1)
            )
            out.append((
                "Monthly burn — FY26",
                f"SELECT month, actual_inr FROM ({union}) ORDER BY ord"
            ))

    # Region breakdown
    if region_col and table != "open_po":
        if measure:
            out.append((
                "Spend by region",
                f"SELECT COALESCE({region_col},'(unspecified)') AS region, "
                f"ROUND(SUM(TRY_CAST({measure} AS DOUBLE))::DOUBLE, 2) AS {measure} "
                f"FROM {table} GROUP BY 1 ORDER BY 2 DESC NULLS LAST LIMIT 12"
            ))
        else:
            out.append((
                "Rows by region",
                f"SELECT COALESCE({region_col},'(unspecified)') AS region, "
                f"COUNT(*) AS row_count "
                f"FROM {table} GROUP BY 1 ORDER BY 2 DESC NULLS LAST LIMIT 12"
            ))
    elif region_col:  # open_po: count rows by region
        out.append((
            "Open POs by region",
            f"SELECT COALESCE({region_col},'(unspecified)') AS region, "
            f"COUNT(*) AS row_count FROM {table} GROUP BY 1 ORDER BY 2 DESC LIMIT 12"
        ))

    # Category mix
    if "category_name" in cmap and measure:
        cat_col = cmap["category_name"]
        out.append((
            "Top categories",
            f"SELECT {cat_col} AS category_name, "
            f"ROUND(SUM(TRY_CAST({measure} AS DOUBLE))::DOUBLE, 2) AS {measure} "
            f"FROM {table} WHERE {cat_col} IS NOT NULL "
            f"GROUP BY 1 ORDER BY 2 DESC NULLS LAST LIMIT 15"
        ))

    # Funder mix (if not already the dimension)
    if "funder_name" in cmap and not re.search(r"funder", question, re.I):
        f_col = cmap["funder_name"]
        if measure:
            out.append((
                "Top funders",
                f"SELECT {f_col} AS funder_name, "
                f"ROUND(SUM(TRY_CAST({measure} AS DOUBLE))::DOUBLE, 2) AS {measure} "
                f"FROM {table} WHERE {f_col} IS NOT NULL AND {f_col} <> 'NONE' "
                f"GROUP BY 1 ORDER BY 2 DESC NULLS LAST LIMIT 10"
            ))
        else:
            out.append((
                "Funder distribution",
                f"SELECT {f_col} AS funder_name, COUNT(*) AS row_count "
                f"FROM {table} WHERE {f_col} IS NOT NULL AND {f_col} <> 'NONE' "
                f"GROUP BY 1 ORDER BY 2 DESC LIMIT 10"
            ))

    # Approval status mix for open_po
    if "approval_status" in cmap:
        out.append((
            "PO approval status mix",
            f"SELECT COALESCE({cmap['approval_status']},'(unknown)') AS approval_status, "
            f"COUNT(*) AS row_count FROM {table} GROUP BY 1 ORDER BY 2 DESC LIMIT 10"
        ))

    return out[:3]


# --------------------------------------------------------------------------- #
# Public entry: assemble a UIPlan from result sets                            #
# --------------------------------------------------------------------------- #

@dataclass
class ResultSet:
    title: str
    sql: str
    rows: list[dict[str, Any]]
    table: str = ""


_PERSONA_PALETTE = {
    "executive":   ("#A8470F", "calm",          "spacious"),
    "analyst":     ("#F47C2C", "neutral",       "comfortable"),
    "investigator":("#E63946", "urgent",        "compact"),
    "storyteller": ("#FB9A2E", "celebratory",   "spacious"),
    "auditor":     ("#0EA5E9", "investigative", "compact"),
    "strategist":  ("#A855F7", "calm",          "comfortable"),
}


def infer_persona(question: str, intent: str = "exploratory") -> Persona:
    """Pick a persona from the question wording + intent."""
    q = question.lower()
    name = "analyst"
    if any(k in q for k in ("forecast", "strategy", "scenario", "future",
                             "what if", "5-year", "projection")):
        name = "strategist"
    elif any(k in q for k in ("over budget", "anomaly", "spike", "fraud",
                              "irregular", "alert", "violation", "risk")):
        name = "investigator"
    elif any(k in q for k in ("compliance", "audit", "ytd", "balance",
                              "reconcile", "validate", "verify")):
        name = "auditor"
    elif any(k in q for k in ("summary", "ceo", "executive", "headline",
                              "overview", "snapshot", "kpi")):
        name = "executive"
    elif any(k in q for k in ("story", "highlight", "showcase", "impact",
                              "achievement", "win", "celebrate")):
        name = "storyteller"
    elif intent == "comparative":
        name = "analyst"
    elif intent == "diagnostic":
        name = "investigator"
    accent, mood, density = _PERSONA_PALETTE[name]
    return Persona.model_validate({
        "name": name, "accent": accent, "mood": mood, "density": density,
    })


def assemble_plan(question: str, results: list[ResultSet],
                  intent: str = "exploratory") -> UIPlan:
    components: list[Component] = []
    citations: list[Citation] = []
    for i, rs in enumerate(results, 1):
        # Primary gets the full treatment; secondaries stay lean so the full
        # dashboard stays varied & digestible rather than repeating the same
        # KPI+bar+table+SQL quartet for every slice.
        comps = _components_for_result(i, rs.title, rs.sql, rs.rows,
                                       lean=(i > 1))
        components.extend(comps)
        citations.append(Citation(
            id=f"src{i}", source=rs.table or "duckdb",
            row_ids=list(range(len(rs.rows))),
            snippet=rs.sql[:280],
        ))

    # Build a flowing 12-col grid: KPIs share rows; other charts vary by type.
    # Chart-type → (preferred width, height) for visual variety
    _CHART_SIZE: dict[str, tuple[int, int]] = {
        "kpi":          (4,  2),
        "data-table":   (12, 5),
        "raw-sql":      (12, 3),
        "callout":      (12, 3),
        "narrative":    (12, 4),
        "line":         (8,  5),
        "area":         (8,  5),
        "bar":          (6,  5),
        "stacked-bar":  (8,  6),
        "pie":          (4,  5),
        "scatter":      (6,  5),
        "gauge":        (4,  4),
        "treemap":      (6,  6),
        "sunburst":     (6,  6),
        "sankey":       (12, 7),
        "heatmap":      (8,  6),
        "funnel":       (4,  5),
        "radar":        (6,  5),
        "waterfall":    (8,  5),
        "histogram":    (6,  5),
        "network":      (8,  6),
        "box":          (6,  5),
        "parallel":     (12, 6),
    }
    grid: list[GridCell] = []
    cursor_y = 0
    cursor_x = 0
    for c in components:
        w, h = _CHART_SIZE.get(c.type, (6, 5))
        if c.type == "kpi":
            if cursor_x + w > 12:
                cursor_x = 0
                cursor_y += h
            grid.append(GridCell(id=f"g_{c.id}", component=c.id, w=w, h=h,
                                 x=cursor_x, y=cursor_y))
            cursor_x += w
        else:
            if cursor_x != 0 and cursor_x + w > 12:
                cursor_y += 5
                cursor_x = 0
            grid.append(GridCell(id=f"g_{c.id}", component=c.id, w=w, h=h,
                                 x=cursor_x, y=cursor_y))
            cursor_x += w
            if cursor_x >= 12:
                cursor_y += h
                cursor_x = 0

    # Narrative summary built from results
    bits: list[str] = []
    for rs in results[:3]:
        if not rs.rows:
            continue
        prof = profile(rs.rows)
        nums = [c for c, p in prof.items() if p.kind == "numeric"]
        cats = [c for c, p in prof.items() if p.kind == "categorical"]
        if nums and cats:
            tot = sum(float(r.get(nums[0]) or 0) for r in rs.rows)
            if tot == 0:
                # Skip zero-total results from the narrative
                continue
            top = max(rs.rows, key=lambda r: float(r.get(nums[0]) or 0))
            bits.append(
                f"**{rs.title}** — {len(rs.rows)} rows, Σ{nums[0]} ≈ ₹{tot:,.0f}; "
                f"top {cats[0]}: *{top.get(cats[0])}*."
            )
        else:
            bits.append(f"**{rs.title}** — {len(rs.rows)} rows.")
    if not bits:
        narrative = (
            f"The query \"{question}\" could not be fully answered from the "
            f"available data. Some aspects may be outside the scope of the "
            f"current database. The results shown above are the closest "
            f"match available."
        )
    else:
        narrative = " ".join(bits)

    return UIPlan(
        persona=infer_persona(question, intent),
        layout=Layout(grid=grid),
        components=components,
        narrative=narrative,
        citations=citations,
    )


def dynamic_offline_plan(question: str) -> UIPlan:
    """Heuristic, schema-aware plan when no LLM keys are configured."""
    table = _pick_table(question)
    title_p, sql_p = heuristic_primary_sql(question, table)
    primary_rows = run_sql(sql_p).rows
    rs_primary = ResultSet(title=title_p, sql=sql_p, rows=primary_rows, table=table)
    extras = []
    for ti, sq in heuristic_secondaries(question, table):
        r = run_sql(sq).rows
        if r:
            extras.append(ResultSet(title=ti, sql=sq, rows=r, table=table))
    all_results = [rs_primary, *extras]
    # Filter out result sets with all-zero/null numeric data
    all_results = [rs for rs in all_results if rs.rows and not _all_zero_numeric(rs.rows)]
    if not all_results:
        # Nothing meaningful found — return a plan with a scope message
        all_results = [rs_primary]  # keep primary for SQL provenance
    return assemble_plan(question, all_results)