"""The STATE stage: per-(unit, day) moment materialization (m9 WP3).

The write-only half of cumulative-intervals.md §4's committed v1 strategy:
for every STATE-eligible metric, each not-yet-materialized closed local day
``[tz-midnight(d), tz-midnight(d+1))`` is rendered through the m8 cohort
factory backend (a single-day, non-cumulative window — never a hand-rolled
cohort join, m9 §0.2) and replaced into ``_ab_unit_state`` via the §5.2
replace-not-sum primitive. The opening day is clamped to ``grid.start_ts``
when the window starts mid-day (m10), and a trailing day the horizon would
truncate is not materialized at all — see :func:`closed_state_days`. No reader exists in this milestone — WP4's
``IncrementalBackend`` flips the read path.

Eligibility (per metric, within one experiment):

- non-stratified — no stratum dimension exists in the state key (m9 §8 Q2);
- referenced by at least one closed-form (unseeded) comparison — a
  bootstrap-only metric never pays the write cost (WP3 step 5);
- the SQL body does not reference ``ab_cov_*`` — such a render depends on
  the comparison's covariate window, so its day moments would not be
  comparison-independent; the metric stays on full-window recompute;
- no explicit ``columns.covariate`` role (an R2 review exclusion) — that
  role is documented as an author-computed column that may be a static
  per-unit snapshot, which is NOT additive across day renders: per-day
  ``sum_cov`` rows would inflate by the unit's active-day count once
  summed. Additivity cannot be verified from config, so such metrics stay
  on full-window recompute;
- **the metric DECLARES ``state_additive: true``** — the author's promise
  that per-day partials add up to the window total. This is a declaration
  because it cannot be read off the SQL: three review rounds each found a
  new textual shape that looks additive and is not (a dead staging CTE, an
  outer re-aggregation of an inner sum, a ``UNION`` branch that binds
  positionally, an identity ``sum()`` over a renamed per-unit ``max()``).
  The hazard is real — the scaffolded ``example_signup_cr`` projects
  ``max(signed_up)`` and a literal ``1 AS visits``, both of which inflate
  when eleven daily renders are summed;
- **and the projections do not visibly contradict that promise**:
  ``_role_projections_are_additive`` is a VETO-only sanity filter — every
  ``AS <role column>`` must alias exactly one additive aggregate call, with
  no ``DISTINCT``/``OVER``, and multi-branch SQL is refused outright. It can
  only take eligibility away, never grant it, so its remaining blind spots
  cost an unnoticed author mistake, not a silent corruption — and
  ``abk verify-incremental`` is the empirical oracle that catches those.

Contiguity invariant (the WP4 gap-detection contract): days advance strictly
in order and a failed day TRUNCATES the series from that day before the run
aborts the metric's advancement, so every day ``<= get_last_state_day()`` is
materialized (possibly legitimately empty — trailing empty days are
re-rendered each run, the accepted v1 cost) and days past it are absent, not
stale. Non-finite moments truncate from the failing day with a loud warning
— reads past the last valid day stay on full recompute, never a
partially-written or stale day the future reader would trust.
"""

from __future__ import annotations

import hashlib
import re
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo

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 Grid, local_date, tz_midnight_utc
from abkit.database.internal_tables import (
    InternalTablesManager,
    compute_metric_state_id,
    compute_state_source_id,
    normalize_sql_for_identity,
)
from abkit.loaders.query_template import RenderWindow
from abkit.loaders.state_loader import StateMomentError, day_moments
from abkit.stats import get_method_class

Logger = Callable[[str], None]


def _noop_log(_: str) -> None:  # pragma: no cover - trivial
    return None


@dataclass(frozen=True)
class StateDay:
    """One closed local day and its UTC render window."""

    day: date
    window: RenderWindow


@dataclass
class StateOutcome:
    """One experiment's STATE-stage summary."""

    days_materialized: int = 0
    rows_written: int = 0
    warnings: list[str] = field(default_factory=list)


def closed_state_days(
    experiment: ExperimentConfig, grid: Grid, watermark_ts: datetime
) -> list[StateDay]:
    """The closed local days of the grid window, in order.

    A day is closed iff its full window ``[tz-midnight(d), tz-midnight(d+1))``
    has passed the compute watermark (§6.4: the state stage advances only at
    day close; the same ``end_ts <= watermark_ts`` completeness rule the
    cutoff planner uses) AND lies inside the grid window. The two edges are
    treated asymmetrically, for the reason spelled out below.

    The enumeration is keyed off the GRID, not off ``experiment.start_ts``:
    the config field may carry a time-of-day (m10), in which case the opening
    local day is legitimately PARTIAL and its window starts at
    ``grid.start_ts`` — no state row may ever cover an instant outside the
    experiment window, or the additive read would count pre-experiment rows.
    That clamp is safe on the START side only because ``start_ts`` is folded
    into the series identity: moving it re-keys the series.

    A sub-day ``horizon_ts`` gets the opposite treatment — the trailing
    partial day is **not materialized at all**. Clamping it would freeze a
    part-day row that no later run re-renders (a materialized day is never
    revisited) while ``horizon_ts`` does NOT generally re-key the series, so
    extending the experiment would keep summing the truncated day. Nothing
    needs it: a cutoff reads closed days only through ``cutoff_day - 1``, and
    no cutoff exceeds the horizon, so the trailing day is served by the
    reader's own sub-day tail render.
    """
    zone = ZoneInfo(experiment.timezone)
    days: list[StateDay] = []
    day = local_date(grid.start_ts, zone)
    while True:
        window_start = max(tz_midnight_utc(day, zone), grid.start_ts)
        window_end = tz_midnight_utc(day + timedelta(days=1), zone)
        if window_start >= grid.horizon_ts or window_end > grid.horizon_ts:
            break
        if window_end > watermark_ts:
            break
        days.append(StateDay(day=day, window=RenderWindow(window_start, window_end)))
        day += timedelta(days=1)
    return days


def _needs_seed(method_name: str) -> bool:
    method_cls = get_method_class(method_name)
    return any(spec.name == "seed" for spec in method_cls.param_specs)


def _cohort_identity(experiment: ExperimentConfig, project_root: Path | None) -> dict[str, Any]:
    """The cohort-shaping config folded into the series identity (R1 fix).

    Everything here changes which units (or which day boundaries) the per-day
    render sees; an edit therefore orphans the series and re-materializes it
    under the current definition — the same self-consistency the full-window
    recompute gets for free by re-rendering the whole window every cutoff.
    """
    assignment_sql = experiment.assignment.get_query_text(project_root)
    identity: dict[str, Any] = {
        # the SAME literal-preserving normalization the metric SQL gets: a
        # whitespace edit INSIDE a quoted literal (`'Summer  Sale'`) changes
        # which units the cohort holds and must orphan the series
        "assignment_sql_sha256": hashlib.sha256(
            normalize_sql_for_identity(assignment_sql).encode("utf-8")
        ).hexdigest(),
        "added_filters": experiment.assignment.added_filters,
        "unit_key": experiment.unit_key,
        "variants": list(experiment.assignment.variants),
        "timezone": experiment.timezone,
        # str() of a date and of a midnight datetime differ ('2024-07-01' vs
        # '2024-07-01 00:00:00'), which is why the config layer preserves the
        # written type instead of normalizing: a re-parse must not re-hash.
        # `interval_anchor` is deliberately NOT folded in — it moves cutoffs,
        # never day boundaries, and a day's contents are what this key covers.
        "start_ts": str(experiment.start_ts),
    }
    # The assignment render's window ends at the GRID HORIZON, so a cohort
    # SQL referencing ab_end_date/ab_end_ts renders differently when
    # horizon_ts moves (R2 review) — fold it in for exactly those SQLs.
    # End-invariant assignment SQL (the common case) skips it, so the most
    # routine edit there is — extending an experiment — never orphans state.
    if "ab_end_" in assignment_sql:
        identity["horizon_ts"] = str(experiment.horizon_ts)
    return identity


def state_series_key(
    experiment: ExperimentConfig,
    metric: MetricConfig,
    metric_sql: str,
    project_root: Path | None = None,
) -> tuple[str, str]:
    """The ``(source_table, column_set_id)`` key of one metric's state series.

    THE identity function (WP3 steps 1-2, as corrected in review) — the WP4
    reader and every test must derive the key through here, never by
    composing the hash inputs by hand.
    """
    source_id = compute_state_source_id(experiment.name, metric.name)
    series_id = compute_metric_state_id(
        metric.columns.role_map(),
        metric_sql,
        cohort_config=_cohort_identity(experiment, project_root),
    )
    return source_id, series_id


#: aggregates whose per-day results SUM to the whole-window result. A positive
#: allowlist on purpose: an unrecognised projection makes a metric INELIGIBLE
#: (a missed optimisation), never eligible-and-wrong.
_ADDITIVE_AGGREGATES = ("sum", "sumif", "count", "countif")

#: an additive aggregate call and NOTHING else — anchored at both ends of the
#: aliased expression, so ``sum(x)/count(*)`` (a per-unit ratio) is rejected
#: even though it ends in an additive call. One level of paren nesting keeps
#: ``sum(if(x, 1, 0))`` readable as additive.
_ADDITIVE_CALL = r"(?:{aggs})\s*\((?:[^()]|\([^()]*\))*\)"

#: aggregates are not additive over a DISTINCT set (a per-day distinct count
#: does not sum to the window's), and a windowed aggregate is not a per-unit
#: value at all.
_NON_ADDITIVE_MODIFIERS = re.compile(r"\bdistinct\b|\bover\b", re.IGNORECASE)

#: a multi-branch query binds later branches to the output columns POSITIONALLY
#: — they need no alias at all, so an alias scan cannot see what they project
#: (an R2 review finding). Refuse the whole shape rather than guess.
_MULTI_BRANCH = re.compile(r"\bunion\b|\bintersect\b|\bexcept\b", re.IGNORECASE)

#: ``--`` / ``/* */`` comments and single-quoted literals: text that looks like
#: SQL but never executes. Stripped before any structural question is asked.
_SQL_NOISE = re.compile(r"--[^\n]*|/\*.*?\*/|'(?:[^']|'')*'", re.DOTALL)


def _strip_sql_noise(sql: str) -> str:
    return _SQL_NOISE.sub(" ", sql)


def _role_projections_are_additive(metric: MetricConfig, metric_sql: str) -> bool:
    """Is EVERY projection of every summed role column an additive aggregate?

    The STATE stage stores one row per (unit, DAY) and the reader SUMS those
    days, so a role column that is not additive across days is silently
    inflated — the project's own scaffolded fraction metric proves the
    hazard: ``max(signed_up) AS signed_up`` and ``1 AS visits`` both satisfy
    the one-row-per-unit contract, yet summing eleven daily renders gives
    eleven trials where the window has one (found by
    ``abk verify-incremental``, m9 WP5).

    Additivity cannot be *proven* from SQL text, so this asks the strictest
    cheap question instead: after stripping comments and string literals,
    **every** ``AS <role column>`` in the body must alias exactly one additive
    aggregate call, with no ``DISTINCT`` and no ``OVER``. Requiring *every*
    alias — not merely the existence of one — is what makes the check
    unfoolable by a dead CTE, a staging subquery whose result an outer
    ``max()`` re-aggregates, or a second projection of the same name (all
    three were review findings against the earlier "does an additive
    projection exist?" form of this check). A role column that is never aliased
    at all is likewise refused: we cannot see what produces it.

    It remains a NECESSARY, not sufficient, condition — SQL is not parsed
    here, and the body is the pre-Jinja source. The authoritative additivity
    oracle is ``abk verify-incremental`` (m9 WP5), which is precisely why
    ``compute.incremental_reads`` defaults off and the flip criteria
    (cumulative-intervals.md §4.1) demand clean reconciliation runs first.
    """
    body = _strip_sql_noise(metric_sql)
    if _MULTI_BRANCH.search(body):
        return False
    call = _ADDITIVE_CALL.format(aggs="|".join(_ADDITIVE_AGGREGATES))
    for role, column in metric.columns.role_map().items():
        if role in ("variant", "stratum"):
            continue
        aliases = list(re.finditer(rf"(?<![\w.])AS\s+{re.escape(column)}(?![\w])", body, re.I))
        if not aliases:
            return False
        for alias in aliases:
            expression = body[: alias.start()]
            match = re.search(rf"{call}\s*$", expression, flags=re.IGNORECASE | re.DOTALL)
            if match is None:
                return False
            # The aggregate must be the WHOLE aliased expression: anything
            # before it other than a select-list/paren boundary means the
            # alias names a compound expression (a ratio, a difference, a
            # CASE) whose value is not the aggregate itself.
            prefix = expression[: match.start()].rstrip()
            if prefix and prefix[-1] not in ",(":
                return False
            if _NON_ADDITIVE_MODIFIERS.search(match.group(0)):
                return False
    return True


def comparison_state_eligible(
    comparison: ComparisonConfig, metric: MetricConfig, metric_sql: str
) -> bool:
    """Is this (comparison, metric) on the additive-state contract?

    THE eligibility predicate (module docstring) — shared by the WP3 writer
    (below) and the WP4 read-path resolver in ``pipeline/driver.py``: a
    comparison the writer never materializes must never route to the
    incremental reader, and vice versa, so both sides ask one function.
    """
    return (
        metric.state_additive
        and not _needs_seed(comparison.method.name)
        and metric.columns.stratum is None
        and metric.columns.covariate is None
        and "ab_cov_" not in metric_sql
        and _role_projections_are_additive(metric, metric_sql)
    )


def state_eligible_metrics(
    experiment: ExperimentConfig,
    metrics_by_name: dict[str, MetricConfig],
    project_root: Path | None,
) -> list[tuple[MetricConfig, str]]:
    """The ``(metric, metric_sql)`` list this run materializes.

    One series per metric: the experiment config binds each metric to at
    most one comparison, so the comparison's method decides eligibility.
    """
    chosen: dict[str, tuple[MetricConfig, str]] = {}
    for comparison in experiment.comparisons:
        metric = metrics_by_name.get(comparison.metric)
        if metric is None:
            # A dangling metric reference. `abk run` cannot reach this (its
            # config lint fails first), but `abk clean` walks the whole
            # project WITHOUT that lint, and a half-finished rename must not
            # crash housekeeping (an R1 review finding). A comparison whose
            # metric is gone has no state to materialize either way.
            continue
        metric_sql = metric.get_query_text(project_root)
        if not comparison_state_eligible(comparison, metric, metric_sql):
            continue
        chosen[metric.name] = (metric, metric_sql)
    return list(chosen.values())


def materialize_state(
    tables: InternalTablesManager,
    experiment: ExperimentConfig,
    metrics_by_name: dict[str, MetricConfig],
    backend: RecomputeBackend,
    grid: Grid,
    watermark_ts: datetime,
    project_root: Path | None = None,
    full_refresh_window: tuple[datetime, datetime] | None = None,
    force_rebuild: bool = False,
    metric_filter: str | None = None,
    log: Logger = _noop_log,
) -> StateOutcome:
    """Materialize every pending closed day for the eligible metrics.

    ``full_refresh_window`` (``abk run --full-refresh --from/--to``) forces
    re-materialization from the first day its half-open window touches
    through the end of the series (truncate-then-advance — see the inline
    comment) — otherwise a backfill would leave stale day state the WP4
    reader would trust (WP3 step 6).

    ``force_rebuild`` (``abk run --resync-cohort``, copy mode): day state was
    derived from the persisted copy the resync just declared poisoned, so it
    must not outlive it — every eligible series is dropped and re-rendered
    from the rebuilt cohort.

    ``metric_filter`` (``abk run --metric``, m11 DASH-4a) narrows the stage to
    ONE metric's series, so a per-metric recompute neither renders nor writes
    what it will not read. Every series is independent (its own
    ``(source_table, column_set_id)`` key and its own contiguous day range), so
    narrowing cannot break another metric's contiguity invariant — the withheld
    series simply keeps the days it already had, and the next run that includes
    the metric advances it. Two exceptions, both about a withheld series that
    must not be trusted afterwards: the driver passes ``None`` alongside
    ``force_rebuild`` (see its call site), and a ``full_refresh_window``
    TRUNCATES the withheld series over the refreshed window (see below).

    In copy mode the caller passes ``watermark_ts`` CLAMPED to the copy's
    coverage: the day render joins the persisted copy, and a day
    materialized from a partial cohort would freeze that way (state days,
    unlike results cutoffs, are never re-planned once written).
    """
    outcome = StateOutcome()
    days = closed_state_days(experiment, grid, watermark_ts)
    if not days:
        return outcome

    eligible = state_eligible_metrics(experiment, metrics_by_name, project_root)
    withheld: list[tuple[MetricConfig, str]] = []
    if metric_filter is not None:
        withheld = [item for item in eligible if item[0].name != metric_filter]
        eligible = [item for item in eligible if item[0].name == metric_filter]

    if withheld and full_refresh_window is not None:
        # A SCOPED --full-refresh must not leave a withheld series
        # stale-in-place. The flag exists for facts that changed inside an
        # already-materialized window, and those facts are shared by every
        # metric that reads them — while the WP4 gap check only sees ABSENCE,
        # never wrong content. So the withheld series are TRUNCATED over the
        # refreshed window and deliberately NOT re-rendered here: the operator
        # scoped the render cost away, contiguity survives (a shorter prefix —
        # the delete runs from the first touched day through the END of the
        # series, the same tail semantics the unfiltered truncate-then-advance
        # below has),
        # reads past the new last day fall back to full recompute (correct, and
        # warned by the reader), and the next run that does include the metric
        # re-materializes those days from the current facts. Doing nothing
        # instead was a reviewed defect: with compute.incremental_reads on, a
        # later routine run summed the stale days and silently persisted an
        # undercount (m11 DASH-4a review).
        refresh_from, refresh_to = full_refresh_window
        for metric, metric_sql in withheld:
            source_id, series_id = state_series_key(experiment, metric, metric_sql, project_root)
            last_day = tables.get_last_state_day(source_id, series_id)
            if last_day is None:
                continue
            touched = [
                sd.day
                for sd in days
                if sd.day <= last_day
                and sd.window.start_ts < refresh_to
                and sd.window.end_ts > refresh_from
            ]
            if not touched:
                continue
            tables.delete_state_days_from(source_id, series_id, min(touched))
            outcome.warnings.append(
                f"{experiment.name}/{metric.name}: day state truncated from "
                f"{min(touched)} through the end of the series — a "
                "--full-refresh window this metric was not selected for may have "
                "changed facts it reads; a later run that includes this metric "
                "re-derives those days (reads in between fall back to "
                "full-window recompute)"
            )

    for metric, metric_sql in eligible:
        source_id, series_id = state_series_key(experiment, metric, metric_sql, project_root)

        # The identity invalidation (WP3 step 2 + the R1 cohort fix): an
        # edited SQL body, role map, or cohort-shaping config orphans the old
        # series; superseded series under this source key are dropped so a
        # future reader can never sum a stale definition.
        for stale_id in tables.list_state_column_sets(source_id):
            if stale_id != series_id:
                tables.delete_state_series(source_id, stale_id)
        if force_rebuild:
            tables.delete_state_series(source_id, series_id)

        last_day = tables.get_last_state_day(source_id, series_id)
        if full_refresh_window is not None and last_day is not None:
            # Truncate-then-advance (an R1 crash-safety fix): deleting from
            # the first touched day BEFORE re-rendering guarantees a crash
            # mid-refresh leaves a contiguous, self-healing prefix — never a
            # freshly-covered get_last_state_day() over silently stale rows.
            # The tail past the window re-renders too; that is the price of
            # keeping "every day <= last_state_day is materialized" true
            # without a per-day ledger.
            refresh_from, refresh_to = full_refresh_window
            touched = [
                sd.day
                for sd in days
                if sd.day <= last_day
                and sd.window.start_ts < refresh_to
                and sd.window.end_ts > refresh_from
            ]
            if touched:
                tables.delete_state_days_from(source_id, series_id, min(touched))
                last_day = tables.get_last_state_day(source_id, series_id)
        pending = [sd for sd in days if last_day is None or sd.day > last_day]
        if not pending:
            continue

        n_days = len(pending)
        log(f"STATE {experiment.name}/{metric.name}: {n_days} closed days to materialize")
        beat_every = max(1, n_days // 20)
        for day_index, state_day in enumerate(pending, start=1):
            loaded = backend.load_window(metric, metric_sql, state_day.window)
            try:
                data = day_moments(metric, loaded)
            except StateMomentError as exc:
                # Truncate from the failing day (R1: never drop the whole
                # series — that re-rendered the full history every run with
                # zero retained state). Earlier days stay valid and
                # contiguous; the retry next run costs ONE render; the WP4
                # reader sees days past get_last_state_day as unmaterialized
                # and falls back to full recompute (§0.2 — never a
                # materialized-and-empty silent undercount).
                tables.delete_state_days_from(source_id, series_id, state_day.day)
                outcome.warnings.append(
                    f"{experiment.name}/{metric.name}: day state truncated at "
                    f"{state_day.day} — {exc}; reads past the previous day stay "
                    "on full-window recompute"
                )
                break
            outcome.rows_written += tables.replace_day_state(
                source_id, series_id, state_day.day, data
            )
            outcome.days_materialized += 1
            if n_days > 1 and (day_index % beat_every == 0 or day_index == n_days):
                log(
                    f"STATE {experiment.name}/{metric.name}: "
                    f"{day_index}/{n_days} days materialized"
                )
    return outcome
