#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Retrieval tracking: LTP metrics, reconsolidation hooks, STDP/Hebbian co-activation.

Single responsibility: record that brain items were retrieved and reinforce them.
"""

import hashlib
import atexit
import logging
import threading
import sys
from datetime import datetime, timezone
from pathlib import Path

# Allow direct execution from repo root
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent))

from pulse.core.brain_utils import RETRIEVAL_METRICS

log = logging.getLogger("pulse.brain_search.retrieval")

# Shutdown guard: skip DB work during process exit to prevent hang
_shutting_down = threading.Event()
atexit.register(_shutting_down.set)


def _strip_epistemic(text: str) -> str:
    """Strip [UNCERTAIN], [STALE Nd], [CONTESTED] prefixes added by epistemic.py."""
    import re
    return re.sub(r'^\s*\[(?:UNCERTAIN|STALE \d+d|CONTESTED)\]\s*', '', text).strip()


def _item_key(text: str) -> str:
    """Stable key from first 80 chars — MD5 hash. Aligned with _entry_hash() in consolidation."""
    return hashlib.md5(text[:80].lower().strip().encode("utf-8")).hexdigest()[:12]


def _get_short_timeout_conn():
    """Get a DB connection with 2s busy_timeout for background work.
    
    The main get_conn() uses 60s timeout which blocks process exit when
    a daemon holds the write lock. Background retrieval tracking must not
    block exit.
    """
    import sqlite3, math
    from pulse.core.brain_db import DB_PATH
    db_path = DB_PATH.resolve()
    conn = sqlite3.connect(str(db_path), timeout=2)
    conn.execute("PRAGMA journal_mode=WAL;")
    conn.execute("PRAGMA synchronous=NORMAL;")
    conn.execute("PRAGMA busy_timeout=2000;")  # 2s — fail fast, this is non-critical
    conn.create_function("SQRT", 1, math.sqrt)
    conn.row_factory = sqlite3.Row
    return conn


def _record_retrievals(results: dict):
    """Record that items were retrieved (increments use counter).

    Runs in a background thread — must bail fast if DB is locked to avoid
    blocking process exit. Total deadline: 3 seconds across all phases.
    """
    if _shutting_down.is_set():
        return  # Skip DB work during process exit to prevent hang
    if not results.get("results"):
        return

    import time
    _deadline = time.monotonic() + 8.0  # 8s total across all phases (runs in daemon thread)

    def _past_deadline():
        return time.monotonic() > _deadline or _shutting_down.is_set()

    now_iso = datetime.now(timezone.utc).isoformat()
    all_texts = []

    # Track top hits for reconsolidation (limit 5)
    reconsolidation_candidates = []

    for source_hits in results["results"].values():
        if not isinstance(source_hits, list):
            continue
        for hit in source_hits:
            text = hit.get("text", "") or hit.get("title", "")
            if text:
                clean = _strip_epistemic(text)
                all_texts.append(clean)
                if len(reconsolidation_candidates) < 25:
                    reconsolidation_candidates.append(clean)

    if not all_texts:
        return

    # Phase 2B: Reconsolidation Hook — mark labile + reinforce on re-retrieval
    # FIX: Use LIKE prefix match instead of exact `content = ?` which fails when
    # search sources truncate text to 200 chars but FTS stores full content.
    # resolved_items: [(item_id, table_name, text)] — text kept for per-item key lookup
    resolved_items: list[tuple[str, str, str]] = []
    if not _past_deadline():
        try:
            from pulse.brain.plasticity.reconsolidation import mark_labile, reinforce
            conn = _get_short_timeout_conn()
            try:
                for text in reconsolidation_candidates:
                    if _past_deadline():
                        break
                    # Use first 80 chars as prefix match — stable across truncation
                    prefix = text[:80].replace("'", "''") if text else ""
                    if len(prefix) < 20:
                        continue
                    res = conn.execute(
                        "SELECT item_id, table_name FROM knowledge_fts WHERE content LIKE ? LIMIT 1",
                        (prefix + "%",)
                    ).fetchone()
                    if res and res["item_id"] and res["table_name"]:
                        resolved_items.append((res["item_id"], res["table_name"], text))
                        if not reinforce(res["item_id"]):
                            mark_labile(res["item_id"], res["table_name"], prediction_error=0.1)
            finally:
                conn.close()
        except Exception as e:
            logging.debug(f"Reconsolidation hook failed: {e}")

    # Phase 2B+: Increment access_count and update confidence on source tables for resolved items.
    # This closes the retrieval feedback loop — access_count > 0 enables the
    # reward-gated ranking boost in search_scoring.py (Item 82).
    _valid_tables = frozenset({
        "lessons", "failures", "patterns", "private",
        "intents", "decisions", "facts", "evidence",
        "anti_patterns", "schemas",
    })
    if not _past_deadline():
        try:
            from pulse.core.uncertainty import bayesian_update
            _conn_acc = _get_short_timeout_conn()   
            try:
                _updated_ids = set()
                with _conn_acc:
                    # Helper to update both access count and confidence
                    def _update_retrieval_stats(item_id, table):
                        row = _conn_acc.execute(f"SELECT confidence, evidence_count FROM {table} WHERE id=?", (item_id,)).fetchone()
                        if row:
                            new_conf, new_ec = bayesian_update(row["confidence"], row["evidence_count"], positive_evidence=True)
                            _conn_acc.execute(
                                f"UPDATE {table} SET "
                                "access_count = access_count + 1, "
                                "last_accessed = ?, "
                                "confidence = ?, "
                                "evidence_count = ? "
                                "WHERE id = ?",
                                (now_iso, new_conf, new_ec, item_id),
                            )

                    # Path 1: resolved_items from reconsolidation (FTS-resolved)
                    for item_id, table_name, _text in resolved_items:
                        if _past_deadline():        
                            break
                        if table_name in _valid_tables and item_id not in _updated_ids:        
                            _update_retrieval_stats(item_id, table_name)
                            _updated_ids.add(item_id)
                    # Path 2: direct IDs from search results (broader coverage)
                    for _rid in results.get("_retrieved_item_ids", [])[:100]:
                        if _past_deadline():        
                            break
                        if _rid and _rid not in _updated_ids:
                            # Resolve table from FTS index
                            _fts_row = _conn_acc.execute(
                                "SELECT table_name FROM knowledge_fts "
                                "WHERE item_id = ? LIMIT 1", (_rid,)
                            ).fetchone()
                            if _fts_row and _fts_row[0] in _valid_tables:
                                _update_retrieval_stats(_rid, _fts_row[0])
                                _updated_ids.add(_rid)
            finally:
                _conn_acc.close()
        except Exception as _ae:
            log.debug("access_count/confidence increment failed: %s", _ae)
    # Phase 3: STDP + Hebbian — record retrieval sequence AND co-activation
    retrieved_ids = []
    if not _past_deadline():
        try:
            from pulse.brain.plasticity.stdp import record_retrieval
            _dt = datetime.now(timezone.utc)
            session_id = f"{_dt.strftime('%Y%m%d_%H')}_{(_dt.minute // 5) * 5:02d}"
            for item_id, _table, _text in resolved_items:
                if _past_deadline():
                    break
                record_retrieval(session_id, item_id)
                retrieved_ids.append(item_id)
        except Exception:
            pass
    if len(retrieved_ids) >= 2 and not _past_deadline():
        try:
            from pulse.brain.plasticity.hebbian_knowledge import record_coactivation
            record_coactivation(retrieved_ids)
        except Exception:
            pass

    # Phase 2C: Grounding — promote frequently-retrieved items to CONFIRMED
    # item_key is derived from each item's own text — no shared-key bug
    if resolved_items and not _past_deadline():
        try:
            from pulse.brain.integrity.grounding import maybe_promote, RETRIEVAL_THRESHOLD
            _conn = _get_short_timeout_conn()
            try:
                for item_id, table, item_text in resolved_items:
                    if _past_deadline():
                        break
                    key = _item_key(item_text)
                    rm = _conn.execute(
                        "SELECT retrieval_count FROM retrieval_metrics WHERE item_key = ?",
                        (key,)
                    ).fetchone()
                    count = (rm["retrieval_count"] or 0) if rm else 0
                    if count >= RETRIEVAL_THRESHOLD:
                        maybe_promote(item_id, table, trigger="retrieval",
                                      evidence={"retrieval_count": count})
            finally:
                _conn.close()
        except Exception as _ge:
            log.debug("Grounding retrieval hook failed: %s", _ge)

    # Dual-write: SQLite (primary) + JSON (backward compat)
    try:
        conn = _get_short_timeout_conn()
        try:
            with conn:
                for text in all_texts:
                    key = _item_key(text)
                    conn.execute(
                        "INSERT INTO retrieval_metrics (item_key, retrieval_count, last_accessed) "
                        "VALUES (?, 1, ?) "
                        "ON CONFLICT(item_key) DO UPDATE SET "
                        "retrieval_count = retrieval_count + 1, last_accessed = ?",
                        (key, now_iso, now_iso)
                    )
                    # Item 94: Retrieval practice — bayesian confidence update on retrieval
                    try:
                        from pulse.core.uncertainty import bayesian_update as _bu94
                        for _tbl94 in ("lessons", "failures", "patterns"):
                            _row94 = conn.execute(
                                f"SELECT id, confidence, access_count FROM {_tbl94} "
                                "WHERE fingerprint IN ("
                                "SELECT fingerprint FROM knowledge_fts WHERE content LIKE ? LIMIT 1)",
                                (text[:80] + "%",)
                            ).fetchone()
                            if _row94:
                                _nc, _ = _bu94(
                                    float(_row94["confidence"] or 0.5),
                                    int(_row94["access_count"] or 1),
                                    True,
                                )
                                conn.execute(
                                    f"UPDATE {_tbl94} SET confidence = ? WHERE id = ?",
                                    (round(_nc, 4), _row94["id"]),
                                )
                    except Exception:
                        pass  # fail-open: non-critical
                    # PP1 (Item 19): Active forgetting — penalize items retrieved 5+ times without reward
                    try:
                        rm = conn.execute(
                            "SELECT retrieval_count FROM retrieval_metrics WHERE item_key = ?",
                            (key,)
                        ).fetchone()
                        rc = (rm["retrieval_count"] if rm else 0)
                        if rc >= 5:
                            for tbl in ("lessons", "failures", "patterns"):
                                conn.execute(
                                    f"UPDATE {tbl} SET confidence = MAX(0.05, confidence - 0.05) "
                                    "WHERE fingerprint IN (SELECT fingerprint FROM knowledge_fts WHERE content LIKE ? LIMIT 1) "
                                    "AND (access_count IS NULL OR access_count < 1)",
                                    (text[:80] + "%",)
                                )
                    except Exception:
                        pass  # fail-open: active forgetting is non-critical
        finally:
            conn.close()
    except Exception as e:
        log.debug("Retrieval metrics SQLite write failed: %s", e)

    try:
        from pulse.core.brain_utils import atomic_update_json
        # P47: JSON backward-compat removed — SQLite retrieval_metrics table is authoritative
        # # P47: SQLite is primary — JSON backward compat disabled
        # atomic_update_json(RETRIEVAL_METRICS, _updater, default={})  # REMOVED: retrieval_metrics table is source of truth
        pass  # REMOVED
    except Exception as e:
        log.debug("Retrieval metrics JSON write failed: %s", e)
