"""Candidate narrowing for the Pick LLM call.

The composer covers L1's schema-pollution problem — the LLM never sees raw
SQL. But the LLM still needs to choose from *some* candidate set when
picking a metric / dimension / filter column. Sending the full schema for
every question is wasteful (cost, latency) and noisier (more candidates
the model can pick wrong from).

v1: **structural neighborhood**. Given the candidate metrics, compute the
union of:
  - columns of each candidate metric's `from_table`
  - columns of every table reachable within `max_hops` via declared
    relationships.

This covers every column that could legitimately group/filter that metric
without hand-curating a list. Cheap, deterministic, easy to reason about.

v2 (future): embedding-based retrieval against question + metric/column
descriptions, with usage-derived signal (pgvector). Out of scope here.
"""

from __future__ import annotations

from dataclasses import dataclass

from api.recon.types import Metric, Recon


@dataclass
class Candidates:
    metrics: list[Metric]
    # column_refs: each entry is either a bare name (when unambiguous) or
    # `table.column` (when the name lives on more than one reachable table).
    column_refs: list[str]


def neighborhood_tables(recon: Recon, seeds: list[str], max_hops: int = 2) -> set[str]:
    """BFS from `seeds` over recon.relationships, returning every table
    reachable within `max_hops`. Edges are undirected (FKs go both ways)."""
    adj: dict[str, set[str]] = {t: set() for t in recon.tables}
    for r in recon.relationships:
        adj.setdefault(r.from_table, set()).add(r.to_table)
        adj.setdefault(r.to_table, set()).add(r.from_table)
    visited: set[str] = set()
    frontier: set[str] = {s for s in seeds if s in recon.tables}
    for _ in range(max_hops + 1):  # +1 to include the seed itself at hop 0
        visited |= frontier
        next_frontier: set[str] = set()
        for t in frontier:
            next_frontier |= adj.get(t, set())
        frontier = next_frontier - visited
        if not frontier:
            break
    return visited


def narrow_candidates(
    recon: Recon,
    *,
    allowed_metrics: list[str] | None = None,
    max_hops: int = 2,
) -> Candidates:
    """Return the Pick-candidate set for the given metrics.

    `allowed_metrics`: if provided, restricts the candidate metrics to this
    list (intersected with recon.metrics). If None, every metric is in scope.
    `max_hops`: how far to walk from each candidate metric's from_table.
    """
    metric_names = (
        [m for m in allowed_metrics if m in recon.metrics]
        if allowed_metrics is not None
        else list(recon.metrics)
    )
    metrics = [recon.metrics[m] for m in metric_names]
    seeds = [m.from_table for m in metrics]
    tables = neighborhood_tables(recon, seeds, max_hops=max_hops)

    # Build the column ref list. A column name appearing in more than one
    # reachable table is exposed as `table.column` for each owner to keep
    # the LLM's pick unambiguous.
    appearances: dict[str, list[str]] = {}
    for tname in tables:
        for c in recon.tables[tname].columns:
            appearances.setdefault(c.name, []).append(tname)

    refs: list[str] = []
    for col, owners in appearances.items():
        if len(owners) == 1:
            refs.append(col)
        else:
            for t in owners:
                refs.append(f"{t}.{col}")
    refs.sort()

    return Candidates(metrics=metrics, column_refs=refs)