"""Canonical unified Claude/Codex account selection policy."""

from __future__ import annotations

from datetime import datetime
from decimal import Decimal, localcontext

from selector_primitives import (
    DEFAULT_HEADROOM_BAND,
    DEFAULT_SESSION_GATE,
    SCHEMA,
    SELECTOR_VERSION,
    canonical_sha256,
    collision_key,
    decimal_value,
    identity,
    request_error,
    timestamp,
)


def _parse_times(raw: dict) -> tuple[dict, set[str]]:
    parsed: dict[str, tuple[datetime, str] | None] = {}
    invalid: set[str] = set()
    for field in ("seen_at", "limited_until", "last_selected_at", "resets_at"):
        try:
            parsed[field] = timestamp(raw[field]) if raw.get(field) is not None else None
        except ValueError:
            parsed[field] = None
            invalid.add(field)
    return parsed, invalid


def _reservation_state(history: dict, now: datetime) -> tuple[str | None, bool]:
    value = history.get("active_expires_at")
    if value is None:
        return None, False
    parsed, canonical = timestamp(value)
    return canonical, parsed > now


def normalize_candidate(raw: dict, now: datetime, ordinal: int) -> dict:
    runner, generation, engine, account, valid_identity = identity(raw)
    weekly, weekly_text = decimal_value(raw.get("weekly_pct"))
    session, session_text = decimal_value(raw.get("session_pct"))
    quota_known = weekly is not None and session is not None
    decimals = [value for value in (weekly, session) if value is not None]
    precision = max([6] + [len(value.as_tuple().digits)
                           + max(0, -value.as_tuple().exponent) for value in decimals])
    with localcontext() as context:
        context.prec = precision + 3
        weekly_left = Decimal(100) - weekly if weekly is not None else None
        session_left = Decimal(100) - session if session is not None else None
        effective = min(weekly_left, session_left) if quota_known else None
    history = raw["reservation_history"]
    timed = {**raw, "last_selected_at": history["last_selected_at"]}
    parsed, invalid = _parse_times(timed)
    active_expiry, reserved = _reservation_state(history, now)
    seen, limited, last = parsed["seen_at"], parsed["limited_until"], parsed["last_selected_at"]
    status = str(raw.get("status") or "").strip().lower()
    conditions = [
        ("invalid_identity", not valid_identity), ("status_not_active", status != "active"),
        ("invalid_seen_at", "seen_at" in invalid),
        ("future_telemetry", bool(seen and (seen[0] - now).total_seconds() > 30)),
        ("stale_telemetry", bool(seen and (now - seen[0]).total_seconds() > 300)),
        ("invalid_limited_until", "limited_until" in invalid),
        ("account_limited", bool(limited and limited[0] > now)),
        ("invalid_last_selected_at", "last_selected_at" in invalid),
        ("future_last_selected_at", bool(last and (last[0] - now).total_seconds() > 30)),
        ("provider_incapable", raw.get("provider_capable") is not True),
        ("active_reservation", reserved)]
    exclusion = next((code for code, applies in conditions if applies), None)
    canon = lambda value: decimal_value(value)[1] if value is not None else None
    return {
        "input_ordinal": ordinal, "identity_valid": valid_identity, "runner_id": runner,
        "runner_generation": generation, "engine": engine, "account_id": account, "status": status,
        "weekly_pct": weekly_text, "session_pct": session_text, "quota_known": quota_known,
        "effective_headroom": canon(effective), "weekly_remaining": canon(weekly_left),
        "session_remaining": canon(session_left),
        "resets_at": parsed["resets_at"][1] if parsed["resets_at"] else None,
        "limited_until": limited[1] if limited else None, "seen_at": seen[1] if seen else None,
        "provider_capable": raw.get("provider_capable") is True,
        "reservation_history": {"active_expires_at": active_expiry,
                                "last_selected_at": last[1] if last else None},
        "last_selected_at": last[1] if last else None, "eligible": exclusion is None,
        "exclusion_code": exclusion}


def _policy_context(request: dict, rows: list[dict]) -> tuple[str | None, tuple, set[tuple], int]:
    policy, required = request["policy"], request.get("required_engine")
    if policy == "default_claude":
        required = "claude"
    elif policy == "default_codex":
        required = "codex"
    elif policy not in {"explicit", "producer_retry", "reviewer"}:
        required = None
    producer = identity(request.get("producer_identity") or {})[:4]
    excluded = {identity(value)[:4] for value in request.get("excluded_identities", [])}
    alternatives = sum(
        row["eligible"] and (required is None or row["engine"] == required)
        and _row_identity(row) != producer and _row_identity(row) not in excluded for row in rows)
    return required, producer, excluded, alternatives


def _row_identity(row: dict) -> tuple:
    return tuple(row[name] for name in ("runner_id", "runner_generation", "engine", "account_id"))


def _apply_policy(request: dict, rows: list[dict]) -> tuple[tuple, int]:
    required, producer, excluded, alternatives = _policy_context(request, rows)
    for row in rows:
        key = _row_identity(row)
        policy_excluded = required is not None and row["engine"] != required
        policy_excluded |= key in excluded
        policy_excluded |= request["policy"] == "reviewer" and alternatives > 0 and key == producer
        if row["eligible"] and policy_excluded:
            row["eligible"] = False
            row["exclusion_code"] = "policy_excluded"
    return producer, alternatives


def _winner_key(row: dict) -> tuple:
    descending_headroom = Decimal(row["effective_headroom"] or "-1").copy_negate()
    return (
        not row["quota_known"], descending_headroom,
        row["last_selected_at"] is not None, row["last_selected_at"] or "",
        row["engine"], row["account_id"], row["runner_id"], row["runner_generation"])


def _session_gate(rows: list[dict], gate: Decimal | None) -> list[dict]:
    """The FIRST cut: accounts whose 5h session bucket still has room.

    The operator (2026-09-03): "among accounts where high session limits it must
    choose randomly from ones where highest weekly limits." A nearly-spent session
    bucket is about to reject the launch whatever the weekly headroom says, so it
    disqualifies the account outright instead of merely nudging a tie-break. The
    gate COMPARES, it never empties the pool: when nobody clears it the caller
    falls back to every known row. Unknown quota never clears it — a missing
    reading is not evidence of room.
    """
    if gate is None:
        return []
    return [row for row in rows if Decimal(row["session_pct"]) <= gate]


def _band_floor(rows: list[dict], band: Decimal) -> Decimal | None:
    """The lowest WEEKLY remaining still counted as "as good as the best" this round.

    Ranking strictly by headroom hands every task to whichever account is on top,
    which is precisely how one account's limit gets burned to zero while three
    others idle: the runner-up only ever wins after the leader has been spent
    below it. Accounts within ``band`` points of the leader are treated as equally
    good, and the tie-break below spreads work across them.

    The band is measured on weekly remaining, not on min(weekly, session): session
    headroom already had its say in ``_session_gate`` above, and letting it back in
    here both pushed the best weekly account out of the band over a half-spent 5h
    bucket and let a fully spent one stay in. Rows reaching here are the gated set,
    so every one of them has a known reading.
    """
    known = [Decimal(row["weekly_remaining"]) for row in rows
             if row["quota_known"] and row["weekly_remaining"] is not None]
    return max(known) - band if known else None


def _spread(reservation_key: str, row: dict) -> str:
    """A stable pseudo-random ordinal for this row IN THIS REQUEST.

    Real randomness cannot live here — the response carries a reproducible
    digest of its own inputs, and two runs of the same request must agree. A
    digest over (reservation key, identity) gives every task its own order over
    the band while staying a pure function of the request, so a burst of
    parallel launches spreads instead of stacking on one account.
    """
    return canonical_sha256({"reservation_key": reservation_key,
                             "identity": list(_row_identity(row))})


def _band_key(reservation_key: str):
    """Order INSIDE the band: least-recently-used first, then the spread digest.

    Headroom deliberately does not appear — inside the band it is what we are
    choosing to ignore. LRU is what makes this rotate rather than merely jitter:
    with three banded accounts, three sequential tasks touch all three.
    """
    def key(row: dict) -> tuple:
        return (row["last_selected_at"] is not None, row["last_selected_at"] or "",
                _spread(reservation_key, row),
                row["engine"], row["account_id"], row["runner_id"], row["runner_generation"])
    return key


def _snapshot_key(row: dict) -> tuple:
    if row["identity_valid"]:
        return 0, row["runner_id"], row["runner_generation"], row["engine"], row["account_id"]
    return 1, row["input_ordinal"], 0, "", ""


def _error(code: str, detail: dict) -> dict:
    return {"schema": SCHEMA, "selector_version": SELECTOR_VERSION, "ok": False,
            "error_code": code, "error_detail": detail}


def select(request: object) -> dict:
    if error := request_error(request):
        return _error(error[0], error[1])
    assert isinstance(request, dict)
    now, canonical_now = timestamp(request["database_now"])
    rows = [normalize_candidate(item, now, index)
            for index, item in enumerate(request["candidates"])]
    identities: dict[tuple, list[int]] = {}
    for row in rows:
        if row["identity_valid"]:
            identities.setdefault(collision_key(_row_identity(row)), []).append(row["input_ordinal"])
    duplicate = next((items for items in identities.values() if len(items) > 1), None)
    if duplicate:
        return _error("duplicate_candidate_identity", {"input_ordinals": duplicate})
    producer, alternatives = _apply_policy(request, rows)
    eligible = [row for row in rows if row["eligible"]]
    if not eligible:
        return _error("no_candidate", {"eligible_count": 0})
    band, band_text = decimal_value(request.get("headroom_band", DEFAULT_HEADROOM_BAND))
    gate, gate_text = decimal_value(request.get("session_gate", DEFAULT_SESSION_GATE))
    # Two cuts, in this order (operator's 2026-09-03 decision): the session gate says
    # WHO may be considered, the weekly band says which of those count as equally
    # good. "healthy or known" is the gate stepping aside when nobody clears it.
    known = [row for row in eligible if row["quota_known"]]
    healthy = _session_gate(known, gate)
    ranked = healthy or known
    floor = _band_floor(ranked, band) if band is not None else None
    if floor is None:
        # No usable quota anywhere: nothing to band, so keep the plain ordering
        # (which already ranks unknown-quota rows last and rotates on ties).
        banded = eligible
        winner = min(banded, key=_winner_key)
    else:
        banded = [row for row in ranked if Decimal(row["weekly_remaining"]) >= floor]
        winner = min(banded, key=_band_key(request["reservation_key"]))
    rows.sort(key=_snapshot_key)
    snapshot = canonical_sha256(rows)
    chosen = dict(zip(("runner_id", "runner_generation", "engine", "account_id"), _row_identity(winner)))
    digest_input = {
        "schema": SCHEMA, "selector_version": SELECTOR_VERSION,
        "database_now": canonical_now, "policy": request["policy"],
        "required_engine": request.get("required_engine"), "reservation_key": request["reservation_key"],
        # The band and the gate both change which account wins, so both belong in
        # the proof. (This is why the golden digests moved on 2026-09-03.)
        "headroom_band": band_text,
        "session_gate": gate_text,
        "candidate_snapshot_digest": snapshot, "selected_identity": chosen}
    score_fields = ("quota_known", "weekly_pct", "session_pct", "weekly_remaining",
                    "session_remaining", "effective_headroom", "last_selected_at")
    response = {"schema": SCHEMA, "selector_version": SELECTOR_VERSION, "ok": True, **chosen,
                "score_basis": {name: winner[name] for name in score_fields},
                "eligible_count": len(eligible),
                "eligible_alternative_count": alternatives,
                # Observability: which accounts were considered interchangeable,
                # so "why did it not pick the emptiest one" has an answer.
                "headroom_band": band_text,
                "band_floor": format(floor, "f") if floor is not None else None,
                "band_count": len(banded),
                # ...and how many accounts the session gate let through at all
                # (0 means nobody cleared it and the gate stepped aside).
                "session_gate": gate_text,
                "session_ok_count": len(healthy),
                "candidate_snapshot_digest": snapshot,
                "selection_digest": canonical_sha256(digest_input)}
    if request["policy"] == "reviewer" and not alternatives and _row_identity(winner) == producer:
        response["fallback_reason"] = "sole_eligible_account"
    return response