"""Value lists of the catalog's small categorical columns, for resolving scope.

Go's catalog keeps 3 sample values per column, so "COAL LOADER" and site codes like
"SDJ" are often invisible to the planner, and the planner resolved an invisible focus
value on 0 of 3 gate runs (2026-10-29) while visible ones passed. Code therefore reads
the full value list of each SMALL string column — once per process per catalog —
and matches the objective and the user's words against it.

Read-only and through the normal query pipeline: one `retrieve_data` GROUP BY per
(table, column), validated and executed like any planner query (all five defense
layers apply), never SQL. Capped, time-boxed and never-throw. The result says whether
the load was COMPLETE: a focus resolved from a partial index is not pinned, and a
partial index is not cached (adversarial review 2026-10-29: a timed-out site list
otherwise pinned a focus without the site for the whole analysis). When Go emits full
value lists for low-cardinality columns (handoff), this becomes a catalog read.
"""

from __future__ import annotations

import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import Any

from src.middlewares.logging import get_logger

from ...catalog.models import Catalog, Column, Source, Table

logger = get_logger("scope_values")

#: A column is listed only when the catalog says it has at most this many values.
MAX_DISTINCT = 25
#: Tables queried per column name (largest first): `section` sits on 15 tables with
#: 4-6 values each; three cover the union without 15 queries.
TABLES_PER_COLUMN = 3
MAX_QUERIES = 12
#: Measured on the production source (2026-10-29): 9.7 s for the first query on a
#: cold connection, ~2 s each after. The load runs alongside the planner call
#: (5-20 s), so this cap rarely adds wait.
TIMEOUT_S = 20.0
#: Per-process cache keyed by the catalog's sources and `updated_at`; the TTL bounds
#: staleness of rows Go does not re-catalog. Only complete loads are cached.
CACHE_TTL_S = 900.0

Invoke = Callable[[str, dict[str, Any]], Awaitable[Any]]


class ValueIndex(dict):
    """`{column name: [values]}` plus how it was loaded.

    `complete`: every queried column answered in time. `sources`: lower-cased column
    name → the source ids whose lists were read, so a same-named column in another
    source (a file's own `section`) is not scoped by values it never held.
    """

    def __init__(self, *args: Any, complete: bool = True, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.complete = complete
        self.sources: dict[str, set[str]] = {}


_CACHE: dict[tuple, tuple[float, ValueIndex]] = {}


def _small(col: Column) -> bool:
    dc = col.stats.distinct_count if col.stats else None
    return (
        col.data_type == "string"
        and not col.pii_flag
        and dc is not None
        and 0 < dc <= MAX_DISTINCT
    )


def _targets(catalog: Catalog) -> list[tuple[Source, Table, Column]]:
    """Every column NAME's best table first, then second-best tables, … — so the cap
    can never starve a column of its only query (review: 4 columns on 3 tables each
    used all 12 queries and `site_code` was never read)."""
    by_name: dict[str, list[tuple[Source, Table, Column]]] = {}
    for s in catalog.sources:
        if s.source_type not in ("schema", "tabular"):
            continue
        for t in s.tables:
            for c in t.columns:
                if _small(c):
                    by_name.setdefault(c.name.lower(), []).append((s, t, c))
    for items in by_name.values():
        items.sort(
            key=lambda x: (x[2].stats.distinct_count or 0, x[1].row_count or 0),
            reverse=True,
        )
    out: list[tuple[Source, Table, Column]] = []
    for rank in range(TABLES_PER_COLUMN):
        out.extend(items[rank] for items in by_name.values() if len(items) > rank)
    return out[:MAX_QUERIES]


async def _one(invoke: Invoke, s: Source, t: Table, c: Column) -> list[str] | None:
    """The column's values, or None when the query failed."""
    ir = {
        "source_id": s.source_id,
        "table_id": t.table_id,
        "select": [{"kind": "column", "column_id": c.column_id, "alias": c.name}],
        "group_by": [c.column_id],
        "limit": MAX_DISTINCT * 2,
    }
    out = await invoke("retrieve_data", {"ir": ir})
    if getattr(out, "kind", None) != "table":
        return None
    rows = getattr(out, "rows", None) or []
    return [str(r[0]) for r in rows if r and isinstance(r[0], str) and r[0].strip()]


def _cache_key(catalog: Catalog) -> tuple:
    return (
        catalog.user_id,
        tuple(sorted((s.source_id, str(s.updated_at)) for s in catalog.sources)),
    )


async def load_value_index(catalog: Catalog, invoke: Invoke) -> ValueIndex:
    """The value lists of the small string columns. Never raises."""
    try:
        key = _cache_key(catalog)
        hit = _CACHE.get(key)
        if hit is not None and time.monotonic() - hit[0] < CACHE_TTL_S:
            return hit[1]
        index = await _load(catalog, invoke)
        if index and index.complete:  # a partial or failed load is never cached
            _CACHE[key] = (time.monotonic(), index)
        return index
    except Exception as e:
        logger.warning(
            "scope value index failed", degraded_seam="scope_values", error=repr(e)
        )
        return ValueIndex(complete=False)


async def _load(catalog: Catalog, invoke: Invoke) -> ValueIndex:
    targets = _targets(catalog)
    if not targets:
        return ValueIndex()
    tasks = [asyncio.create_task(_one(invoke, s, t, c)) for s, t, c in targets]
    done, pending = await asyncio.wait(tasks, timeout=TIMEOUT_S)
    for p in pending:
        p.cancel()
    index = ValueIndex(complete=not pending)
    for task, (s, _t, c) in zip(tasks, targets, strict=True):
        if task not in done:
            continue
        if task.exception() is not None:
            logger.info("scope value list failed", error=repr(task.exception()))
            index.complete = False
            continue
        values = task.result()
        if values is None:
            index.complete = False
            continue
        bucket = index.setdefault(c.name, [])
        for v in values:
            if v not in bucket:
                bucket.append(v)
        index.sources.setdefault(c.name.lower(), set()).add(s.source_id)
    for name in [k for k, v in index.items() if not v]:
        del index[name]
    logger.info(
        "scope value index loaded",
        columns=len(index),
        queries=len(targets),
        pending=len(pending),
        complete=index.complete,
        values=sum(len(v) for v in index.values()),
    )
    return index