"""Retrieval-tuning step logic and named constants.

Pure-function module: takes analytics outputs (doc_utilization, score_summary)
and a tuning state dict, returns the next state. Does no I/O.

See docs/superpowers/specs/2026-05-08-retrieval-tuning-loop-design.html.
"""
from __future__ import annotations
import copy
from datetime import datetime, timezone


#: Maximum cumulative penalty per doc. 0.40 = 40% score haircut. Keeps a
#: stuck doc retrievable when nothing else is in the room.
MAX_PENALTY = 0.40

#: Per-step penalty increment. With once-daily cadence, full saturation
#: takes >= 8 consecutive steps on the same doc.
PENALTY_DELTA = 0.05

#: A doc is eligible for a penalty increment iff its util_pct (over the
#: window) is below this threshold. Same 0-100 scale as
#: analytics.doc_utilization()'s util_pct field.
UTIL_PENALTY_THRESHOLD = 30.0

#: Minimum retrievals in the window before a doc is eligible for action.
#: Cuts noise from one-off retrievals.
MIN_RETRIEVALS_FOR_ACTION = 10

#: Per-dimension regression floor on a 1-5 scale. If any rubric dimension's
#: post-step mean drops by >= this value vs. pre-step, freeze + rollback.
REGRESSION_THRESHOLD = 0.30

#: Minimum judged turns in the post-step window before regression check
#: is meaningful. Below this, skip the check (and skip proposing a step).
MIN_TURNS_FOR_REGRESSION_CHECK = 20

#: Tuning state schema version. Bump on incompatible schema change.
TUNING_STATE_VERSION = 1

#: Penalty value at or above which a doc is excluded from kind-quota fill.
#: Below the floor, penalties only adjust within-kind ordering (already true).
#: 0.30 is six consecutive penalty steps, signalling the tuner is sure.
QUOTA_INELIGIBILITY_FLOOR = 0.30


def empty_state() -> dict:
    """Default tuning state: no penalties, not frozen."""
    return {
        "version": TUNING_STATE_VERSION,
        "updated_at": None,
        "frozen": False,
        "frozen_reason": None,
        "step_count": 0,
        "penalties": {},
    }


def now_utc_iso() -> str:
    """ISO-8601 UTC timestamp with second precision: '2026-05-08T18:45:03Z'."""
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def propose_step(
    util: dict[str, dict],
    *,
    current_state: dict,
) -> dict | None:
    """Propose at most one penalty increment for the lowest-util eligible doc.

    `util`: output of analytics.doc_utilization() filtered to the working
            window. Each value has retrieved_count, used_count, util_pct.
    `current_state`: the live retrieval_tuning state (read-only here).
            Caller (cli._cmd_tune) is responsible for short-circuiting on
            current_state["frozen"] before invoking this function — the
            frozen check is not performed here.

    Returns a step_action dict matching the audit-log shape, or None when
    no doc is eligible (all above threshold, or all already at MAX_PENALTY,
    or below MIN_RETRIEVALS_FOR_ACTION).
    """
    pens = current_state.get("penalties", {})
    eligible: list[tuple[str, dict]] = []
    for doc_id, info in util.items():
        if info["retrieved_count"] < MIN_RETRIEVALS_FOR_ACTION:
            continue
        if info["util_pct"] >= UTIL_PENALTY_THRESHOLD:
            continue
        if pens.get(doc_id, 0.0) >= MAX_PENALTY:
            continue
        eligible.append((doc_id, info))

    if not eligible:
        return None

    # Lowest util_pct first; ties broken by highest retrieved_count.
    eligible.sort(key=lambda kv: (kv[1]["util_pct"], -kv[1]["retrieved_count"]))
    doc_id, info = eligible[0]
    before = float(pens.get(doc_id, 0.0))
    after = min(before + PENALTY_DELTA, MAX_PENALTY)
    return {
        "kind": "penalty_increment",
        "doc_id": doc_id,
        "delta": round(after - before, 4),
        "penalty_before": round(before, 4),
        "penalty_after": round(after, 4),
        "reason": (
            f"util_pct={info['util_pct']} over {info['retrieved_count']} "
            f"retrievals"
        ),
    }


def check_regression(
    scores_before: dict[str, float],
    scores_after: dict[str, float],
    *,
    judged_turns: int,
    before_stats: dict[str, dict] | None = None,
    after_stats: dict[str, dict] | None = None,
) -> dict:
    """Pure check: did any dimension drop by >= REGRESSION_THRESHOLD?

    Returns a dict matching the audit log's `regression_check` shape.
    Sets `skipped=True` (and `any_dimension_regressed=False`) when there
    aren't enough turns to evaluate.

    When `before_stats` and `after_stats` are provided (each shaped as
    {dim: {mean, std, n}} per analytics.score_summary_with_stats), the
    return value also includes a confidence block with sample sizes and a
    Welch's t on the max-drop dimension. Telemetry only — does not gate
    the freeze decision.
    """
    if judged_turns < MIN_TURNS_FOR_REGRESSION_CHECK:
        return {
            "any_dimension_regressed": False,
            "max_drop_dimension": None,
            "max_drop_value": 0.0,
            "threshold": REGRESSION_THRESHOLD,
            "skipped": True,
        }

    shared = set(scores_before.keys()) & set(scores_after.keys())
    drops = {
        d: scores_before[d] - scores_after[d]
        for d in shared
    }
    if not drops:
        return {
            "any_dimension_regressed": False,
            "max_drop_dimension": None,
            "max_drop_value": 0.0,
            "threshold": REGRESSION_THRESHOLD,
            "skipped": False,
        }
    max_dim, max_drop = max(drops.items(), key=lambda kv: kv[1])
    out = {
        "any_dimension_regressed": max_drop >= REGRESSION_THRESHOLD,
        "max_drop_dimension": max_dim,
        "max_drop_value": round(max_drop, 4),
        "threshold": REGRESSION_THRESHOLD,
        "skipped": False,
    }
    if (
        before_stats and after_stats
        and max_dim in before_stats and max_dim in after_stats
    ):
        b = before_stats[max_dim]
        a = after_stats[max_dim]
        n_b, n_a = b["n"], a["n"]
        s_b, s_a = b["std"], a["std"]
        # Welch's t and pooled std (descriptive). Guard against n<2 / std=0.
        denom_sq = (
            (s_b * s_b / n_b if n_b else 0.0)
            + (s_a * s_a / n_a if n_a else 0.0)
        )
        denom = denom_sq ** 0.5 if denom_sq > 0 else 0.0
        t_stat = round(max_drop / denom, 4) if denom > 0 else None
        std_pooled = round(((s_b * s_b + s_a * s_a) / 2) ** 0.5, 4)
        out["confidence"] = {
            "n_before": n_b,
            "n_after": n_a,
            "max_drop_t_stat": t_stat,
            "max_drop_std_pooled": std_pooled,
            "interpretation": "informational_only",
        }
    return out


def apply_step(state: dict, step_action: dict) -> dict:
    """Return a new state with the step applied. Does not mutate `state`."""
    new = copy.deepcopy(state)
    if step_action["kind"] == "penalty_increment":
        new["penalties"][step_action["doc_id"]] = round(
            float(step_action["penalty_after"]), 4
        )
    new["step_count"] = int(state.get("step_count", 0)) + 1
    new["updated_at"] = now_utc_iso()
    return new


def rollback_step(state: dict, rolled_back_step: dict, *, reason: str) -> dict:
    """Restore the prior penalty value and set frozen=True."""
    new = copy.deepcopy(state)
    sa = rolled_back_step["step_action"]
    doc_id = sa["doc_id"]
    before = float(sa["penalty_before"])
    if before > 0.0:
        new["penalties"][doc_id] = round(before, 4)
    else:
        new["penalties"].pop(doc_id, None)
    new["frozen"] = True
    new["frozen_reason"] = reason
    new["updated_at"] = now_utc_iso()
    # step_count stays as-is — audit log is monotonic.
    return new