#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
RPE-Driven Prompt Adaptation (Phase 4, Architecture 1).

The boot prompt becomes learned, not hardcoded. Items that have historically
helped agents succeed (positive credit in the reward ledger) get priority
in boot briefings. Attention weights become empirical rather than static.

Cold start: Falls back to static weights when insufficient data exists.
Requires ~50+ task outcomes with credit attribution before activation.

Based on: the credit_ledger.json accumulated by reward_signal.py.
"""
import json
import logging
from datetime import datetime, timezone

from pulse.core.paths import get_state_dir

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

LEARNED_WEIGHTS_FILE = get_state_dir() / "reward" / "learned_attention_weights.json"
MIN_DATA_POINTS = 20  # Need this many credited items before adapting
CREDIT_DECAY_PER_DAY = 0.95  # Temporal decay on credit (stale credit fades)
BLEND_RATIO = 0.6  # 60% learned + 40% static when enough data exists


def prioritize_by_credit(items: list[dict], limit: int = 10) -> list[dict]:
    """Rerank brain items by accumulated reward credit.

    Items with positive credit history get moved to the front.
    Items with negative credit get demoted. Unscored items keep
    their original position (recency-based).

    Args:
        items: List of brain items (dicts with 'id' or 'text' keys)
        limit: Max items to return

    Returns: Reranked list, credit-positive items first.
    """
    ledger = _load_credit_ledger()
    if not ledger or len(ledger) < MIN_DATA_POINTS:
        return items[:limit]  # Cold start: no reranking

    # Score each item by credit
    scored = []
    for item in items:
        item_id = item.get("id", "")
        text_key = (item.get("text", "") or item.get("title", ""))[:40].lower()
        # Try direct ID match first, then fuzzy text match
        credit_entry = ledger.get(item_id, {})
        if not credit_entry and text_key:
            # Fuzzy: check if any ledger key contains this text
            for lid, ldata in ledger.items():
                if text_key in lid.lower():
                    credit_entry = ldata
                    break
        total_credit = credit_entry.get("total_credit", 0)
        scored.append((total_credit, item))

    # Sort: positive credit first, then negative, then zero
    scored.sort(key=lambda x: x[0], reverse=True)
    return [item for _, item in scored[:limit]]


def get_learned_weights(task_type: str) -> dict | None:
    """Get empirically learned attention weights for a task type.

    Returns a source->weight dict if sufficient data exists, or None
    if cold start (caller should use static weights).

    Learned from: which sources contributed credit-positive items
    in past tasks of this type.
    """
    if not LEARNED_WEIGHTS_FILE.exists():
        return None
    try:
        data = json.loads(LEARNED_WEIGHTS_FILE.read_text(encoding="utf-8"))
    except Exception:
        return None

    task_weights = data.get(task_type)
    if not task_weights:
        return None
    if task_weights.get("_data_points", 0) < MIN_DATA_POINTS:
        return None

    # Apply temporal decay: stale weights drift toward neutral (1.0)
    # CREDIT_DECAY_PER_DAY^days interpolates learned weight back to baseline
    decay = 1.0
    last_updated = task_weights.get("_last_updated")
    if last_updated:
        try:
            then = datetime.fromisoformat(last_updated)
            days_old = (datetime.now(timezone.utc) - then).total_seconds() / 86400
            decay = CREDIT_DECAY_PER_DAY ** days_old
        except Exception:
            pass

    weights = {}
    for k, v in task_weights.items():
        if not k.startswith("_"):
            # Decay toward 1.0 (neutral baseline) as data ages
            weights[k] = round(1.0 + (v - 1.0) * decay, 3)
    return weights or None


def update_learned_weights(task_type: str, source_credits: dict[str, float]):
    """Update learned attention weights after a task outcome.

    Called by reward_signal after credit assignment. Tracks which search
    sources contributed positively/negatively credited items.

    Args:
        task_type: The task type (debug, architecture, feature, etc.)
        source_credits: {source_name: total_credit_from_that_source}
    """
    if not task_type or not source_credits:
        return

    LEARNED_WEIGHTS_FILE.parent.mkdir(parents=True, exist_ok=True)
    try:
        data = {}
        if LEARNED_WEIGHTS_FILE.exists():
            data = json.loads(LEARNED_WEIGHTS_FILE.read_text(encoding="utf-8"))
    except Exception:
        data = {}

    if task_type not in data:
        data[task_type] = {"_data_points": 0, "_last_updated": None}

    entry = data[task_type]
    entry["_data_points"] = entry.get("_data_points", 0) + 1
    entry["_last_updated"] = datetime.now(timezone.utc).isoformat()

    # Exponential moving average of source credits
    alpha = 0.1  # Learning rate for weight updates
    for source, credit in source_credits.items():
        # Normalize credit to weight adjustment: positive → boost, negative → suppress
        current = entry.get(source, 1.0)
        # Small update toward the credit signal direction
        adjustment = alpha * (1.0 + credit)  # credit > 0 → weight > 1.0
        new_weight = current * (1 - alpha) + adjustment
        entry[source] = round(max(0.3, min(2.5, new_weight)), 3)

    try:
        LEARNED_WEIGHTS_FILE.write_text(
            json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    except Exception as e:
        log.debug("Failed to save learned weights: %s", e)


def update_from_credits(domain: str, credits: dict[str, float]):
    """Called by reward_signal after credit assignment.

    Maps credited items back to search sources and updates learned weights.
    Combines _aggregate_credits_by_source + _infer_task_type + update.
    """
    if not credits:
        return
    _TABLE_TO_SOURCE = {
        "lessons": "wins", "failures": "fails", "patterns": "patterns",
        "facts": "facts", "schemas": "schemas", "procedures": "procedures",
        "decisions": "decisions", "episodic": "episodic",
        "private": "private", "intents": "domains", "evidence": "evidence",
    }
    # Aggregate credits by source — per-item try/except so one bad lookup
    # doesn't abort the entire batch (previously a single FTS miss zeroed totals)
    source_totals = {}
    try:
        from pulse.core.brain_db import get_conn
        conn = get_conn()
    except Exception as e:
        log.warning("update_from_credits: DB unavailable: %s", e)
        return
    # Prefix mapping for file-based sources not in knowledge_fts
    # (schemas/procedures/episodic emit SCH-/PROC-/EP- prefixed IDs)
    _PREFIX_TO_SOURCE = {"SCH-": "schemas", "PROC-": "procedures", "EP-": "episodic"}
    for item_id, credit in credits.items():
        try:
            row = conn.execute(
                "SELECT table_name FROM knowledge_fts WHERE item_id = ? LIMIT 1",
                (item_id,)
            ).fetchone()
            if row:
                source = _TABLE_TO_SOURCE.get(row["table_name"], "domains")
                source_totals[source] = source_totals.get(source, 0) + credit
            else:
                # Fallback: file-based sources use prefix-encoded IDs
                for prefix, source in _PREFIX_TO_SOURCE.items():
                    if item_id.startswith(prefix):
                        source_totals[source] = source_totals.get(source, 0) + credit
                        break
        except Exception as e:
            log.debug("FTS lookup failed for %s: %s", item_id, e)
            continue
    if not source_totals:
        return
    # Infer task type from domain
    d = (domain or "").lower()
    if any(w in d for w in ("debug", "fix", "bug")):
        task_type = "debug"
    elif any(w in d for w in ("arch", "design", "system")):
        task_type = "architecture"
    elif any(w in d for w in ("feature", "impl", "build")):
        task_type = "feature"
    elif any(w in d for w in ("review", "audit", "security")):
        task_type = "review"
    else:
        task_type = "general"
    update_learned_weights(task_type, source_totals)


def resolve_items_from_text(text: str) -> list[str]:
    """Resolve brain item IDs from text via FTS MATCH + LIKE fallback.

    Used by pulse_save.py when working memory is empty — finds items
    related to the task outcome/title for credit assignment.
    """
    if not text or len(text) < 5:
        return []
    try:
        from pulse.core.brain_db import get_conn
        conn = get_conn()
        words = [w for w in text.lower().split() if len(w) > 4][:8]
        item_ids = []
        for word in words:
            found = False
            try:
                rows = conn.execute(
                    "SELECT item_id FROM knowledge_fts WHERE content MATCH ? LIMIT 3",
                    (word,)).fetchall()
                for r in rows:
                    if r["item_id"] and r["item_id"] not in item_ids:
                        item_ids.append(r["item_id"])
                        found = True
            except Exception:
                pass
            if not found:
                for tbl in ("lessons", "failures", "patterns"):
                    try:
                        rows = conn.execute(
                            f"SELECT id FROM {tbl} WHERE content LIKE ? LIMIT 2",
                            (f"%{word}%",)).fetchall()
                        for r in rows:
                            if r["id"] not in item_ids:
                                item_ids.append(r["id"])
                    except Exception:
                        pass
        return item_ids[:20]
    except Exception:
        return []


def _load_credit_ledger() -> dict:
    """Load the credit ledger from reward_signal."""
    try:
        from pulse.core.reward.reward_signal import CREDIT_LEDGER_FILE
        if not CREDIT_LEDGER_FILE.exists():
            return {}
        return json.loads(CREDIT_LEDGER_FILE.read_text(encoding="utf-8"))
    except Exception:
        return {}
