#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Workspace Integrator: Cross-source synthesis for brain search results.

Based on Global Workspace Theory — takes flat multi-source search results
and produces a coherent synthesized briefing with causal links between items.
Solves the "binding problem": turns isolated memories into understanding.

Two paths:
  - Heuristic (always runs, <10ms, free) — template-based synthesis
  - LLM (opt-in, 1-3s) — deeper semantic synthesis via llm_router

Dependencies: brain_analysis (text_similarity), llm_router (optional)
"""

import hashlib
import time

from pulse.core.brain_analysis import text_similarity, search_relevance

# Source type mapping for template selection
_SOURCE_TYPE = {
    "wins": "lesson", "fails": "failure", "anti_patterns": "warning",
    "patterns": "pattern", "predictions": "prediction", "evidence": "evidence",
    "graph": "causal", "schemas": "schema", "procedures": "procedure",
    "facts": "fact", "episodic": "episode", "decisions": "decision",
    "hot_cache": "recent", "index": "indexed", "working_memory": "session",
    "domains": "domain", "private": "private",
}

# Synthesis templates for known cross-source combinations
_TEMPLATES = {
    ("failure", "warning"): "Known risk: {a}. Guard: {b}.",
    ("failure", "lesson"): "Past failure ({a}) addressed by: {b}.",
    ("lesson", "warning"): "Best practice ({a}) reinforced by anti-pattern: {b}.",
    ("lesson", "pattern"): "Proven approach: {a} (pattern: {b}).",
    ("failure", "causal"): "Root cause: {b} led to failure: {a}.",
    ("prediction", "lesson"): "Predicted risk: {a}. Mitigate with: {b}.",
    ("prediction", "failure"): "Prediction ({a}) echoes past failure: {b}.",
    ("prediction", "warning"): "Predicted risk ({a}) aligns with known trap: {b}.",
    ("lesson", "evidence"): "Verified: {a} (evidence: {b}).",
    ("fact", "lesson"): "Factual basis: {a} — applied as: {b}.",
    ("pattern", "warning"): "Pattern ({a}) conflicts with anti-pattern: {b}.",
    ("episode", "failure"): "Incident: {a} — categorized failure: {b}.",
    ("failure", "decision"): "Past failure ({a}) — check if decision still valid: {b}.",
    ("decision", "warning"): "Decision ({a}) conflicts with anti-pattern: {b}.",
    ("decision", "lesson"): "Decision ({a}) supported by lesson: {b}.",
    ("decision", "prediction"): "Decision ({a}) — predicted outcome: {b}.",
    ("decision", "failure"): "Decision ({a}) — past failure in same area: {b}.",
}

# In-memory synthesis cache (query fingerprint -> lines, with TTL)
_cache: dict = {}
_CACHE_MAX = 50
_CACHE_TTL = 300  # 5 minutes


def integrate(results: dict, query: str, fast_only: bool = True) -> list:
    """Synthesize cross-source connections from brain search results.

    Args:
        results: Raw output from brain_search() — {results: {source: [hits]}}
        query: The original search query
        fast_only: If True, skip LLM synthesis (use for hook path)

    Returns:
        List of synthesis strings (0-5 lines). Empty = no cross-source binding found.
    """
    hits_dict = results.get("results", results)
    if not isinstance(hits_dict, dict):
        return []

    top_items = _top_k_global(hits_dict, k=8)
    if len(top_items) < 3:
        return []

    # Check cache
    cache_key = _cache_key(top_items, query)
    cached = _cache.get(cache_key)
    if cached and (time.time() - cached[1]) < _CACHE_TTL:
        return cached[0]

    clusters = cluster_results(top_items, threshold=0.2, query=query)

    # Filter to multi-source clusters only
    multi_source = [c for c in clusters if len({i["_source_type"] for i in c}) >= 2]
    if not multi_source:
        return []

    lines = heuristic_synthesis(multi_source, query)

    # LLM synthesis (opt-in, non-hook path only) — only replace if >= heuristic count
    if not fast_only and lines:
        try:
            llm_lines = _llm_synthesis(multi_source, query)
            if llm_lines and len(llm_lines) >= len(lines):
                lines = llm_lines
        except Exception:
            pass  # Keep heuristic output

    # Cache and return
    lines = lines[:5]
    _cache_put(cache_key, lines)
    return lines


def cluster_results(items: list, threshold: float = 0.2, query: str = "") -> list:
    """Group items into cross-source clusters for synthesis.

    Strategy: all items already passed brain_search relevance filtering, so they
    share the query topic. We use a hybrid approach:
      1. Pairwise word-overlap (Jaccard) — catches items sharing vocabulary
      2. Query-relevance grouping — items both relevant to the query cluster
         together even with zero shared words (solves the binding problem)
      3. Fallback — if no multi-source clusters form, treat all items as one
         cluster (they all matched the same search query)

    Returns list of clusters (each cluster is a list of item dicts).
    """
    n = len(items)
    if n < 2:
        return [items] if items else []

    # Precompute pairwise similarity: Jaccard + query co-relevance boost
    sims = {}
    for i in range(n):
        for j in range(i + 1, n):
            ti = items[i].get("_text_norm", "")[:80]
            tj = items[j].get("_text_norm", "")[:80]

            # Base: word overlap between items
            jaccard = text_similarity(ti, tj)

            # Boost: rescue pairs below threshold that both match the query
            query_boost = 0.0
            if query and jaccard < threshold:
                ri = search_relevance(query, ti)
                rj = search_relevance(query, tj)
                query_boost = min(ri, rj) * 0.3

            sims[(i, j)] = min(jaccard + query_boost, 1.0)

    # Greedy complete-linkage clustering
    assigned = [False] * n
    clusters = []

    for i in range(n):
        if assigned[i]:
            continue
        cluster = [i]
        assigned[i] = True
        for j in range(i + 1, n):
            if assigned[j]:
                continue
            if all(sims.get((min(k, j), max(k, j)), 0) >= threshold for k in cluster):
                cluster.append(j)
                assigned[j] = True
        clusters.append([items[idx] for idx in cluster])

    # Fallback: if no multi-source clusters, merge all into one mega-cluster
    # (all items matched the same query — they ARE topically related)
    has_multi = any(
        len({it["_source_type"] for it in c}) >= 2
        for c in clusters
    )
    if not has_multi and n >= 3:
        clusters = [items]

    return clusters


def heuristic_synthesis(clusters: list, query: str) -> list:
    """Template-based synthesis for multi-source clusters.

    Picks the best template based on source type combinations.
    Returns formatted synthesis lines.
    """
    lines = []
    for cluster in clusters[:3]:  # Max 3 clusters
        types_in_cluster = {}
        for item in cluster:
            stype = item.get("_source_type", "unknown")
            if stype not in types_in_cluster:
                types_in_cluster[stype] = item

        if len(types_in_cluster) < 2:
            continue

        # Try all type pairs, find best template
        type_keys = list(types_in_cluster.keys())
        best_line = None
        for i in range(len(type_keys)):
            for j in range(i + 1, len(type_keys)):
                pair = (type_keys[i], type_keys[j])
                template = _TEMPLATES.get(pair) or _TEMPLATES.get((pair[1], pair[0]))
                if template:
                    a_text = types_in_cluster[pair[0]].get("_text_short", "?")
                    b_text = types_in_cluster[pair[1]].get("_text_short", "?")
                    best_line = template.format(a=a_text, b=b_text)
                    break
            if best_line:
                break

        if not best_line:
            # Default template for unknown combinations
            items_sorted = sorted(types_in_cluster.items())
            a_type, a_item = items_sorted[0]
            b_type, b_item = items_sorted[1]
            best_line = (f"Cross-ref: {a_item.get('_text_short', '?')} ({a_type}) "
                         f"relates to {b_item.get('_text_short', '?')} ({b_type}).")

        if best_line and len(best_line) > 40:
            lines.append(best_line[:150])

    return lines


def _top_k_global(hits_dict: dict, k: int = 8) -> list:
    """Flatten all sources, tag with source type, return top K by score."""
    all_items = []
    for source_name, hits in hits_dict.items():
        if not isinstance(hits, list):
            continue
        stype = _SOURCE_TYPE.get(source_name, source_name)
        for h in hits:
            if isinstance(h, dict) and not h.get("error"):
                text = h.get("text") or h.get("title") or h.get("description", "")
                text_clean = text[:120].strip()
                if text_clean:
                    item = dict(h)
                    item["_source_type"] = stype
                    item["_text_norm"] = text_clean.lower()
                    item["_text_short"] = text_clean[:60]
                    item["_score"] = h.get("salience", 1.0) * h.get("confidence", 0.5)
                    all_items.append(item)

    all_items.sort(key=lambda x: x["_score"], reverse=True)
    return all_items[:k]


def _cache_key(items: list, query: str) -> str:
    """MD5 of sorted item fingerprints + query for cache lookup."""
    parts = sorted(i.get("_text_norm", "")[:40] for i in items)
    raw = query[:200] + "|" + "|".join(parts)
    return hashlib.md5(raw.encode()).hexdigest()[:16]


def _cache_put(key: str, lines: list) -> None:
    """Store in cache with TTL, evict if over max."""
    global _cache
    if len(_cache) > _CACHE_MAX:
        # Evict oldest half
        by_age = sorted(_cache.items(), key=lambda x: x[1][1])
        _cache = dict(by_age[_CACHE_MAX // 2:])
    _cache[key] = (lines, time.time())


def _llm_synthesis(clusters: list, query: str) -> list:
    """LLM-based cross-source synthesis via llm_router."""
    from pulse.core.llm_router import dispatch

    # Format items for LLM
    formatted = []
    for cluster in clusters[:3]:
        for item in cluster[:4]:
            formatted.append(
                f"[{item.get('_source_type', '?').upper()}] {item.get('_text_short', '?')}")

    prompt = (
        f'Given these brain search results about "{query[:60]}", produce 1-3 '
        f"concise synthesis sentences connecting items from different sources.\n\n"
        f"Results:\n" + "\n".join(formatted) + "\n\n"
        "Rules: reference 2+ items from different sources per sentence. "
        "Note contradictions explicitly. Max 150 chars per sentence. Be specific.\n"
        'Output JSON: {"synthesis": ["sentence1", "sentence2"]}'
    )

    import json
    resp = dispatch(prompt, task_type="fast", max_tokens=300)
    text = resp.get("text", "") if isinstance(resp, dict) else str(resp)

    # Parse JSON from response
    try:
        start = text.index("{")
        end = text.rindex("}") + 1
        data = json.loads(text[start:end])
        return [s[:150] for s in data.get("synthesis", []) if isinstance(s, str) and len(s) > 30]
    except (ValueError, json.JSONDecodeError):
        return []
