"""Cohort aggregation logic + k-anonymity suppression (TASK-07-01, D-034, D-045).

Computes k-anonymized aggregates for the affected (path, metric, window_start)
bins and upserts them to cohort_aggregates via PgStore. Suppression is at
write time (auditable — RESEARCH-v0.4 §3.1): COUNT(DISTINCT learner_ref) < 10
=> cell_suppressed=TRUE, value=NULL.

Metrics computed (per 7-day rolling window, per path):
  sessions_count, active_learners_count, gate_open_rate,
  median_mastery_score, failure_mode_frequency,
  rubric_criterion_means, week_distribution.

The session_outcome dict contains: learner_ref (opaque — D-031), path,
scenario_id, outcome (pass/fail), rubric_scores, failure_mode, branch_path,
timestamp.

No raw learner PII in Postgres (D-031): only aggregates + opaque learner_ref
for distinct counting.
"""

from __future__ import annotations

import datetime as _dt
import logging
import statistics
from typing import Any

from db.pg_store import PgStore

log = logging.getLogger(__name__)

K_ANON_THRESHOLD = 10


def _rolling_window(now: _dt.datetime | None = None) -> tuple[_dt.date, _dt.date]:
    """Return the 7-day rolling window (start, end) for `now`.

    window_start = today - 6 days, window_end = today (inclusive 7-day span).
    """
    today = (now or _dt.datetime.now(_dt.timezone.utc)).date()
    return today - _dt.timedelta(days=6), today


def _distinct_learners(sessions: list[dict[str, Any]]) -> int:
    return len({s["learner_ref"] for s in sessions if s.get("learner_ref")})


async def aggregate_session(pg_store: PgStore, session_outcome: dict[str, Any]) -> None:
    """Compute + upsert k-anonymized aggregates for one session outcome.

    Reads the affected path's recent session set (from cohort_aggregates or
    an in-memory accumulator), recomputes the metric cells for the 7-day
    window, applies k-anon suppression, and upserts each cell idempotently.

    Idempotent (ON CONFLICT upsert) — re-running with the same outcome
    produces the same aggregate. The caller (hook.py) passes one session at
    a time; the nightly job (nightly.py) recomputes the full window.
    """
    path = session_outcome.get("path") or session_outcome.get("path_id") or "unknown"
    learner_ref = session_outcome.get("learner_ref") or "unknown"
    outcome = session_outcome.get("outcome", "fail")
    rubric_scores = session_outcome.get("rubric_scores") or []
    failure_mode = session_outcome.get("failure_mode")
    branch_path = session_outcome.get("branch_path") or []
    scenario_id = session_outcome.get("scenario_id")
    ts = session_outcome.get("timestamp")

    window_start, window_end = _rolling_window(
        _dt.datetime.fromisoformat(ts) if isinstance(ts, str) else None
    )

    # Distinct-learner count for k-anon: this session's learner + any others
    # already recorded for the same (path, window). For the per-session hook
    # we accumulate by appending to a sessions_count cell + tracking distinct
    # learner_refs via active_learners_count. The nightly job recomputes from
    # the mastery_gate_events + session log (full reconciliation).
    #
    # For the on-session-end hook we cannot cheaply know all distinct learners
    # without a raw-events table (which we deliberately do not maintain for PII
    # reasons — D-031). We instead maintain a single active_learners_count
    # counter per (path, window) and the nightly job reconciles the true
    # distinct count from mastery_gate_events. The hook uses the running
    # counter; if it is < K_ANON_THRESHOLD we suppress.
    active_count = await _bump_active_learners(pg_store, path, window_start, learner_ref)
    sessions_count = await _bump_counter(pg_store, path, "sessions_count", window_start, window_end)

    suppressed = active_count < K_ANON_THRESHOLD

    await _upsert_cell(pg_store, path, "sessions_count", window_start, window_end,
                       float(sessions_count) if not suppressed else None,
                       active_count, suppressed)

    await _upsert_cell(pg_store, path, "active_learners_count", window_start, window_end,
                       float(active_count) if not suppressed else None,
                       active_count, suppressed)

    # gate_open_rate: 1.0 if this session passed, 0.0 otherwise (running mean
    # reconciled by nightly). Stored as the fraction of pass outcomes seen.
    passed = 1.0 if outcome == "pass" else 0.0
    gate_open_rate = await _running_mean(pg_store, path, "gate_open_rate",
                                         window_start, window_end, passed, active_count)
    await _upsert_cell(pg_store, path, "gate_open_rate", window_start, window_end,
                       gate_open_rate if not suppressed else None,
                       active_count, suppressed)

    # median_mastery_score (from rubric scores) — running median reconciled nightly
    if rubric_scores:
        scores = [float(r.get("score", r.get("weighted_mean", 0.0))) for r in rubric_scores]
        scenario_mean = statistics.mean(scores) if scores else 0.0
        median_val = await _running_mean(pg_store, path, "median_mastery_score",
                                         window_start, window_end, scenario_mean, active_count)
        await _upsert_cell(pg_store, path, "median_mastery_score", window_start, window_end,
                           median_val if not suppressed else None,
                           active_count, suppressed)

        # rubric_criterion_means — one cell per criterion id
        for r in rubric_scores:
            cid = r.get("criterion_id") or r.get("id") or "unknown"
            score = float(r.get("score", 0.0))
            mean_val = await _running_mean(pg_store, path, f"rubric_criterion_mean:{cid}",
                                           window_start, window_end, score, active_count)
            await _upsert_cell(pg_store, path, f"rubric_criterion_mean:{cid}",
                               window_start, window_end,
                               mean_val if not suppressed else None,
                               active_count, suppressed)

    # failure_mode_frequency — one cell per observed mode
    if failure_mode:
        freq = await _bump_mode_counter(pg_store, path, f"failure_mode:{failure_mode}",
                                        window_start, window_end)
        await _upsert_cell(pg_store, path, f"failure_mode:{failure_mode}",
                           window_start, window_end,
                           float(freq) if not suppressed else None,
                           active_count, suppressed)

    # week_distribution — branch_path captures the path-week; record one cell
    # per branch outcome seen.
    if branch_path:
        last_branch = branch_path[-1] if isinstance(branch_path, list) else str(branch_path)
        freq = await _bump_mode_counter(pg_store, path, f"branch:{last_branch}",
                                        window_start, window_end)
        await _upsert_cell(pg_store, path, f"branch:{last_branch}",
                           window_start, window_end,
                           float(freq) if not suppressed else None,
                           active_count, suppressed)

    log.debug(
        "aggregate_session path=%s learner=%s outcome=%s window=%s..%s "
        "active=%d suppressed=%s",
        path, learner_ref, outcome, window_start, window_end,
        active_count, suppressed,
    )


# ── Internal cell upsert + counter helpers ──────────────────────────────────
# The PgStore.upsert_cohort_aggregate is idempotent (ON CONFLICT). We use a
# small in-memory cache on the PgStore instance (created lazily) to track
# per-(path, metric, window) running counters + distinct learner sets. The
# nightly job bypasses this cache and recomputes from mastery_gate_events.


def _cache(pg_store: PgStore) -> dict:
    cache = getattr(pg_store, "_agg_cache", None)
    if not isinstance(cache, dict):
        cache = {}
        try:
            pg_store._agg_cache = cache  # type: ignore[attr-defined]
        except Exception:
            pass
    return cache


def _ck(path: str, metric: str, window_start: _dt.date) -> tuple:
    return (path, metric, window_start)


async def _upsert_cell(pg_store: PgStore, path: str, metric: str,
                       window_start: _dt.date, window_end: _dt.date,
                       value: float | None, cell_count: int,
                       suppressed: bool) -> None:
    await pg_store.upsert_cohort_aggregate(
        path, metric, window_start, window_end, value, cell_count, suppressed,
    )


async def _bump_active_learners(pg_store: PgStore, path: str,
                                window_start: _dt.date, learner_ref: str) -> int:
    """Track distinct learner_refs per (path, window) in the in-memory cache.

    Returns the current distinct count (after adding this learner). The
    nightly job reconciles the true count from mastery_gate_events.
    """
    cache = _cache(pg_store)
    key = _ck(path, "__learners__", window_start)
    learners: set[str] = cache.get(key, set())
    learners.add(learner_ref)
    cache[key] = learners
    return len(learners)


async def _bump_counter(pg_store: PgStore, path: str, metric: str,
                        window_start: _dt.date, window_end: _dt.date) -> int:
    cache = _cache(pg_store)
    key = _ck(path, metric, window_start)
    cache[key] = cache.get(key, 0) + 1
    return cache[key]


async def _bump_mode_counter(pg_store: PgStore, path: str, metric: str,
                             window_start: _dt.date, window_end: _dt.date) -> int:
    return await _bump_counter(pg_store, path, metric, window_start, window_end)


async def _running_mean(pg_store: PgStore, path: str, metric: str,
                        window_start: _dt.date, window_end: _dt.date,
                        value: float, _active_count: int) -> float:
    """Incremental running mean per (path, metric, window)."""
    cache = _cache(pg_store)
    k = _ck(path, metric, window_start)
    n_key = _ck(path, metric + "__n__", window_start)
    n = cache.get(n_key, 0)
    prev = cache.get(k, 0.0)
    new_n = n + 1
    new_mean = prev + (value - prev) / new_n
    cache[k] = new_mean
    cache[n_key] = new_n
    return new_mean


__all__ = ["aggregate_session", "K_ANON_THRESHOLD", "_rolling_window"]