#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Phase 1: Reward Signal — Close the Learning Cycle.

When a task completes (success or failure), this module:
1. Computes Reward Prediction Error (RPE): delta = actual - expected
2. Assigns credit/blame to the knowledge items that informed the task
3. Propagates reward → reinforces good knowledge, contradicts bad knowledge
4. Publishes reward.signal event for downstream consumers

Usage:
    from pulse.core.reward_signal import record_task_outcome
    record_task_outcome(task_id="T-abc", agent="claude", success=True,
                        domain="architecture", retrieved_item_ids=["L-123", "F-456"])
"""
import json
import logging
import math
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")

# State files (in state/ subdirectory to avoid compliance validator deletion)
_REWARD_DIR = get_state_dir() / "reward"
REWARD_SIGNALS_FILE = _REWARD_DIR / "reward_signals.json"
CREDIT_LEDGER_FILE = _REWARD_DIR / "credit_ledger.json"

# Limits
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)

# Reward magnitudes
REINFORCE_DELTA = 0.3   # Salience boost on positive credit
CONTRADICT_DELTA = 0.15 # Confidence penalty on negative credit
MIN_CREDIT = 0.05       # Items below this threshold get no update


@dataclass
class RewardSignal:
    task_id: str
    agent: str
    success: bool
    domain: str
    rpe_delta: float  # Reward Prediction Error
    items_credited: dict  # {item_id: credit_amount}
    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 outcome - expected.

    Returns positive delta for surprising success, negative for surprising failure.
    Range: approximately [-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) -> dict[str, float]:
    """Temporal credit assignment with exponential positional decay.

    Items earlier in the list (retrieved first, most influential) get more credit.
    Returns {item_id: signed_credit} where sign matches rpe_delta direction.
    """
    if not item_ids or rpe_delta == 0:
        return {}
    credits = {}
    direction = 1.0 if rpe_delta > 0 else -1.0
    magnitude = abs(rpe_delta)
    for i, item_id in enumerate(item_ids):
        # Exponential decay: first item gets full credit, subsequent decay
        positional = CREDIT_DECAY ** i
        credit = direction * magnitude * positional
        if abs(credit) >= MIN_CREDIT:
            credits[item_id] = round(credit, 4)
    return credits


def propagate_reward(credits: dict[str, float]) -> int:
    """Apply credit/blame to brain knowledge items via reconsolidation.

    Positive credit → reinforce (salience boost)
    Negative credit → contradict (confidence penalty)
    Returns count of items updated.
    """
    if not credits:
        return 0
    updated = 0
    try:
        from pulse.brain.reconsolidation import reinforce, contradict, mark_labile
        from pulse.core.brain_db import get_conn

        conn = get_conn()
        for item_id, credit in credits.items():
            # Find the item's table
            try:
                row = conn.execute(
                    "SELECT table_name FROM knowledge_fts WHERE item_id = ? LIMIT 1",
                    (item_id,)
                ).fetchone()
                if not row:
                    continue
                table = row["table_name"]
            except Exception:
                continue

            if credit > 0:
                # Positive: mark labile then reinforce
                mark_labile(item_id, table)
                if reinforce(item_id):
                    updated += 1
                    log.debug("Reinforced %s (credit=%.3f)", item_id, credit)
            elif credit < 0:
                # Negative: mark labile then contradict
                mark_labile(item_id, table)
                if contradict(item_id):
                    updated += 1
                    log.debug("Contradicted %s (credit=%.3f)", item_id, credit)
    except ImportError:
        log.warning("Reconsolidation module not available for reward propagation")
    except Exception as e:
        log.warning("Reward propagation failed: %s", e)
    return updated


def _get_expected_success_rate(domain: str) -> float:
    """Get expected success rate from metamemory domain confidence."""
    try:
        from pulse.brain.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

            # Trim to MAX_CREDIT_ITEMS (keep highest absolute credit)
            if len(ledger) > MAX_CREDIT_ITEMS:
                sorted_items = sorted(ledger.items(),
                                       key=lambda x: abs(x[1].get("total_credit", 0)),
                                       reverse=True)
                ledger = dict(sorted_items[: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 []

    # 1. Compute RPE
    expected = _get_expected_success_rate(domain)
    rpe_delta = compute_rpe(success, expected)

    # 2. Assign credit
    credits = assign_credit(retrieved_item_ids, rpe_delta)

    # 3. Propagate to brain
    items_updated = propagate_reward(credits)

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

    # 5. Log signal
    _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)

    # 6. Update credit ledger
    _update_credit_ledger(credits)

    # 7. Publish event
    try:
        from pulse.core import event_bus
        event_bus.publish("reward.signal", "reward_signal", {
            "task_id": task_id, "agent": agent, "success": success,
            "rpe_delta": signal.rpe_delta, "items_updated": items_updated,
        })
    except Exception:
        pass

    log.info("Reward signal: task=%s agent=%s success=%s rpe=%.3f items_credited=%d propagated=%d",
             task_id, agent, success, rpe_delta, len(credits), items_updated)
    return signal
