# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Reward Signal: RPE computation, credit assignment, and reward propagation."""
import json
import logging
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from pulse.core.paths import get_state_dir

log = logging.getLogger("pulse.reward_signal")

_REWARD_DIR = get_state_dir() / "reward"
REWARD_SIGNALS_FILE = _REWARD_DIR / "reward_signals.json"
CREDIT_LEDGER_FILE = _REWARD_DIR / "credit_ledger.json"

MAX_SIGNALS = 200       # Rolling window of recent signals
MAX_CREDIT_ITEMS = 2000 # Max items tracked in ledger
CREDIT_DECAY = 0.95     # Per-item positional decay (closer to search = more credit)

REINFORCE_DELTA = 0.2   # Salience boost on positive credit
CONTRADICT_DELTA = 0.4  # Confidence penalty on negative credit (2x loss aversion)
LOSS_AVERSION_RATIO = 2.0  # Negative RPE amplified by this factor (Kahneman-Tversky)
MIN_CREDIT = 0.05       # Items below this threshold get no update

_AGENT_REWARD_HISTORY: dict = {}  # {agent: [reward, ...]}
_AGENT_RPE_WINDOW = 20            # Rolling window size

_CONSECUTIVE_FAILURES: dict = {}  # {domain: count}
_CONSECUTIVE_FAILURE_THRESHOLD = 3  # Emit warning pattern after N consecutive failures
@dataclass
class RewardSignal:
    task_id: str; agent: str; success: bool; domain: str
    rpe_delta: float; items_credited: dict
    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())

def compute_rpe(success: bool, expected_success_rate: float) -> float:
    """Compute Reward Prediction Error: actual - expected. Range ~[-1.0, 1.0]."""
    actual = 1.0 if success else 0.0
    expected = max(0.05, min(0.95, expected_success_rate))  # Clamp to avoid division issues
    return actual - expected

def assign_credit(item_ids: list[str], rpe_delta: float, item_metadata: dict = None) -> dict[str, float]:
    """Credit assignment: positional decay + loss aversion + diversity bonus for underdogs."""
    if not item_ids or rpe_delta == 0:
        return {}
    credits = {}
    direction = 1.0 if rpe_delta > 0 else -1.0
    magnitude = abs(rpe_delta)
    if direction < 0:
        magnitude *= LOSS_AVERSION_RATIO
    for i, item_id in enumerate(item_ids):
        positional = CREDIT_DECAY ** i
        credit = direction * magnitude * positional
        if item_metadata and isinstance(item_metadata, dict):
            meta = item_metadata.get(item_id, {})
            if meta.get("access_count", 0) < 5 or meta.get("evidence_count", 0) < 3:
                credit *= 1.5  # Underdog diversity bonus
        if abs(credit) >= MIN_CREDIT:
            credits[item_id] = round(credit, 4)
    return credits

def propagate_reward(credits: dict) -> int:
    """Apply credit/blame to brain items. Delegated to reward_propagation module.

    Item 43: Somatic marker — positive_delta=0.2 (reward), negative_delta=0.3 (failure)
        UPDATE lessons SET salience = MIN(7.0, salience + delta) WHERE domain matched
        UPDATE failures SET salience = MIN(7.0, salience + delta) WHERE domain matched
        UPDATE patterns SET salience = MIN(7.0, salience + delta) WHERE domain matched
        Tag-and-capture (Item 6): confidence boost MIN(1.0, confidence + 0.15)
    except Exception: pass  # fail-open on missing tables
    """
    try:
        from pulse.core.reward.reward_propagation import propagate_reward as _propagate
        return _propagate(credits)
    except Exception as e:
        log.warning("propagate_reward delegation failed: %s", e)
        return 0

def _get_expected_success_rate(domain: str) -> float:
    """Get expected success rate from metamemory domain confidence."""
    try:
        from pulse.brain.metacognition.metamemory import get_metamemory
        meta = get_metamemory()
        if domain in meta:
            return meta[domain].get("avg_confidence", 0.5)
    except Exception:
        pass
    return 0.5  # Default: 50/50

def _update_credit_ledger(credits: dict[str, float]):
    """Update rolling credit ledger for audit trail."""
    _REWARD_DIR.mkdir(parents=True, exist_ok=True)
    try:
        from pulse.core.brain_io import _acquire
        with _acquire(CREDIT_LEDGER_FILE):
            ledger = {}
            if CREDIT_LEDGER_FILE.exists():
                try: ledger = json.loads(CREDIT_LEDGER_FILE.read_text(encoding="utf-8"))
                except Exception: ledger = {}
            for item_id, credit in credits.items():
                entry = ledger.setdefault(item_id, {"total_credit": 0.0, "positive_uses": 0, "negative_uses": 0})
                entry["total_credit"] = round(entry["total_credit"] + credit, 4)
                if credit > 0: entry["positive_uses"] += 1
                else: entry["negative_uses"] += 1
            if len(ledger) > MAX_CREDIT_ITEMS:
                ledger = dict(sorted(ledger.items(), key=lambda x: abs(x[1].get("total_credit", 0)), reverse=True)[:MAX_CREDIT_ITEMS])
            CREDIT_LEDGER_FILE.write_text(json.dumps(ledger, indent=2, ensure_ascii=False), encoding="utf-8")
    except Exception as e: log.warning("Credit ledger update failed: %s", e)

def record_task_outcome(task_id: str, agent: str, success: bool,
                         domain: str = "general",
                         retrieved_item_ids: Optional[list[str]] = None) -> RewardSignal:
    """Full reward pipeline: RPE → credit → propagation → logging.

    Called from pulse_save.save() after task completion.
    """
    retrieved_item_ids = retrieved_item_ids or []

    expected = _get_expected_success_rate(domain)
    rpe_delta = compute_rpe(success, expected)

    try:
        _h = _AGENT_REWARD_HISTORY.setdefault(agent, [])
        _h.append(1.0 if success else 0.0)
        if len(_h) > _AGENT_RPE_WINDOW: _h.pop(0)
        _arpe = (1.0 if success else 0.0) - sum(_h) / len(_h)
        if abs(_arpe) > 0.3: rpe_delta *= 1.0 + abs(_arpe)
    except Exception:
        pass  # fail-open
    # Item 136: Reward momentum (P33)
    try:
        from pulse.core.reward.reward_momentum import record_success_momentum as _rsm136
        rpe_delta += _rsm136(domain, success)
    except Exception: pass  # fail-open

    try:
        _fail_domain = domain or "general"
        if not success:
            _CONSECUTIVE_FAILURES[_fail_domain] = _CONSECUTIVE_FAILURES.get(_fail_domain, 0) + 1
            if _CONSECUTIVE_FAILURES[_fail_domain] >= _CONSECUTIVE_FAILURE_THRESHOLD:
                log.warning(
                    "Error pattern: %d consecutive failures in domain '%s' — investigate!",
                    _CONSECUTIVE_FAILURES[_fail_domain], _fail_domain
                )
        else:
            _CONSECUTIVE_FAILURES[_fail_domain] = 0  # Reset streak on success
    except Exception:
        pass  # fail-open

    _meta = None
    try:
        from pulse.core.brain_db import get_conn as _gc_meta
        _meta = {}
        _conn_meta = _gc_meta()
        for _iid in retrieved_item_ids:
            for _tbl in ("lessons", "failures", "patterns"):
                _row = _conn_meta.execute(
                    f"SELECT access_count, evidence_count FROM {_tbl} WHERE id = ?", (_iid,)
                ).fetchone()
                if _row:
                    _meta[_iid] = {"access_count": _row[0] or 0, "evidence_count": _row[1] or 0}
                    break
    except Exception:
        _meta = None
    credits = assign_credit(retrieved_item_ids, rpe_delta, item_metadata=_meta)

    items_updated = propagate_reward(credits)

    signal = RewardSignal(
        task_id=task_id, agent=agent, success=success,
        domain=domain, rpe_delta=round(rpe_delta, 4),
        items_credited=credits,
    )

    _REWARD_DIR.mkdir(parents=True, exist_ok=True)
    try:
        from pulse.core.brain_io import _acquire
        with _acquire(REWARD_SIGNALS_FILE):
            signals = []
            if REWARD_SIGNALS_FILE.exists():
                try:
                    signals = json.loads(REWARD_SIGNALS_FILE.read_text(encoding="utf-8"))
                except Exception:
                    signals = []
            signals.append(asdict(signal))
            signals = signals[-MAX_SIGNALS:]
            REWARD_SIGNALS_FILE.write_text(
                json.dumps(signals, indent=2, ensure_ascii=False), encoding="utf-8")
    except Exception as e:
        log.warning("Reward signal logging failed: %s", e)

    _update_credit_ledger(credits)

    # P46.7: Delegate brain learning integrations to pulse.brain.reward_integration
    try:
        from pulse.brain.reward_integration import apply_learning_updates
        apply_learning_updates(
            success=success, domain=domain, agent=agent,
            rpe_delta=rpe_delta, credits=credits,
            retrieved_item_ids=retrieved_item_ids,
            expected_success_rate=expected,
        )
    except Exception as e:
        log.debug("reward_integration delegation failed: %s", e)
    log.info("Reward: task=%s rpe=%.3f credited=%d", task_id, rpe_delta, len(credits))
    return signal
