#!/usr/bin/env python3
"""
RLM Bloodstream — Session Context Injector
==========================================
On session start:
1. Queries Qdrant for top-20 most relevant memories based on current task
2. Queries PostgreSQL for recent session state (last 48h entities)
3. Formats injected context as structured markdown
4. Writes to E:\\genesis-system\\data\\context_state\\injected_context.md

Usage:
    python E:\\genesis-system\\core\\session_context_injector.py
    python E:\\genesis-system\\core\\session_context_injector.py --task "Build voice widget for Telnyx"
    python E:\\genesis-system\\core\\session_context_injector.py --top-k 20 --task "GHL modules"

Author: Genesis Systems Builder
Created: 2026-02-23
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Optional

# ── Path setup ────────────────────────────────────────────────────────────────
GENESIS_ROOT = Path("/mnt/e/genesis-system")
DATA_DIR = GENESIS_ROOT / "data"
CONTEXT_STATE_DIR = DATA_DIR / "context_state"
OUTPUT_FILE = CONTEXT_STATE_DIR / "injected_context.md"
KG_ENTITIES_DIR = GENESIS_ROOT / "KNOWLEDGE_GRAPH" / "entities"

sys.path.insert(0, str(GENESIS_ROOT / "data" / "genesis-memory"))
sys.path.insert(0, str(GENESIS_ROOT))

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("ContextInjector")

# ── Default task if none provided ─────────────────────────────────────────────
DEFAULT_TASK = "Genesis session start — load relevant operational memory"

# ── Pattern type display labels ───────────────────────────────────────────────
TYPE_LABELS = {
    "error_fix": "Error Resolution",
    "decision": "Decision",
    "tool_use": "Tool Pattern",
    "learning": "Learning",
    "titan_memory": "Titan Memory",
    "axiom": "Axiom",
}


# ── PostgreSQL helpers ────────────────────────────────────────────────────────
def get_pg_connection():
    try:
        from elestio_config import PostgresConfig
        import psycopg2
        return psycopg2.connect(**PostgresConfig.get_connection_params())
    except ImportError:
        import psycopg2
        return psycopg2.connect(
            host=os.environ.get("PG_HOST", "localhost"),
            port=int(os.environ.get("PG_PORT", 5432)),
            dbname=os.environ.get("PG_DBNAME", "genesis"),
            user=os.environ.get("PG_USER", "genesis"),
            password=os.environ.get("PG_PASSWORD", ""),
        )


def query_recent_pg_entities(conn, hours: int = 48,
                              limit: int = 30) -> List[Dict[str, Any]]:
    """Pull recent high-confidence entities from PostgreSQL digestion table."""
    cutoff = (datetime.utcnow() - timedelta(hours=hours)).isoformat()
    try:
        cur = conn.cursor()
        cur.execute(
            """
            SELECT entity_id, pattern_type, content, confidence, source_file, created_at
            FROM digestion_kg_entities
            WHERE confidence >= 0.75 AND created_at >= %s
            ORDER BY confidence DESC, created_at DESC
            LIMIT %s
            """,
            (cutoff, limit)
        )
        rows = cur.fetchall()
        cur.close()
        return [
            {
                "entity_id": r[0],
                "pattern_type": r[1],
                "content": r[2],
                "confidence": r[3],
                "source_file": r[4],
                "created_at": str(r[5]),
                "source": "postgresql_recent",
            }
            for r in rows
        ]
    except Exception as e:
        logger.warning(f"PostgreSQL recent entity query failed: {e}")
        return []


def query_bloodstream_knowledge(conn, task: str,
                                 limit: int = 10) -> List[Dict[str, Any]]:
    """Pull from bloodstream_knowledge table (KG entities loaded by pipeline)."""
    try:
        cur = conn.cursor()
        # Simple keyword search across content
        keywords = [w.lower() for w in task.split() if len(w) > 3][:5]
        if not keywords:
            keywords = ["genesis"]

        like_clauses = " OR ".join(
            [f"LOWER(content) LIKE %s" for _ in keywords]
        )
        params = [f"%{k}%" for k in keywords] + [limit]

        cur.execute(
            f"""
            SELECT source, type, title, content, confidence, created_at
            FROM bloodstream_knowledge
            WHERE ({like_clauses})
            ORDER BY confidence DESC, created_at DESC
            LIMIT %s
            """,
            params
        )
        rows = cur.fetchall()
        cur.close()
        return [
            {
                "source": r[0],
                "pattern_type": r[1],
                "title": r[2],
                "content": r[3][:400],
                "confidence": r[4],
                "created_at": str(r[5]),
                "source_table": "bloodstream_knowledge",
            }
            for r in rows
        ]
    except Exception as e:
        logger.warning(f"bloodstream_knowledge query failed: {e}")
        return []


# ── Qdrant semantic search ─────────────────────────────────────────────────────
def query_qdrant_semantic(task: str, top_k: int = 20) -> List[Dict[str, Any]]:
    """Semantic search over genesis_memories collection in Qdrant."""
    try:
        from elestio_config import QdrantConfig
        from qdrant_client import QdrantClient

        cfg = QdrantConfig.get_client_params()
        client = QdrantClient(**cfg)

        COLLECTION = "genesis_memories"

        # Check collection exists
        collections = [c.name for c in client.get_collections().collections]
        if COLLECTION not in collections:
            logger.info(f"Qdrant collection '{COLLECTION}' not yet populated")
            return []

        # Embed the task using Gemini (3072d — matches genesis_memories collection)
        try:
            from core.kb.embedder import embed_text
            query_vec = embed_text(task[:512])
        except Exception as exc:
            logger.warning(f"Gemini embedding failed — skipping Qdrant query: {exc}")
            return []

        response = client.query_points(
            collection_name=COLLECTION,
            query=query_vec,
            limit=top_k,
            with_payload=True,
        )
        results = response.points if hasattr(response, "points") else response

        return [
            {
                "entity_id": r.payload.get("entity_id", ""),
                "pattern_type": r.payload.get("pattern_type", "semantic"),
                "content": r.payload.get("content", ""),
                "confidence": round(r.score, 3),
                "source_file": r.payload.get("source_file", ""),
                "created_at": r.payload.get("created_at", ""),
                "source": "qdrant_semantic",
            }
            for r in results
        ]

    except Exception as e:
        logger.warning(f"Qdrant semantic search failed: {e}")
        return []


# ── KG entity file fallback ───────────────────────────────────────────────────
def load_kg_entity_files(limit: int = 20) -> List[Dict[str, Any]]:
    """Fallback: read recent JSONL entity files from KG directory."""
    if not KG_ENTITIES_DIR.exists():
        return []

    entity_files = sorted(
        KG_ENTITIES_DIR.glob("*.jsonl"),
        key=lambda f: f.stat().st_mtime,
        reverse=True
    )[:5]  # Last 5 files

    entities = []
    for ef in entity_files:
        try:
            for line in ef.read_text(encoding="utf-8").splitlines():
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line)
                    obj["source"] = f"kg_file:{ef.name}"
                    entities.append(obj)
                except json.JSONDecodeError:
                    pass
        except Exception as e:
            logger.debug(f"Could not read {ef}: {e}")

    # Sort by confidence desc, take limit
    entities.sort(key=lambda e: e.get("confidence", 0.5), reverse=True)
    return entities[:limit]


# ── Titan Memory loader ───────────────────────────────────────────────────────
def load_titan_memory() -> List[Dict[str, Any]]:
    """Load recent Titan Memory learnings from axioms JSONL."""
    axioms_dir = GENESIS_ROOT / "KNOWLEDGE_GRAPH" / "axioms"
    if not axioms_dir.exists():
        return []

    axiom_files = sorted(
        axioms_dir.glob("*.jsonl"),
        key=lambda f: f.stat().st_mtime,
        reverse=True
    )[:3]

    axioms = []
    for af in axiom_files:
        try:
            for line in af.read_text(encoding="utf-8").splitlines():
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line)
                    obj["source"] = f"axiom:{af.name}"
                    obj["pattern_type"] = "axiom"
                    axioms.append(obj)
                except json.JSONDecodeError:
                    pass
        except Exception as e:
            logger.debug(f"Could not read axiom file {af}: {e}")

    axioms.sort(key=lambda e: e.get("confidence", 0.8), reverse=True)
    return axioms[:10]


# ── Context formatter ─────────────────────────────────────────────────────────
def format_injected_context(
    task: str,
    qdrant_results: List[Dict[str, Any]],
    pg_recent: List[Dict[str, Any]],
    bloodstream: List[Dict[str, Any]],
    kg_fallback: List[Dict[str, Any]],
    titan: List[Dict[str, Any]],
    generated_at: str,
) -> str:
    """Render all memory sources into a structured markdown context block."""

    lines = [
        "# Injected Session Context",
        f"> Generated: {generated_at}",
        f"> Task Signal: {task}",
        "",
        "---",
        "",
    ]

    # Section 1: Titan Memory (always show — highest signal)
    if titan:
        lines.append("## Titan Memory — Core Operational Learnings")
        lines.append("")
        for i, e in enumerate(titan[:8], 1):
            content = e.get("content") or e.get("axiom") or e.get("learning", "")
            label = e.get("id") or e.get("title") or f"axiom_{i}"
            conf = e.get("confidence", 0.8)
            lines.append(f"**{label}** (conf: {conf:.2f})")
            lines.append(f"> {content[:300]}")
            lines.append("")
        lines.append("---")
        lines.append("")

    # Section 2: Semantic search results from Qdrant
    if qdrant_results:
        lines.append("## Top Semantic Memories (Qdrant)")
        lines.append(f"*Relevance-ranked to task: \"{task[:80]}\"*")
        lines.append("")
        for i, e in enumerate(qdrant_results[:10], 1):
            label = TYPE_LABELS.get(e.get("pattern_type", ""), e.get("pattern_type", "Memory"))
            conf = e.get("confidence", 0.0)
            content = e.get("content", "")[:300]
            lines.append(f"**{i}. [{label}]** (relevance: {conf:.3f})")
            lines.append(f"> {content}")
            lines.append("")
        lines.append("---")
        lines.append("")

    # Section 3: Recent PostgreSQL entities
    if pg_recent:
        lines.append("## Recent Session Learnings (Last 48h)")
        lines.append("")
        by_type: Dict[str, List] = {}
        for e in pg_recent[:15]:
            pt = e.get("pattern_type", "other")
            by_type.setdefault(pt, []).append(e)

        for pt, items in by_type.items():
            label = TYPE_LABELS.get(pt, pt.title())
            lines.append(f"### {label}")
            for e in items[:4]:
                content = e.get("content", "")[:250]
                lines.append(f"- {content}")
            lines.append("")
        lines.append("---")
        lines.append("")

    # Section 4: Bloodstream knowledge (task-relevant)
    if bloodstream:
        lines.append("## Task-Relevant Bloodstream Knowledge")
        lines.append("")
        for e in bloodstream[:8]:
            title = e.get("title") or e.get("pattern_type", "entry")
            content = e.get("content", "")[:300]
            conf = e.get("confidence", 0.7)
            lines.append(f"**{title}** (conf: {conf:.2f})")
            lines.append(f"> {content}")
            lines.append("")
        lines.append("---")
        lines.append("")

    # Section 5: KG file fallback
    if kg_fallback and not (qdrant_results or pg_recent):
        lines.append("## KG Entity Fallback (File-Based)")
        lines.append("")
        for e in kg_fallback[:10]:
            content = e.get("content") or e.get("learning", "")
            pt = e.get("type") or e.get("pattern_type", "entity")
            lines.append(f"- **[{pt}]** {str(content)[:250]}")
        lines.append("")
        lines.append("---")
        lines.append("")

    # Footer
    total = len(qdrant_results) + len(pg_recent) + len(bloodstream) + len(titan)
    lines.append(f"*Total memories injected: {total} | Sources: Qdrant, PostgreSQL, KG files, Titan*")
    lines.append("")

    return "\n".join(lines)


# ── Unified Nervous System integration ────────────────────────────────────────

def build_retriever():
    """Build the unified RetrieverOrchestrator with all 3 sources."""
    from core.nervous_system.orchestrator import RetrieverOrchestrator
    from core.nervous_system.qdrant_retriever import QdrantRetriever
    from core.nervous_system.pg_retriever import PGRetriever
    from core.nervous_system.kg_file_retriever import KGFileRetriever

    orch = RetrieverOrchestrator(timeout_ms=8000)
    orch.register_retriever(QdrantRetriever())
    orch.register_retriever(PGRetriever())
    orch.register_retriever(KGFileRetriever())
    return orch


def run_unified(task: str = DEFAULT_TASK, top_k: int = 20) -> Dict[str, Any]:
    """Run context injection via the unified Nervous System retriever."""
    from core.nervous_system.contracts import RetrievalRequest

    generated_at = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
    logger.info(f"=== Session Context Injector (Unified) === task='{task[:60]}...'")

    orch = build_retriever()
    result = orch.query(RetrievalRequest(query=task, top_k=top_k))

    # Titan Memory always loaded separately (file-based, not in orchestrator)
    titan = load_titan_memory()

    # Format output
    lines = [
        "# Genesis Session Context (Injected)",
        f"*Generated: {generated_at} | Task: {task[:80]}*",
        "",
    ]

    # Titan section
    if titan:
        lines.append("## Titan Memory (Hardwired Learnings)")
        for t in titan[:10]:
            lines.append(f"- **{t.get('id', 'unknown')}**: {t.get('content', t.get('axiom', ''))[:200]}")
        lines.append("")

    # Unified retrieval results grouped by source
    source_groups: Dict[str, list] = {}
    for chunk in result.chunks:
        prefix = chunk.source.split(":")[0]
        source_groups.setdefault(prefix, []).append(chunk)

    if "qdrant" in source_groups:
        lines.append("## Semantic Memory (Qdrant)")
        for c in source_groups["qdrant"][:10]:
            lines.append(f"- [score={c.relevance_score:.2f}] {c.content[:200]}")
        lines.append("")

    if "pg" in source_groups:
        lines.append("## Recent Entities (PostgreSQL)")
        for c in source_groups["pg"][:10]:
            lines.append(f"- [{c.source}] {c.content[:200]}")
        lines.append("")

    if "kg" in source_groups:
        lines.append("## Knowledge Graph (Files)")
        for c in source_groups["kg"][:10]:
            lines.append(f"- [{c.source}] {c.content[:200]}")
        lines.append("")

    total = result.total_chunks + len(titan)
    lines.append(f"*Total: {total} | Sources: {', '.join(result.sources_queried)} | Latency: {result.latency_ms:.0f}ms*")

    context_md = "\n".join(lines)
    CONTEXT_STATE_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_FILE.write_text(context_md, encoding="utf-8")
    logger.info(f"Context written to {OUTPUT_FILE}")

    stats = {
        "generated_at": generated_at,
        "task": task,
        "total_chunks": result.total_chunks,
        "titan": len(titan),
        "sources_queried": result.sources_queried,
        "sources_failed": result.sources_failed,
        "latency_ms": result.latency_ms,
        "output_file": str(OUTPUT_FILE),
    }

    manifest_file = CONTEXT_STATE_DIR / "injected_context_manifest.json"
    manifest_file.write_text(json.dumps(stats, indent=2, default=str), encoding="utf-8")

    return stats


# ── Legacy pipeline (kept for backward compat) ──────────────────────────────
def run(task: str = DEFAULT_TASK, top_k: int = 20,
        pg_hours: int = 48) -> Dict[str, Any]:
    """Run the context injection pipeline."""
    generated_at = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
    logger.info(f"=== Session Context Injector === task='{task[:60]}...'")

    stats = {
        "generated_at": generated_at,
        "task": task,
        "qdrant_results": 0,
        "pg_recent": 0,
        "bloodstream": 0,
        "kg_fallback": 0,
        "titan": 0,
        "output_file": str(OUTPUT_FILE),
    }

    # 1. Titan Memory (always first)
    titan = load_titan_memory()
    stats["titan"] = len(titan)
    logger.info(f"Titan Memory: {len(titan)} axioms")

    # 2. Qdrant semantic search
    qdrant_results = query_qdrant_semantic(task, top_k=top_k)
    stats["qdrant_results"] = len(qdrant_results)
    logger.info(f"Qdrant: {len(qdrant_results)} semantic results")

    # 3. PostgreSQL recent entities
    conn = None
    pg_recent = []
    bloodstream = []
    try:
        conn = get_pg_connection()
        pg_recent = query_recent_pg_entities(conn, hours=pg_hours)
        bloodstream = query_bloodstream_knowledge(conn, task)
        stats["pg_recent"] = len(pg_recent)
        stats["bloodstream"] = len(bloodstream)
        logger.info(f"PostgreSQL: {len(pg_recent)} recent, {len(bloodstream)} bloodstream")
    except Exception as e:
        logger.warning(f"PostgreSQL unavailable: {e}")
    finally:
        if conn:
            conn.close()

    # 4. KG file fallback
    kg_fallback = []
    if not qdrant_results and not pg_recent:
        kg_fallback = load_kg_entity_files()
        stats["kg_fallback"] = len(kg_fallback)
        logger.info(f"KG fallback: {len(kg_fallback)} entities")

    # 5. Format and write
    context_md = format_injected_context(
        task=task,
        qdrant_results=qdrant_results,
        pg_recent=pg_recent,
        bloodstream=bloodstream,
        kg_fallback=kg_fallback,
        titan=titan,
        generated_at=generated_at,
    )

    CONTEXT_STATE_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_FILE.write_text(context_md, encoding="utf-8")
    logger.info(f"Context written to {OUTPUT_FILE}")

    # Also write JSON manifest
    manifest_file = CONTEXT_STATE_DIR / "injected_context_manifest.json"
    manifest_file.write_text(
        json.dumps(stats, indent=2, default=str),
        encoding="utf-8"
    )

    return stats


# ── CLI ───────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="RLM Session Context Injector")
    parser.add_argument("--task", default=DEFAULT_TASK,
                        help="Current task description for relevance ranking")
    parser.add_argument("--top-k", type=int, default=20,
                        help="Number of results to retrieve")
    parser.add_argument("--legacy", action="store_true",
                        help="Use legacy pipeline instead of unified nervous system")
    parser.add_argument("--pg-hours", type=int, default=48,
                        help="(Legacy only) Hours back for PostgreSQL query")
    args = parser.parse_args()

    if args.legacy:
        result = run(task=args.task, top_k=args.top_k, pg_hours=args.pg_hours)
    else:
        result = run_unified(task=args.task, top_k=args.top_k)
    print(json.dumps(result, indent=2, default=str))