"""The v2 compute strategy: additive state reads (m9 WP4, opt-in).

``IncrementalBackend`` matches :class:`RecomputeBackend`'s ``load_cutoff``
interface but reads per-unit cumulative moments from ``_ab_unit_state``
(the WP3 STATE stage's materialization) instead of re-scanning the raw fact
table — cumulative-intervals.md §4's committed v1 strategy, read side:
a cutoff costs (one additive SUM over closed-day state rows) + (for sub-day
cutoffs, a fact scan of at most the current-day tail, §6.4) instead of the
O(window) full rescan.

Correctness posture (m9 §0.2 — the load-bearing safety net):

- **Any gap falls back to full recompute, never a silent undercount.** The
  WP3 contiguity invariant makes detection one comparison: every day
  ``<= get_last_state_day()`` is materialized (a trailing day with zero
  qualifying events keeps ``get_last_state_day`` below it — the reader then
  conservatively falls back until a non-empty day advances the series); days
  past it are absent, not stale. If the last state day trails the cutoff's
  last closed day, this backend delegates the WHOLE cutoff to the wrapped
  ``RecomputeBackend`` and surfaces one warning per metric per run.
- **A non-finite tail contribution falls back too**: the closed-day writer
  rejects non-finite moments (``StateMomentError``), so the summed state is
  finite by construction, but the live tail render can carry warehouse NULLs
  (NaN after the loader) whose full-window SQL aggregation would have
  skipped them — recompute is the only faithful answer there.
- **Documented limitation (the R1 review's backfill finding, adjudicated as
  the m8 copy-mode precedent):** an event backfilled into an already-
  materialized day LATER than ``data_lag`` is frozen in day state — the gap
  check sees a contiguous series and cannot see content staleness. This is
  the §6.2 contract ("``data_lag`` bounds how wrong FROZEN points can be")
  applied to day state: the recompute default self-heals such events
  prospectively, the incremental path does not. Detection is WP5's
  ``verify-incremental`` whole-series reconciliation; recovery is ``abk run
  --full-refresh --from/--to`` (truncate-then-advance re-materializes, then
  recomputes). Out-of-SLA exposure backdating freezes the same way.
- **The reshaped result reuses the recompute containers unchanged**: the
  per-unit cumulative totals feed the SAME ``MetricLoadResult`` shape
  ``load_metric`` produces, so ``analyze_cutoff``/``build_container`` and
  ``SufficientStats.from_sample``'s stable reduction stay the single
  numerical path for the arm-level statistic (m9 WP4 step 2e).

The arm split happens at read time (state rows are arm-agnostic by design,
§5.2): tail units take the variant of the LIVE tail render's own cohort join
— exactly the arm the full-window recompute would land the unit's whole
value on — and state-only units join this run's cohort mapping (the LOAD
snapshot in direct mode, the persisted ``_ab_exposures`` in copy mode; both
via the driver-supplied loader). A state unit MISSING from that map (it
enrolled between LOAD and the STATE render — the live source the render
joins is fresher than the snapshot) triggers ONE quiet re-read of the live
source (``variant_map_refresh``, an R1+R2 review fix: never a stale-snapshot
drop, and never a routine second validation query either). A unit absent
from the refreshed cohort is dropped, mirroring the recompute INNER JOIN.

The CUPED pre-period covariate is untouched (m9 WP4 step 3): it loads
through the wrapped backend's one ``preperiod_covariate`` cache, so the
metric's own value load is the only thing this class replaces.
"""

from __future__ import annotations

from collections.abc import Callable
from datetime import date, datetime, timedelta
from pathlib import Path
from zoneinfo import ZoneInfo

import numpy as np

from abkit.compute.recompute_backend import RecomputeBackend
from abkit.config.experiment_config import ComparisonConfig, ExperimentConfig
from abkit.config.metric_config import MetricConfig
from abkit.core.period_planner import Cutoff, Grid, local_date, tz_midnight_utc
from abkit.database.internal_tables import InternalTablesManager
from abkit.database.manager import BaseDatabaseManager
from abkit.loaders.exposure_source import ExposureSnapshot, load_variant_map
from abkit.loaders.metric_loader import MetricLoadResult
from abkit.loaders.query_template import RenderWindow

#: metric type -> {role: moment column} — how per-unit cumulative moments
#: reshape into the loader's role arrays (the inverse of state_loader's
#: day_moments mapping; sum_value_sq and friends are day-level squares and
#: deliberately unused — the cross-unit variance comes from the per-unit
#: totals through SufficientStats, never from summed day squares)
_ROLE_MOMENTS: dict[str, dict[str, str]] = {
    "sample": {"value": "sum_value"},
    "fraction": {"count": "sum_value", "nobs": "n"},
    "ratio": {"numerator": "sum_value", "denominator": "sum_denominator"},
}


def build_incremental_backend(
    manager: BaseDatabaseManager,
    tables: InternalTablesManager,
    experiment: ExperimentConfig,
    recompute: RecomputeBackend,
    snapshot: ExposureSnapshot | None,
    grid: Grid,
    project_root: Path | None = None,
    on_warning: Callable[[str], None] | None = None,
    on_fallback: Callable[[str, str, datetime], None] | None = None,
) -> IncrementalBackend:
    """THE one way to construct the incremental reader (m9 WP5).

    Both readers — the driver's COMPUTE loop and `abk verify-incremental`'s
    reconciliation — go through here, so the command that certifies the
    incremental path certifies exactly the backend the pipeline runs (the
    same "one switch" discipline m8's ``build_cohort_backend`` established
    for cohort resolution). It wires the cohort map for the active mode:

    - **copy mode**: the persisted ``_ab_exposures`` (what the renders join);
      no refresh loader — the copy is this run's frozen cohort by design.
    - **direct mode**: the FREE validated LOAD snapshot, plus the
      refresh-on-miss re-read of the live source for a unit enrolled between
      LOAD and the STATE render.
    """
    copy_enabled = experiment.assignment.cohort_copy.enabled

    def _variant_map() -> dict[str, str]:
        if copy_enabled:
            return tables.get_exposure_variant_map(experiment.name)
        if snapshot is None:  # pragma: no cover - a caller contract violation
            raise ValueError(
                "direct-mode incremental reads need the run's validated cohort "
                "snapshot (build_cohort_backend(..., with_snapshot=True))"
            )
        return {str(unit): variant for unit, (variant, _, _) in snapshot.by_unit.items()}

    def _refresh_variant_map() -> dict[str, str]:
        return load_variant_map(manager, experiment, project_root, grid)

    return IncrementalBackend(
        tables,
        recompute,
        experiment,
        variant_map_loader=_variant_map,
        variant_map_refresh=None if copy_enabled else _refresh_variant_map,
        project_root=project_root,
        on_warning=on_warning,
        on_fallback=on_fallback,
    )


class IncrementalBackend:
    """Loads one comparison's data per cutoff from materialized day state.

    Constructed once per run by the driver alongside the factory-built
    ``RecomputeBackend`` it wraps (fallback + tail renders + covariate cache
    all go through that ONE instance, so the m8 cohort-mode threading is
    inherited, never reimplemented — m9 §0.2). ``variant_map_loader`` is
    called lazily at most once; ``on_warning`` receives the fallback
    disclosures the driver surfaces in ``RunOutcome.warnings``.
    """

    def __init__(
        self,
        tables: InternalTablesManager,
        recompute: RecomputeBackend,
        experiment: ExperimentConfig,
        variant_map_loader: Callable[[], dict[str, str]],
        variant_map_refresh: Callable[[], dict[str, str]] | None = None,
        project_root: Path | None = None,
        on_warning: Callable[[str], None] | None = None,
        on_fallback: Callable[[str, str, datetime], None] | None = None,
    ) -> None:
        self._tables = tables
        self._recompute = recompute
        self._experiment = experiment
        self._variant_map_loader = variant_map_loader
        self._variant_map_refresh = variant_map_refresh
        self._map_refreshed = False
        self._project_root = project_root
        self._on_warning = on_warning or (lambda _: None)
        # Fires on EVERY fallback (metric, kind, cutoff end_ts) — undeduped,
        # unlike the user-facing warning. `abk verify-incremental` (m9 WP5)
        # needs per-cutoff resolution: a cutoff served by the fallback was
        # NOT actually verified, and reporting it as "pass" would be a lie.
        self._on_fallback = on_fallback or (lambda _m, _k, _t: None)
        self._zone = ZoneInfo(experiment.timezone)
        self._series_keys: dict[str, tuple[str, str]] = {}
        self._last_state_day: dict[tuple[str, str], date | None] = {}
        # One entry per series: (required_last, the RAW per-unit moments) —
        # every same-day cutoff of a sub-day grid reuses one state read
        # instead of re-aggregating the full history per look (an R1 review
        # fix). Safe under the run lock: nothing writes the series during
        # COMPUTE, and the cached dict is never mutated (totals are rebuilt
        # per cutoff).
        self._cumulative_cache: dict[tuple[str, str], tuple[date, dict[str, dict[str, float]]]] = {}
        self._variant_map: dict[str, str] | None = None
        self._warned: set[tuple[str, str]] = set()

    # -- delegation: provenance render + covariate stay the recompute path --

    def render(self, metric_sql: str, window: RenderWindow) -> str:
        """The provenance copy of the FULL-WINDOW SQL (what recompute runs).

        Persisted rows keep the full-window render either way: it documents
        the window/cohort semantics of the number, not the read strategy.
        """
        return self._recompute.render(metric_sql, window)

    # -- the incremental read ----------------------------------------------

    def _series_key(self, metric: MetricConfig, metric_sql: str) -> tuple[str, str]:
        # Deferred like the loaders→compute edge (m8 precedent): the pipeline
        # package imports this module, so a module-level `pipeline.state`
        # import makes `abkit.compute.*` unimportable as an entry point.
        from abkit.pipeline.state import state_series_key

        key = self._series_keys.get(metric.name)
        if key is None:
            key = state_series_key(self._experiment, metric, metric_sql, self._project_root)
            self._series_keys[metric.name] = key
        return key

    def _cached_last_state_day(self, key: tuple[str, str]) -> date | None:
        # Stable within one run: STATE (when selected) runs before COMPUTE
        # under the run lock, and nothing else writes the series.
        if key not in self._last_state_day:
            self._last_state_day[key] = self._tables.get_last_state_day(*key)
        return self._last_state_day[key]

    def _cohort_variants(self) -> dict[str, str]:
        if self._variant_map is None:
            self._variant_map = self._variant_map_loader()
        return self._variant_map

    def _warn_once(self, metric_name: str, kind: str, message: str) -> None:
        # Deduped per (metric, reason-kind) — a benign early gap warning must
        # not swallow a later, operationally distinct disclosure (e.g. a
        # non-finite tail — a real data-quality signal; an R1 review fix).
        if (metric_name, kind) not in self._warned:
            self._warned.add((metric_name, kind))
            self._on_warning(message)

    def _fallback(
        self,
        comparison: ComparisonConfig,
        metric: MetricConfig,
        metric_sql: str,
        grid: Grid,
        cutoff: Cutoff,
        kind: str,
        reason: str,
    ) -> MetricLoadResult:
        self._warn_once(
            metric.name,
            kind,
            f"{self._experiment.name}/{metric.name}: incremental read fell back "
            f"to full recompute — {reason}",
        )
        self._on_fallback(metric.name, kind, cutoff.end_ts)
        return self._recompute.load_cutoff(comparison, metric, metric_sql, grid, cutoff)

    def load_cutoff(
        self,
        comparison: ComparisonConfig,
        metric: MetricConfig,
        metric_sql: str,
        grid: Grid,
        cutoff: Cutoff,
    ) -> MetricLoadResult:
        """Load one (comparison, cutoff) from state + at most one day of tail."""
        end_ts = cutoff.end_ts
        cutoff_day = local_date(end_ts, self._zone)
        last_midnight = tz_midnight_utc(cutoff_day, self._zone)
        required_last = cutoff_day - timedelta(days=1)
        role_moments = _ROLE_MOMENTS[metric.type]

        totals: dict[str, dict[str, float]] = {}
        # The first materialized day is the GRID's opening local day, never
        # the raw config field: with a sub-day start_ts the field carries a
        # time-of-day, and `date >= datetime` is a TypeError, not a comparison.
        first_state_day = local_date(grid.start_ts, self._zone)
        if required_last >= first_state_day:
            key = self._series_key(metric, metric_sql)
            last_state = self._cached_last_state_day(key)
            if last_state is None or last_state < required_last:
                have = (
                    "no materialized days" if last_state is None else f"state through {last_state}"
                )
                return self._fallback(
                    comparison,
                    metric,
                    metric_sql,
                    grid,
                    cutoff,
                    "gap",
                    f"{have}, cutoffs need closed days through {required_last} "
                    "(run the STATE step to advance the series)",
                )
            cached = self._cumulative_cache.get(key)
            if cached is not None and cached[0] == required_last:
                moments = cached[1]
            else:
                moments = self._tables.per_unit_cumulative(
                    key[0], key[1], first_state_day, required_last
                )
                self._cumulative_cache[key] = (required_last, moments)
            for unit, unit_moments in moments.items():
                totals[unit] = {role: unit_moments[m] for role, m in role_moments.items()}

        tail_variant: dict[str, str] = {}
        # Clamped to the window's left edge, the mirror of the STATE writer's
        # opening-day clamp: with a sub-day start_ts a cutoff can land INSIDE
        # the opening local day, and an unclamped [midnight, end_ts) tail
        # would sum hours of pre-experiment facts that full recompute — which
        # renders [grid.start_ts, end_ts) — never sees.
        tail_start = max(last_midnight, grid.start_ts)
        if tail_start < end_ts:
            tail = self._recompute.load_window(metric, metric_sql, RenderWindow(tail_start, end_ts))
            for variant in tail.variants():
                units = tail.units_by_variant[variant]
                roles = tail.roles_by_variant[variant]
                for role in role_moments:
                    values = roles.get(role)
                    if values is None or (values.size and not np.isfinite(values).all()):
                        # NULL-bearing tail: full-window SQL would have
                        # skipped the NULLs inside its aggregates — only the
                        # recompute render answers that faithfully.
                        return self._fallback(
                            comparison,
                            metric,
                            metric_sql,
                            grid,
                            cutoff,
                            "tail",
                            f"the current-day tail render carries "
                            f"{'a missing' if values is None else 'non-finite'} "
                            f"'{role}' role",
                        )
                for index, unit in enumerate(units):
                    tail_variant[unit] = variant
                    bucket = totals.setdefault(unit, dict.fromkeys(role_moments, 0.0))
                    for role in role_moments:
                        bucket[role] += float(roles[role][index])

        loaded = self._reshape(metric, totals, tail_variant)

        # CUPED: the fixed pre-period covariate stays the recompute-side load
        # (one shared cache — m9 WP4 step 3); eligibility already excluded
        # explicit ``columns.covariate`` metrics, so only the lookback form
        # can reach here.
        lookback = comparison.method.covariate_lookback
        if lookback is not None and metric.columns.covariate is None:
            loaded.attach_covariate(
                self._recompute.preperiod_covariate(metric, metric_sql, lookback, grid)
            )
        return loaded

    def _reshape(
        self,
        metric: MetricConfig,
        totals: dict[str, dict[str, float]],
        tail_variant: dict[str, str],
    ) -> MetricLoadResult:
        """Per-unit cumulative totals -> the ``load_metric`` result shape.

        Tail units carry the live render's own arm; state-only units join the
        run's cohort mapping; unmapped units drop (the INNER JOIN mirror).
        Units sort per variant — the loader's canonical order (m3 D11), so
        downstream containers are order-identical to the recompute path.
        """
        variant_map = self._cohort_variants()
        # Refresh-on-miss (an R2 review fix): the initial map is the FREE
        # LOAD snapshot; a state unit missing from it means the live source
        # gained the unit after LOAD (the STATE render saw it) — pay ONE
        # quiet re-read then, and only then.
        #
        # ONE refresh is sufficient by construction (the R3 review's
        # per-instance-budget question): the whole run holds one lock and
        # STATE precedes COMPUTE, so a refresh taken here reads a source at
        # least as fresh as the render that WROTE every state row this run —
        # and rows from earlier runs are older still. Every unit that can
        # appear in `totals` is therefore in the refreshed map unless it has
        # since LEFT the cohort, and dropping those is exactly what the
        # recompute render's INNER JOIN does. Units enrolling later in
        # COMPUTE have no state rows at all; their tail values (if any) carry
        # the tail render's own arm.
        disagreeing: list[str] = []
        if self._variant_map_refresh is not None and not self._map_refreshed:
            if any(u not in variant_map and u not in tail_variant for u in totals):
                self._map_refreshed = True
                refreshed = self._variant_map_refresh()
                # The refresh REPLACES the snapshot, so a unit whose arm moved
                # would silently start agreeing with the tail and the drift
                # disclosure below would never fire (an R2 review finding: one
                # legitimate mid-run enrollee masks every other unit's flip).
                # Diff the two maps here, while both are still in hand.
                disagreeing.extend(
                    unit
                    for unit, variant in variant_map.items()
                    if unit in refreshed and refreshed[unit] != variant
                )
                self._variant_map = variant_map = refreshed
        units_of: dict[str, list[str]] = {}
        for unit in totals:
            snapshot_variant = variant_map.get(unit)
            live_variant = tail_variant.get(unit)
            if live_variant is not None and snapshot_variant is not None:
                if live_variant != snapshot_variant:
                    disagreeing.append(unit)
            variant = live_variant or snapshot_variant
            if variant is None:
                continue
            units_of.setdefault(variant, []).append(unit)
        if disagreeing:
            # An already-exposed unit changed arm between this run's cohort
            # snapshot and the tail render. That is out of contract — the
            # assignment source is an append-shaped exposure log deduped by
            # MIN(exposure_ts), so a unit's arm is immutable — and it is
            # exactly the shape whose IN-render form m8 hard-errors on. The
            # live arm wins (never the stale snapshot), but the event is
            # surfaced rather than resolved silently: it means the two read
            # paths sampled the cohort at different instants and their numbers
            # may legitimately differ (an R1 exit-gate finding, adjudicated
            # SPLIT — the divergence is real, the trigger is out of contract).
            self._warn_once(
                metric.name,
                "variant-drift",
                f"{self._experiment.name}/{metric.name}: {len(disagreeing)} unit(s) "
                f"changed variant mid-run (e.g. '{sorted(disagreeing)[0]}') — the live "
                "render's arm was used; re-check the assignment source and run "
                "`abk verify-incremental`",
            )

        roles = list(_ROLE_MOMENTS[metric.type])
        units_by_variant: dict[str, np.ndarray] = {}
        roles_by_variant: dict[str, dict[str, np.ndarray]] = {}
        strata_by_variant: dict[str, np.ndarray | None] = {}
        for variant, units in units_of.items():
            units.sort()
            units_by_variant[variant] = np.array(units, dtype=object)
            roles_by_variant[variant] = {
                role: np.array([totals[u][role] for u in units], dtype=np.float64) for role in roles
            }
            strata_by_variant[variant] = None

        return MetricLoadResult(
            metric=metric.name,
            units_by_variant=units_by_variant,
            roles_by_variant=roles_by_variant,
            strata_by_variant=strata_by_variant,
        )
