"""Policy selector — Phase 3 active: Thompson sampling over eligible reflexes.

Contract (stable across phases — callers can rely on the return shape):

    select_policy(skill_name, caller, trigger_type, ...) -> PolicyChoice

Fallback rules — deterministic ``incumbent_default`` returns when:
  1. ``caller`` is in ``PROTECTED_CALLERS`` (recovery, rollback)
  2. ``skill_name`` is in ``PROTECTED_SKILLS`` (governance substrates)
  3. ``parent_run_id`` is set — a nested ``invoke_skill`` inherits the
     parent's policy so experimentation doesn't cascade
  4. The configured ``kill_switch`` is truthy
  5. Fewer than ``min_eligible`` candidates are available
  6. The evolution ledger is unavailable (defensive)

Selection algorithm (Phase 3):
  - Fetch ACTIVE reflexes whose ``scope.skill_name`` matches the current
    skill (and whose ``caller_in`` matches the current caller if set).
  - Always include ``incumbent_default`` and ``baseline_random`` as
    synthetic candidates — the first is posterior-tracked so regressions
    in the incumbent itself become visible; the second is a flat-prior
    exploration floor that can't be promoted.
  - Draw one Thompson sample per candidate from Beta(α, β) scaled by
    utility_ema (so history of success acts as a soft prior).
  - Pick the highest-sampled candidate; write selector telemetry.

A reflex that exists but is in status `draft` / `candidate` / `quarantined`
/ `retired` is NEVER sampled — only `active` status reaches the sampler.
This keeps the selector deterministic during experimentation (a candidate
is evaluated by the PromotionArena, not by running live with live traffic).
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from agntspace.evolution.bayes import posterior_mean, sample_beta
from agntspace.evolution.kernel import is_protected_caller, is_protected_skill
from agntspace.evolution.schemas import (
    POLICY_BASELINE_RANDOM,
    POLICY_INCUMBENT_DEFAULT,
    PolicyChoice,
)

if TYPE_CHECKING:
    from agntspace.evolution.ledger import EvolutionLedger

log = logging.getLogger(__name__)


# ── Synthetic candidates ──
# These two are always eligible; they anchor the uncertainty model.
#
# ``incumbent_default`` is the current deterministic path — posterior
# tracked so a regression in the incumbent itself (e.g. an LLM model
# provider degraded) becomes visible via its utility_ema declining.
#
# ``baseline_random`` has a flat prior and a small utility_ema floor
# so it almost never wins Thompson sampling but exists as an honest
# "we have no signal either way" anchor. Never promoted, never retired.


def _kill_switch_active(config: dict[str, Any] | None) -> bool:
    """Read ``selection.kill_switch`` from the loaded evolution config."""
    if not config or not isinstance(config, dict):
        return False
    sel = config.get("selection")
    if not isinstance(sel, dict):
        return False
    return bool(sel.get("kill_switch", False))


def _min_eligible(config: dict[str, Any] | None) -> int:
    default = 3
    if not config or not isinstance(config, dict):
        return default
    sel = config.get("selection", {})
    try:
        return int(sel.get("min_eligible_before_thompson", default))
    except (TypeError, ValueError):
        return default


def _reflex_scope_matches(
    reflex_scope: dict[str, Any], skill_name: str, caller: str,
) -> bool:
    """A reflex is eligible only if its scope covers (skill, caller)."""
    if not isinstance(reflex_scope, dict):
        return False
    scope_skill = str(reflex_scope.get("skill_name", ""))
    if scope_skill and scope_skill != skill_name:
        return False
    callers = reflex_scope.get("caller_in") or []
    if callers and caller and caller not in callers:
        return False
    return True


async def _eligible_reflexes(
    ledger: EvolutionLedger, skill_name: str, caller: str,
) -> list[tuple[str, float, float, float]]:
    """Return [(reflex_id, alpha, beta, utility_ema), ...] for active reflexes.

    Only ``status='active'`` reflexes are sampled — draft and candidate
    reflexes are evaluated by the promotion arena, not by the selector.
    Quarantined + retired reflexes never return here.
    """
    reflexes = await ledger.list_reflexes(kind="reflex", status="active")
    out: list[tuple[str, float, float, float]] = []
    for r in reflexes:
        if not _reflex_scope_matches(r.scope, skill_name, caller):
            continue
        post = await ledger.get_posterior(r.reflex_id)
        out.append((r.reflex_id, post.alpha, post.beta_param, post.utility_ema))
    return out


async def select_policy(
    skill_name: str,
    *,
    caller: str = "system",
    trigger_type: str = "",
    parent_run_id: str = "",
    context_keys: dict[str, str] | None = None,
    ledger: EvolutionLedger | None = None,
    config: dict[str, Any] | None = None,
) -> PolicyChoice:
    """Select a policy for the current invoke_skill call.

    Phase 3 behavior:
      - Fallback rules first (protected skill/caller/nested/kill switch/no ledger)
      - Thompson-sample over active reflexes scoped to (skill, caller) PLUS
        ``incumbent_default`` + ``baseline_random``
      - Return the highest-sampled policy with telemetry

    The telemetry (top-5 samples) is captured by the caller via
    ``EvolutionService.persist_selector_telemetry`` using the returned
    ``PolicyChoice`` plus sidecar data the caller assembles.
    """
    _ = context_keys, trigger_type  # reserved for future scope narrowing

    # Rule 1 + 2: protected skills / callers never experiment.
    if is_protected_skill(skill_name):
        return PolicyChoice.fallback(f"protected_skill:{skill_name}")
    if is_protected_caller(caller):
        return PolicyChoice.fallback(f"protected_caller:{caller}")

    # Rule 3: nested invocations inherit parent policy. Experimentation
    # in nested calls would cascade — you'd end up with a candidate
    # cluster for every intermediate agent in a chain, not for the
    # top-level skill the caller actually invoked.
    if parent_run_id:
        return PolicyChoice(
            policy_id=POLICY_INCUMBENT_DEFAULT,
            rationale=f"inherit_parent:{parent_run_id}",
            sampled_score=0.0,
            sampled_rank=0,
            is_fallback=True,
        )

    # Rule 4: kill switch.
    if _kill_switch_active(config):
        return PolicyChoice.fallback("kill_switch")

    # Rule 6: no ledger, no selection.
    if ledger is None:
        return PolicyChoice.fallback("no_ledger")

    # ── Phase 3: Thompson sampling ──
    # Gather active reflexes scoped to this (skill, caller), plus the
    # two synthetic candidates.
    try:
        active = await _eligible_reflexes(ledger, skill_name, caller)
    except Exception:
        log.debug("eligible_reflexes failed", exc_info=True)
        active = []

    # Synthetic candidates — posteriors optional, may not exist yet
    try:
        incumbent_post = await ledger.get_posterior(POLICY_INCUMBENT_DEFAULT)
        incumbent_sample = (
            POLICY_INCUMBENT_DEFAULT,
            incumbent_post.alpha,
            incumbent_post.beta_param,
            incumbent_post.utility_ema,
        )
    except Exception:
        incumbent_sample = (POLICY_INCUMBENT_DEFAULT, 1.0, 1.0, 0.0)

    try:
        baseline_post = await ledger.get_posterior(POLICY_BASELINE_RANDOM)
        baseline_sample = (
            POLICY_BASELINE_RANDOM,
            baseline_post.alpha,
            baseline_post.beta_param,
            baseline_post.utility_ema,
        )
    except Exception:
        baseline_sample = (POLICY_BASELINE_RANDOM, 1.0, 1.0, 0.0)

    eligible = active + [incumbent_sample, baseline_sample]

    # Minimum eligible count — below this, fall back deterministically.
    # With only the two synthetics, there's no real choice to make.
    min_e = _min_eligible(config)
    if len(eligible) < min_e:
        # Honest rationale: we need at least one active reflex before
        # sampling is meaningful. Stay on incumbent until mining + the
        # promotion arena has produced candidates that passed the gates.
        return PolicyChoice(
            policy_id=POLICY_INCUMBENT_DEFAULT,
            rationale=f"fallback:insufficient_candidates:{len(eligible)}<{min_e}",
            sampled_score=0.0,
            sampled_rank=0,
            is_fallback=True,
        )

    # Thompson sampling: one draw per candidate, weighted by utility_ema.
    # Scaling factor ``(0.5 + 0.5 * utility_ema)`` keeps every candidate
    # viable at the start (multiplier 0.5 when utility_ema=0) but favors
    # ones with a track record of producing good outcomes.
    samples: list[tuple[str, float, float]] = []
    for policy_id, alpha, beta, ema in eligible:
        beta_draw = sample_beta(alpha, beta)
        scaled = beta_draw * (0.5 + 0.5 * max(0.0, min(1.0, ema)))
        mean = posterior_mean(alpha, beta)
        samples.append((policy_id, scaled, mean))

    # Pick the highest-sampled candidate; sort descending.
    samples.sort(key=lambda t: -t[1])
    chosen_id, chosen_score, _chosen_mean = samples[0]

    # If the chosen policy is the incumbent and nothing else beat it,
    # the rationale reflects that rather than misreporting as "thompson".
    is_incumbent_choice = chosen_id == POLICY_INCUMBENT_DEFAULT

    return PolicyChoice(
        policy_id=chosen_id,
        rationale=(
            "thompson:incumbent_won" if is_incumbent_choice
            else f"thompson:n={len(eligible)}"
        ),
        sampled_score=round(chosen_score, 6),
        sampled_rank=0,
        is_fallback=is_incumbent_choice,  # incumbent wins = not experimental
    )


def build_telemetry_samples(
    *,
    chosen_id: str,
    eligible_snapshot: list[tuple[str, float, float, float]],
    sampled_scores: list[tuple[str, float, float]],
) -> list[dict[str, Any]]:
    """Format the top-5 samples for selector telemetry storage.

    Helper so EvolutionService + invoke_skill stay decoupled from
    telemetry format. Returns a list of dicts in
    ``[{reflex_id, sampled_score, posterior_mean}, ...]`` shape.
    """
    top = sorted(sampled_scores, key=lambda t: -t[1])[:5]
    return [
        {
            "reflex_id": pid,
            "sampled_score": round(float(score), 6),
            "posterior_mean": round(float(mean), 6),
            "chosen": pid == chosen_id,
        }
        for pid, score, mean in top
    ]


async def explain_policy(
    skill_name: str,
    *,
    caller: str = "system",
    trigger_type: str = "",
    parent_run_id: str = "",
    context_keys: dict[str, str] | None = None,
    ledger: EvolutionLedger | None = None,
    config: dict[str, Any] | None = None,
    rng_seed: int | None = None,
) -> dict[str, Any]:
    """Simulate what ``select_policy`` would do, WITHOUT persisting anything.

    Returns a rich diagnostic dict so users + developers can answer
    "why did the selector pick X here?" and "what would happen right now?"
    without having to trigger an actual invoke_skill call.

    Shape::

        {
          "skill_name": "kb-ingest",
          "caller": "watcher",
          "chosen": {
            "policy_id": "reflex-abc123",
            "rationale": "thompson:n=4",
            "sampled_score": 0.812,
            "is_fallback": false
          },
          "fallback_rule_triggered": null,   # or name of the fallback rule
          "eligible": [
            {
              "policy_id": "reflex-abc123",
              "kind": "active_reflex" | "incumbent_default" | "baseline_random",
              "alpha": 8.4, "beta": 2.1, "utility_ema": 0.84,
              "posterior_mean": 0.800,
              "sampled_score": 0.812,       # None when fallback short-circuits
              "in_top_5": true,
              "is_chosen": true,
              "scope": {...}                # for active_reflex only
            },
            ...
          ],
          "config": {
            "kill_switch": false,
            "min_eligible_before_thompson": 3,
            "protected_skill": false,
            "protected_caller": false
          }
        }

    ``rng_seed`` makes the sampling reproducible for testing / showing
    the same preview twice. When omitted, uses system random.
    """
    from agntspace.evolution.schemas import POLICY_BASELINE_RANDOM

    kill_switch = _kill_switch_active(config)
    min_e = _min_eligible(config)
    protected_skill = is_protected_skill(skill_name)
    protected_caller = is_protected_caller(caller)

    out: dict[str, Any] = {
        "skill_name": skill_name,
        "caller": caller,
        "trigger_type": trigger_type or skill_name,
        "parent_run_id": parent_run_id,
        "chosen": None,
        "fallback_rule_triggered": None,
        "eligible": [],
        "config": {
            "kill_switch": kill_switch,
            "min_eligible_before_thompson": min_e,
            "protected_skill": protected_skill,
            "protected_caller": protected_caller,
        },
    }

    # Short-circuit rules — match the ordering in select_policy so users
    # see the ACTUAL rule that would fire.
    if protected_skill:
        out["fallback_rule_triggered"] = f"protected_skill:{skill_name}"
        out["chosen"] = {
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "rationale": f"fallback:protected_skill:{skill_name}",
            "sampled_score": 0.0,
            "is_fallback": True,
        }
        return out
    if protected_caller:
        out["fallback_rule_triggered"] = f"protected_caller:{caller}"
        out["chosen"] = {
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "rationale": f"fallback:protected_caller:{caller}",
            "sampled_score": 0.0,
            "is_fallback": True,
        }
        return out
    if parent_run_id:
        out["fallback_rule_triggered"] = f"inherit_parent:{parent_run_id}"
        out["chosen"] = {
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "rationale": f"inherit_parent:{parent_run_id}",
            "sampled_score": 0.0,
            "is_fallback": True,
        }
        return out
    if kill_switch:
        out["fallback_rule_triggered"] = "kill_switch"
        out["chosen"] = {
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "rationale": "fallback:kill_switch",
            "sampled_score": 0.0,
            "is_fallback": True,
        }
        return out
    if ledger is None:
        out["fallback_rule_triggered"] = "no_ledger"
        out["chosen"] = {
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "rationale": "fallback:no_ledger",
            "sampled_score": 0.0,
            "is_fallback": True,
        }
        return out

    # Gather eligible reflexes + the two synthetic anchors.
    try:
        active = await _eligible_reflexes(ledger, skill_name, caller)
    except Exception:
        active = []

    # Fetch reflex rows so we can include their scope in the explanation
    active_rows: dict[str, Any] = {}
    try:
        reflex_rows = await ledger.list_reflexes(kind="reflex", status="active")
        for r in reflex_rows:
            active_rows[r.reflex_id] = r
    except Exception:
        pass

    try:
        incumbent_post = await ledger.get_posterior(POLICY_INCUMBENT_DEFAULT)
    except Exception:
        incumbent_post = None  # defensive
    try:
        baseline_post = await ledger.get_posterior(POLICY_BASELINE_RANDOM)
    except Exception:
        baseline_post = None

    def _pm(alpha: float, beta: float) -> float:
        total = alpha + beta
        return round(alpha / total, 6) if total > 0 else 0.5

    # Build the eligible list WITHOUT sampling first — gives full visibility.
    eligible_entries: list[dict[str, Any]] = []
    for rid, alpha, beta, ema in active:
        row = active_rows.get(rid)
        entry = {
            "policy_id": rid,
            "kind": "active_reflex",
            "alpha": round(alpha, 4),
            "beta": round(beta, 4),
            "utility_ema": round(ema, 4),
            "posterior_mean": _pm(alpha, beta),
            "sampled_score": None,
            "in_top_5": False,
            "is_chosen": False,
            "scope": dict(row.scope) if row is not None else {},
            "name": row.name if row is not None else "",
        }
        eligible_entries.append(entry)

    if incumbent_post is not None:
        eligible_entries.append({
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "kind": "incumbent_default",
            "alpha": round(incumbent_post.alpha, 4),
            "beta": round(incumbent_post.beta_param, 4),
            "utility_ema": round(incumbent_post.utility_ema, 4),
            "posterior_mean": _pm(incumbent_post.alpha, incumbent_post.beta_param),
            "sampled_score": None,
            "in_top_5": False,
            "is_chosen": False,
            "scope": {},
            "name": "incumbent_default",
        })
    if baseline_post is not None:
        eligible_entries.append({
            "policy_id": POLICY_BASELINE_RANDOM,
            "kind": "baseline_random",
            "alpha": round(baseline_post.alpha, 4),
            "beta": round(baseline_post.beta_param, 4),
            "utility_ema": round(baseline_post.utility_ema, 4),
            "posterior_mean": _pm(baseline_post.alpha, baseline_post.beta_param),
            "sampled_score": None,
            "in_top_5": False,
            "is_chosen": False,
            "scope": {},
            "name": "baseline_random",
        })

    if len(eligible_entries) < min_e:
        out["fallback_rule_triggered"] = f"insufficient_candidates:{len(eligible_entries)}<{min_e}"
        out["chosen"] = {
            "policy_id": POLICY_INCUMBENT_DEFAULT,
            "rationale": f"fallback:insufficient_candidates:{len(eligible_entries)}<{min_e}",
            "sampled_score": 0.0,
            "is_fallback": True,
        }
        out["eligible"] = eligible_entries
        return out

    # Simulate Thompson sampling deterministically when rng_seed is given
    import random as _random
    rng = _random.Random(rng_seed) if rng_seed is not None else _random.Random()

    # Sample each candidate, match select_policy math exactly
    scored: list[tuple[str, float]] = []
    for e in eligible_entries:
        beta_draw = sample_beta(float(e["alpha"]), float(e["beta"]), rng=rng)
        scaled = beta_draw * (0.5 + 0.5 * max(0.0, min(1.0, float(e["utility_ema"]))))
        e["sampled_score"] = round(scaled, 6)
        scored.append((e["policy_id"], scaled))

    scored.sort(key=lambda t: -t[1])
    chosen_id, chosen_score = scored[0]
    top5_ids = {pid for pid, _ in scored[:5]}
    for e in eligible_entries:
        if e["policy_id"] in top5_ids:
            e["in_top_5"] = True
        if e["policy_id"] == chosen_id:
            e["is_chosen"] = True

    is_incumbent_choice = chosen_id == POLICY_INCUMBENT_DEFAULT
    out["chosen"] = {
        "policy_id": chosen_id,
        "rationale": (
            "thompson:incumbent_won" if is_incumbent_choice
            else f"thompson:n={len(eligible_entries)}"
        ),
        "sampled_score": round(chosen_score, 6),
        "is_fallback": is_incumbent_choice,
    }
    # Sort eligible by sampled_score descending so the UI presentation is stable
    eligible_entries.sort(key=lambda e: -(e["sampled_score"] or 0))
    out["eligible"] = eligible_entries
    return out
