#!/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.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,
    search_meta_schemas, search_capture_log,
)
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

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") -> dict:
    """Hybrid search: fast-path + parallel live + D/E/A/B scoring."""
    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()

    # 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(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(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.rule_embeddings import search_similar, EMBED_INDEX_FILE
            if EMBED_INDEX_FILE.exists():
                embed_hits = search_similar(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)

    # --- Source selection (Item 68: complexity-based) ---
    _is_complex = len(query.split()) >= 5 or any(op in query for op in ["AND", "OR", "NOT", '"'])
    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,
        "meta_schemas": search_meta_schemas,
        "capture_log": search_capture_log,   # N0: conversational memory recall
    }

    # Check if index is stale (>30min behind capture logs)
    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:
                    index_stale = True
        except Exception:
            index_stale = True

    if not results.get("index") or index_stale:
        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:
        search_fns.update({
            "wins": search_wins, "fails": search_fails,
            "patterns": search_patterns, "evidence": search_evidence,
        })

    # Item 68: simple queries -> core sources (P47: always include episodic + working_memory)
    # Fix 1F: also include graph, schemas, procedures for simple queries
    if not _is_complex:
        _simple_sources = (
            "wins", "fails", "patterns", "anti_patterns", "decisions",
            "episodic", "working_memory", "facts", "predictions",
            "graph", "schemas", "procedures", "capture_log",
        )
        search_fns = {k: v for k, v in search_fns.items() if k in _simple_sources}
    if exclude_sources:
        search_fns = {k: v for k, v in search_fns.items() if k not in exclude_sources}

    # Item 129: Knowledge type routing -- boost relevant sources by query intent (P30)
    try:
        from pulse.brain.brain_search.query_routing import boost_sources_by_intent
        search_fns = boost_sources_by_intent(query, search_fns)
    except Exception:
        pass  # fail-open

    # Item 69: cognitive load-adaptive source pruning
    try:
        from pulse.core.interoception import get_signals
        load = get_signals().get("agent_load", 0)
        mx = 5 if load > 10 else (8 if load > 5 else len(search_fns))
        if len(search_fns) > mx:
            priority = ["anti_patterns", "fails", "wins", "patterns",
                        "decisions", "predictions", "graph", "facts"]
            search_fns = {k: v for k, v in search_fns.items() if k in priority[:mx]}
    except Exception: pass

    # --- Parallel fetch ---
    futures = {_EXECUTOR.submit(fn, 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

    # --- Finalize ---
    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
