#!/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 10 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
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
)


# === MODULE-LEVEL THREAD POOL (singleton, reused across searches) ===
_EXECUTOR = ThreadPoolExecutor(max_workers=10)
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 = []
    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 not all_texts:
        return

    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()

    # 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 + sources that need fresh data
    # If index returned results, skip redundant static sources
    search_fns = {
        "predictions": search_predictions,
        "anti_patterns": search_anti_patterns,
        "graph": search_graph,
    }
    # If index had no hits, fall back to full live search
    if not results.get("index"):
        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,
        })
    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)}]

    # Quality filter: drop low-confidence hits (purification gate)
    for source_name in list(results.keys()):
        hits = results[source_name]
        if isinstance(hits, list):
            results[source_name] = [
                h for h in hits
                if h.get("confidence", 0.5) >= 0.4
            ]

    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)

    return search_result


# === CLI ===

def main():
    if len(sys.argv) < 2:
        print("Usage: brain_search.py \"query\" [--json] [--top N]")
        return

    query = sys.argv[1]
    as_json = "--json" in sys.argv
    exclude_private = "--exclude-private" in sys.argv
    top_n = 10
    if "--top" in sys.argv:
        idx = sys.argv.index("--top")
        if idx + 1 < len(sys.argv):
            top_n = int(sys.argv[idx + 1])

    print(f"=== Brain Search V43.13: \"{query}\" ===")
    excl = {"private"} if exclude_private else None
    result = brain_search(query, exclude_sources=excl)
    print(f"  Found {result['total_results']} results in {result['elapsed_ms']}ms")

    if as_json:
        print(json.dumps(result, indent=2))
    else:
        for source, hits in result["results"].items():
            if hits:
                print(f"\n  [{source.upper()}] ({len(hits)} hits)")
                for h in hits[:top_n]:
                    sal = h.get("salience", 0)
                    txt = h.get("text", "")[:80]
                    print(f"    [{sal:.3f}] {txt}")


if __name__ == "__main__":
    main()
