#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Brain Search Engine: Parallel search orchestrator (V43.13)

Orchestrates parallel search across 15 brain sources.
Returns salience-scored results with recency decay.

Usage:
    python brain_search.py "query text"
    python brain_search.py "query" --json
    python brain_search.py "query" --top 5
Dependencies: Python stdlib + search_sources.py
"""

import atexit
import hashlib
import json
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from pulse.core.brain_utils import load_config, load_json_dict, RETRIEVAL_METRICS, text_similarity
from pulse.brain.search_sources import (
    search_anti_patterns, search_wins, search_fails, search_domains,
    search_patterns, search_evidence, search_private, search_schemas,
    search_procedures, search_graph, search_predictions, search_facts,
    search_episodic, search_working_memory, search_decisions
)
from pulse.brain.hot_cache import search_hot


# === MODULE-LEVEL THREAD POOL (singleton, reused across searches) ===
_EXECUTOR = ThreadPoolExecutor(max_workers=16)
atexit.register(_EXECUTOR.shutdown, wait=False)


# === RETRIEVAL TRACKING (LTP) ===

def _item_key(text: str) -> str:
    """Stable key from first 60 chars — MD5 hash."""
    return hashlib.md5(text[:60].lower().encode("utf-8")).hexdigest()[:12]


def _record_retrievals(results: dict):
    """Record that items were retrieved (increments use counter)."""
    if not results.get("results"):
        return
    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:
                all_texts.append(text)
                if len(reconsolidation_candidates) < 5:
                    reconsolidation_candidates.append(text)

    if not all_texts:
        return
        
    # Phase 2B: Reconsolidation Hook
    try:
        from pulse.brain.reconsolidation import mark_labile
        from pulse.core.brain_db import get_conn
        conn = get_conn()
        for text in reconsolidation_candidates:
            res = conn.execute("SELECT item_id, table_name FROM knowledge_fts WHERE content = ? LIMIT 1", (text,)).fetchone()
            if res and res["item_id"] and res["table_name"]:
                mark_labile(res["item_id"], res["table_name"])
    except Exception as e:
        import logging
        logging.warning(f"Reconsolidation hook failed: {e}")

    def _updater(data):
        if not isinstance(data, dict):
            data = {}
        items = data.setdefault("items", {})
        data["total_queries"] = data.get("total_queries", 0) + 1
        for text in all_texts:
            key = _item_key(text)
            entry = items.setdefault(key, {
                "retrieval_count": 0,
                "first_retrieved": now_iso,
            })
            entry["retrieval_count"] = entry.get("retrieval_count", 0) + 1
            entry["last_retrieved"] = now_iso
        return data

    try:
        from pulse.core.brain_utils import atomic_update_json
        atomic_update_json(RETRIEVAL_METRICS, _updater, default={})
    except Exception:
        pass  # Non-critical — don't break search if metrics write fails


# === PARALLEL SEARCH ===

def brain_search(query: str, config: dict = None, max_per_source: int = None,
                 exclude_sources: set | None = None) -> dict:
    """Hybrid search: inverted index fast-path + parallel live search.

    1. Index fast-path: instant keyword lookup for static sources
    2. Live parallel: dynamic sources (predictions, anti-patterns) + fills gaps
    3. Merge and deduplicate results
    """
    if config is None:
        config = load_config()
    if max_per_source is None:
        max_per_source = config.get("search", {}).get("max_results_per_source", 10)
    timeout = config.get("search", {}).get("thread_timeout_seconds", 5)

    results = {}
    start = time.time()

    # Priority 0: Hot cache — instant (<10ms), catches recent high-salience items
    try:
        hot_hits = search_hot(query, config)
        if hot_hits:
            results["hot_cache"] = hot_hits[:max_per_source]
    except Exception:
        pass  # Never break search

    # Fast-path: inverted index (covers lessons, failures, patterns, knowledge,
    # schemas, procedures, evidence — all static sources)
    try:
        from pulse.brain.search_index import query_index
        index_hits = query_index(query, config, max_results=max_per_source)
        if index_hits:
            results["index"] = index_hits[:max_per_source]
    except Exception:
        pass

    # Live parallel: always run dynamic sources + core freshness sources.
    # Even when index has hits, run live scan for evidence/wins/fails/patterns
    # to catch data newer than the last index build.
    search_fns = {
        "predictions": search_predictions,
        "anti_patterns": search_anti_patterns,
        "graph": search_graph,
        "facts": search_facts,
        "episodic": search_episodic,
        "working_memory": search_working_memory,
        "decisions": search_decisions,
    }

    # Check index staleness: if index is >24h behind latest capture, force full scan
    index_stale = False
    if results.get("index"):
        try:
            from pulse.core.brain_utils import STATE_DIR
            idx_file = STATE_DIR / "search_index.json"
            cap_files = list(STATE_DIR.glob("pulse_captured_*.log"))
            if idx_file.exists() and cap_files:
                idx_mtime = idx_file.stat().st_mtime
                latest_cap = max(f.stat().st_mtime for f in cap_files)
                if (latest_cap - idx_mtime) > 1800:  # 30 minutes
                    index_stale = True
        except Exception:
            index_stale = True  # Assume stale if check fails

    if not results.get("index") or index_stale:
        # No index hits or stale index — full live search
        search_fns.update({
            "wins": search_wins,
            "fails": search_fails,
            "domains": search_domains,
            "patterns": search_patterns,
            "evidence": search_evidence,
            "private": search_private,
            "schemas": search_schemas,
            "procedures": search_procedures,
        })
    else:
        # Index has fresh hits — still scan core sources for recent data
        search_fns.update({
            "wins": search_wins,
            "fails": search_fails,
            "patterns": search_patterns,
            "evidence": search_evidence,
        })
    if exclude_sources:
        search_fns = {k: v for k, v in search_fns.items() if k not in exclude_sources}

    futures = {_EXECUTOR.submit(fn, query, config): name
               for name, fn in search_fns.items()}
    for future in as_completed(futures, timeout=timeout):
        name = futures[future]
        try:
            hits = future.result()
            results[name] = hits[:max_per_source]
        except Exception as e:
            results[name] = [{"error": str(e)}]

    # Dimension A: load contradiction registry for flagging
    _ctr_texts = set()
    try:
        from pulse.admin.audit.truth_verifier import get_registry
        for c in get_registry().get("contradictions", []):
            if not c.get("resolution"):
                for s in ("item_a", "item_b"):
                    _ctr_texts.add(c[s].get("text", "")[:80].lower())
    except Exception:
        pass

    # Dimension B: load agent trust data once for confidence modulation
    _trust_data = {}
    try:
        from pulse.admin.audit.agent_trust import get_trust as _get_trust
        from pulse.admin.audit.agent_trust import _load_trust
        _trust_data = _load_trust().get("agents", {})
    except Exception:
        _get_trust = None

    # Dimension D+E: apply temporal decay + Wilson intervals + filter expired
    from pulse.core.temporal import temporal_confidence
    from pulse.core.uncertainty import wilson_interval
    for source_name in list(results.keys()):
        hits = results[source_name]
        if not isinstance(hits, list):
            continue
        enriched = []
        for h in hits:
            # Skip expired/historical items
            validity = h.get("validity", "valid")
            if validity in ("expired", "historical"):
                continue
            # Apply type-aware temporal decay to confidence
            ts = h.get("timestamp", "")
            if ts:
                ktype = h.get("type", "lesson")
                base_conf = h.get("confidence", 0.5)
                h["confidence"] = temporal_confidence(base_conf, knowledge_type=ktype, timestamp=ts)
            # Attach Wilson interval for ranking
            ec = h.get("evidence_count", 1)
            lo, hi = wilson_interval(h.get("confidence", 0.5), ec)
            h["confidence_interval"] = [lo, hi]
            h["_wilson_lower"] = lo  # For conservative ranking
            # Dimension A: flag contradicted items
            hit_text = (h.get("content", "") or h.get("text", "") or h.get("title", ""))[:80].lower()
            if _ctr_texts and hit_text and any(text_similarity(hit_text, ct) >= 0.5 for ct in _ctr_texts):
                h["_contradiction_flag"] = True
            # Dimension B: modulate confidence by source agent trust
            # Only apply trust to items from known agents, not brain source names
            try:
                if _get_trust is not None and _trust_data:
                    source_agent = h.get("agent", "")
                    if source_agent and source_agent.lower() in _trust_data:
                        trust_score = max(0.1, _get_trust(source_agent))
                        h["confidence"] = h.get("confidence", 0.5) * trust_score
                        h["_source_trust"] = trust_score
                        # Recompute Wilson after trust modulation
                        lo, hi = wilson_interval(h["confidence"], ec)
                        h["confidence_interval"] = [lo, hi]
                        h["_wilson_lower"] = lo
            except Exception:
                pass  # Never crash search for trust lookup
            # Quality gate: drop low temporal-confidence hits
            if h.get("confidence", 0.5) >= 0.4:
                enriched.append(h)
        # Sort by Wilson lower bound (conservative ranking)
        enriched.sort(key=lambda x: x.get("_wilson_lower", 0), reverse=True)
        results[source_name] = enriched

    elapsed_ms = int((time.time() - start) * 1000)
    total = sum(len(v) for v in results.values())
    search_result = {
        "query": query,
        "total_results": total,
        "elapsed_ms": elapsed_ms,
        "results": results
    }

    # Record retrievals for LTP (non-blocking, best-effort)
    _record_retrievals(search_result)

    # Populate working memory (best-effort)
    try:
        from pulse.brain.working_memory import _add_to_session
        _add_to_session("search", [h for src in results.values() for h in (src if isinstance(src, list) else [])])
    except Exception:
        pass

    return search_result


