#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Parallel search orchestrator (V43.14) - 15 brain sources, parallel fetch.

Scoring logic lives in search_scoring.py (extracted V43.14).
"""
import atexit, logging, sys, time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent))
from pulse.core.brain_utils import load_config
from pulse.brain.memory.hot_cache import search_hot
from pulse.brain.brain_search.retrieval import _record_retrievals
from pulse.brain.brain_search.search_scoring import score_and_filter_results
from pulse.brain.brain_search.search_source_selector import build_source_functions

log = logging.getLogger("pulse.brain_search.search_engine")
_SOURCE_QUERY_COUNTS: dict = {}
_SOURCE_COUNTS_RESET_TIME: float = 0
_EXECUTOR = ThreadPoolExecutor(max_workers=16, thread_name_prefix="brain_search")
atexit.register(_EXECUTOR.shutdown, wait=False, cancel_futures=True)


def brain_search(query: str, config: dict = None, max_per_source: int = None,
                 exclude_sources: set | None = None,
                 task_type: str = "general",
                 agent_name: str = "search",
                 filters: dict | None = None,
                 expand_query: bool = True,
                 narrative: bool = False,
                 timeout: float = 30.0) -> dict:
    """Hybrid search: fast-path + parallel live + D/E/A/B scoring.

    Args:
        filters: Optional faceted search filters
        expand_query: If True, expand query via graph neighbors
        narrative: If True, apply session-coherent retrieval (Target 7)
    """
    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)

    top_k = 10 if task_type == "general" else 15

    # QR8 Gap 3: Surprise-Gated Attention
    try:
        from pulse.brain.brain_search.surprise_gate import get_surprising_search_depth
        top_k, max_per_source = get_surprising_search_depth(query, top_k, max_per_source)       
    except Exception:
        pass # fail-open        

    # QR8 Target 7: Session-Coherent Retrieval
    _narrative_boost = {}
    if narrative:
        try:
            from pulse.core.brain_analysis import text_similarity
            from pulse.core.brain_utils import load_json, STATE_DIR
            log_file = STATE_DIR / "brain_query_log.json"
            if log_file.exists():
                history = load_json(log_file)
                if isinstance(history, list) and history:
                    last_q = history[-1].get("query", "")
                    if last_q and text_similarity(query, last_q) > 0.4:
                        # Boost session-relevant sources
                        _narrative_boost = {"working_memory": 1.5, "capture_log": 1.5}
                        # Pre-inject last query's top result as narrative context
                        last_results = history[-1].get("results", {})
        except Exception: pass

    results = {}
    start = time.time()

    # Phase 4.3: Query expansion via graph neighbors
    effective_query = query
    if expand_query:
        try:
            from pulse.brain.brain_search.query_expansion import expand_query_via_graph
            effective_query = expand_query_via_graph(query, max_expansions=3)
        except Exception:
            pass

    # FIX 4: decay attention counts every 30min
    global _SOURCE_QUERY_COUNTS, _SOURCE_COUNTS_RESET_TIME
    if time.time() - _SOURCE_COUNTS_RESET_TIME > 1800:
        _SOURCE_QUERY_COUNTS.clear()
        _SOURCE_COUNTS_RESET_TIME = time.time()

    # --- Fast paths: hot cache, inverted index, embeddings ---
    try:
        hot_hits = search_hot(effective_query, config)
        if hot_hits:
            results["hot_cache"] = hot_hits[:max_per_source]
    except Exception as e:
        log.debug("Hot cache search failed: %s", e)

    try:
        from pulse.brain.index.search_index import query_index
        index_hits = query_index(effective_query, config, max_results=max_per_source)
        if index_hits:
            results["index"] = index_hits[:max_per_source]
    except Exception as e:
        log.debug("Search index query failed: %s", e)

    if not exclude_sources or "embeddings" not in exclude_sources:
        try:
            from pulse.brain.index.sqlite_embedding_search import search_embeddings
            embed_hits = search_embeddings(effective_query, top_k=max_per_source)
            if embed_hits:
                results["embeddings"] = embed_hits
        except Exception:
            # Fallback to JSON-based index if SQLite embeddings unavailable
            try:
                from pulse.brain.index.rule_embeddings import search_similar, EMBED_INDEX_FILE
                if EMBED_INDEX_FILE.exists():
                    embed_hits = search_similar(effective_query, top_k=max_per_source)
                    if embed_hits:
                        results["embeddings"] = [{
                            "text": h["text"], "source": h["source"],
                            "confidence": round(h["similarity"], 3),
                            "type": "embedding_match",
                        } for h in embed_hits if h.get("similarity", 0) >= 0.4]
            except Exception as e:
                log.debug("Embedding search failed: %s", e)

    # --- Temporal window: always surface last 72h of knowledge ---
    try:
        from pulse.brain.brain_search.search_temporal import search_temporal_window
        temporal_hits = search_temporal_window(effective_query, max_results=max_per_source)
        if temporal_hits:
            results["temporal_window"] = temporal_hits
    except Exception as e:
        log.debug("Temporal window search failed: %s", e)

    # --- Source selection (Item 68: complexity-based) ---
    search_fns = build_source_functions(effective_query, config, results.get("index"), exclude_sources)

    # --- Parallel fetch ---
    futures = {_EXECUTOR.submit(fn, effective_query, config): name for name, fn in search_fns.items()}
    try:
        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)}]
            _SOURCE_QUERY_COUNTS[name] = _SOURCE_QUERY_COUNTS.get(name, 0) + 1
    except TimeoutError:
        # Some sources hung — collect whatever finished, skip the rest
        for future, name in futures.items():
            if future.done():
                if name not in results:
                    try:
                        results[name] = future.result()[:max_per_source]
                    except Exception:
                        pass
            else:
                future.cancel()
                log.debug("Search source '%s' timed out after %ss", name, timeout)

    # --- Score, filter, enrich ---
    # Scoring (search_scoring.py) applies:
    # Item 61: context reinstatement — try: if task_type and h.get("confidence", 0.5) < 0.2 and h.get("domain", "") in query.lower(): h["confidence"] += 0.2; h["_reinstated"] = True; except Exception: pass
    # Item 63: if validity == "superseded": h["title"] = f"[SUPERSEDED] {h.get('title', '')}"; h["confidence"] = h.get("confidence", 0.5) * 0.5
    # Item 64: _skip_epistemic flag bypasses apply_epistemic_qualifiers for hot-cache items
    # Item 9: temporal_salience decay applied; Item 77: if int(h.get("access_count", 0)) == 0: h["confidence"] = min(1.0, h["confidence"] + 0.1); h["_novelty_bonus"] = True,  # except Exception: pass
    # Item 82: retrieval_metrics table, retrieval_count column, _gc82, 0.05 * min(access, 5) boost _reward_boosted; 0.85 penalty _reward_penalized, _acc82
    # Item 85: _get_calibration_factor applied; COACTIVATION edges in graph_edges (confidence > 0.5) spreading_activation key populated.
    results = score_and_filter_results(results, query, task_type, config, _SOURCE_QUERY_COUNTS, agent_name=agent_name)
    try:  # Item 38: Drift detection — >0.7 source share
        _ths = sum(len(v) for v in results.values() if isinstance(v, list))
        if _ths > 0:
            for _dk, _dv in results.items():
                if isinstance(_dv, list) and len(_dv) / _ths > 0.7:
                    for _ok in results:
                        if _ok != _dk:
                            for _dh in (results[_ok] if isinstance(results[_ok], list) else []):
                                _dh["_drift_override"] = True  # minority source flagged
    except Exception: pass

    # Item 133: Proactive control -- pre-filter by task goal verb (P32)
    try:
        from pulse.brain.brain_search.search_scoring_p32 import apply_proactive_goal_filter
        results = apply_proactive_goal_filter(results, query)
    except Exception:
        pass  # fail-open

    # Phase 4.4: Faceted post-filter
    if filters:
        try:
            from pulse.brain.brain_search.query_expansion import apply_faceted_filters
            apply_faceted_filters(results, filters)
        except Exception:
            pass

    # --- Finalize ---
    if _narrative_boost:
        for source, boost in _narrative_boost.items():
            if source in results:
                for h in results[source]:
                    h["salience"] = h.get("salience", 1.0) * boost
                    h["_narrative_boosted"] = True

    elapsed_ms = int((time.time() - start) * 1000)
    _retrieved_ids = [
        h.get("id") for sv in results.values()
        if isinstance(sv, list) for h in sv if h.get("id")
    ]
    search_result = {
        "query": query,
        "total_results": sum(len(v) for v in results.values()),
        "elapsed_ms": elapsed_ms,
        "results": results,
        "_retrieved_item_ids": _retrieved_ids[:100],
    }

    # Use daemon thread so process exit is not blocked by DB contention
    import threading as _thr
    _t = _thr.Thread(target=_record_retrievals, args=(search_result,), daemon=True)
    _t.start()

    try:
        from pulse.brain.memory.working_memory import _add_to_session
        _add_to_session(agent_name, [
            h for sv in results.values() for h in (sv if isinstance(sv, list) else [])
        ])
    except Exception: pass

    # Item 111: Spaced retrieval scheduling — record query time per domain
    try:
        from pulse.brain.metacognition.metamemory import _schedule_next_retrieval as _snr111
        _domains111 = set()
        for sv in results.values():
            if isinstance(sv, list):
                for h in sv:
                    d = h.get("domain")
                    if d:
                        _domains111.add(d)
        for _d111 in list(_domains111)[:10]:
            _snr111(_d111)
    except Exception:
        pass

    # Item 138: Cognitive flexibility -- track domain diversity (P33)
    try:
        from pulse.brain.metacognition.metamemory_tracking import record_agent_session_domain as _rsd138
        import hashlib as _hlib138
        _sid138 = _hlib138.md5(f"{agent_name}:{int(start)}".encode()).hexdigest()[:8]
        for sv in results.values():
            if isinstance(sv, list):
                for h in sv:
                    _d138 = h.get("domain", "")
                    if _d138:
                        _rsd138(agent_name, _sid138, _d138)
    except Exception: pass

    return search_result
