#!/usr/bin/env python3
"""
RLM Bloodstream - Session Context Injector
===========================================
Reads top-N most relevant KG entities and axioms, formats as compact context block,
and injects into session memory for Claude/Gemini agents.

The "Bloodstream" = living memory that flows through every agent session.

Architecture:
- Reads JSONL files from KNOWLEDGE_GRAPH/entities/ and axioms/
- Scores entities by relevance (recency, frequency, semantic match to task)
- Formats top-50 entities into compact <5K token context block
- Writes to .gemini/knowledge/bloodstream_context.md for auto-injection

Usage:
    python session_context_injector.py --task "Build voice agent for plumbers"
    python session_context_injector.py --max-entities 100 --output custom.md
"""

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

# =============================================================================
# Configuration
# =============================================================================

GENESIS_ROOT = Path(__file__).parent.parent.parent
KG_DIR = GENESIS_ROOT / "KNOWLEDGE_GRAPH"
ENTITIES_DIR = KG_DIR / "entities"
AXIOMS_DIR = KG_DIR / "axioms"
OUTPUT_DIR = GENESIS_ROOT / ".gemini" / "knowledge"
OUTPUT_FILE = OUTPUT_DIR / "bloodstream_context.md"

MAX_CONTEXT_TOKENS = 5000  # Target token budget
MAX_ENTITIES = 50          # Top N entities to include
ENTITY_DECAY_DAYS = 90     # Older entities decay in relevance

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("bloodstream-injector")

# =============================================================================
# Helper Functions
# =============================================================================

def load_jsonl_files(directory: Path) -> List[Dict[str, Any]]:
    """Load all JSONL files from a directory."""
    entities = []
    if not directory.exists():
        logger.warning(f"Directory does not exist: {directory}")
        return entities

    for jsonl_file in directory.glob("*.jsonl"):
        try:
            with open(jsonl_file, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            entity = json.loads(line)
                            entity["source_file"] = jsonl_file.name
                            entities.append(entity)
                        except json.JSONDecodeError as e:
                            logger.warning(f"JSON decode error in {jsonl_file.name}: {e}")
        except Exception as e:
            logger.error(f"Error reading {jsonl_file}: {e}")

    logger.info(f"Loaded {len(entities)} entities from {directory}")
    return entities

def score_entity_relevance(
    entity: Dict[str, Any],
    task_keywords: Optional[List[str]] = None
) -> float:
    """
    Score entity relevance for current session.

    Scoring factors:
    - Recency (newer = higher score)
    - Keyword match to current task
    - Entity type priority (axioms > entities)
    - Frequency (how often referenced in relationships)
    """
    score = 0.0

    # Recency score (0.0 - 1.0)
    timestamp = entity.get("timestamp") or entity.get("created_at") or entity.get("createdAt")
    if timestamp:
        try:
            created = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            age_days = (datetime.now() - created.replace(tzinfo=None)).days
            recency_score = max(0.0, 1.0 - (age_days / ENTITY_DECAY_DAYS))
            score += recency_score * 0.3
        except Exception as e:
            logger.debug(f"Could not parse timestamp: {timestamp}")

    # Keyword match score (0.0 - 1.0)
    if task_keywords:
        text = str(entity).lower()
        matches = sum(1 for kw in task_keywords if kw.lower() in text)
        keyword_score = min(1.0, matches / len(task_keywords))
        score += keyword_score * 0.5

    # Entity type priority
    entity_type = entity.get("entityType") or entity.get("type") or ""
    if "axiom" in entity_type.lower():
        score += 0.2  # Axioms are high-value knowledge

    # Base score for all entities
    score += 0.1

    return score

def deduplicate_entities(entities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Remove duplicate entities by ID or name."""
    seen_ids = set()
    seen_names = set()
    unique = []

    for entity in entities:
        entity_id = entity.get("id") or entity.get("entityId")
        entity_name = entity.get("name") or entity.get("title")

        # Use ID if available, otherwise name
        key = entity_id if entity_id else entity_name

        if key and key not in seen_ids:
            seen_ids.add(key)
            if entity_name:
                seen_names.add(entity_name)
            unique.append(entity)

    logger.info(f"Deduplicated {len(entities)} → {len(unique)} entities")
    return unique

def format_entity_for_context(entity: Dict[str, Any], index: int) -> str:
    """Format a single entity for the context block."""
    entity_type = entity.get("entityType") or entity.get("type") or "unknown"
    name = entity.get("name") or entity.get("title") or f"Entity-{index}"

    # Extract key info
    observations = entity.get("observations", [])
    description = entity.get("description") or entity.get("content") or ""
    timestamp = entity.get("timestamp") or entity.get("created_at") or ""

    # Compact format
    lines = [f"### {index}. {name} ({entity_type})"]

    if description:
        lines.append(f"**Description**: {description[:200]}...")

    if observations:
        obs_text = "; ".join(observations[:3])  # Top 3 observations
        lines.append(f"**Key Points**: {obs_text[:300]}...")

    if timestamp:
        lines.append(f"**Date**: {timestamp[:10]}")

    return "\n".join(lines)

def generate_bloodstream_context(
    entities: List[Dict[str, Any]],
    axioms: List[Dict[str, Any]],
    task: Optional[str] = None,
    max_entities: int = MAX_ENTITIES
) -> str:
    """Generate the full bloodstream context block."""

    # Parse task keywords
    task_keywords = []
    if task:
        task_keywords = [w.strip() for w in task.lower().split() if len(w) > 3]

    # Score and sort entities
    all_knowledge = entities + axioms
    all_knowledge = deduplicate_entities(all_knowledge)

    for item in all_knowledge:
        item["_relevance_score"] = score_entity_relevance(item, task_keywords)

    all_knowledge.sort(key=lambda x: x["_relevance_score"], reverse=True)

    # Select top N
    top_knowledge = all_knowledge[:max_entities]

    # Build context markdown
    lines = [
        "# Genesis Bloodstream Context",
        f"**Generated**: {datetime.now().isoformat()}",
        f"**Task Context**: {task or 'General session'}",
        f"**Entities Loaded**: {len(top_knowledge)}",
        "",
        "---",
        "",
        "## Active Knowledge Base",
        ""
    ]

    for i, item in enumerate(top_knowledge, 1):
        lines.append(format_entity_for_context(item, i))
        lines.append("")

    lines.extend([
        "---",
        "",
        "## Usage Notes",
        "- This context is auto-injected at session start",
        "- Entities scored by relevance (recency + keyword match + type priority)",
        "- Refresh rate: Updated every session or on-demand",
        "- Source: KNOWLEDGE_GRAPH/entities/ + axioms/",
        ""
    ])

    return "\n".join(lines)

# =============================================================================
# Main Entry Point
# =============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="RLM Bloodstream Session Context Injector"
    )
    parser.add_argument(
        "--task",
        type=str,
        default=None,
        help="Current task/session context for relevance scoring"
    )
    parser.add_argument(
        "--max-entities",
        type=int,
        default=MAX_ENTITIES,
        help=f"Maximum entities to include (default: {MAX_ENTITIES})"
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=OUTPUT_FILE,
        help=f"Output file path (default: {OUTPUT_FILE})"
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging"
    )

    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)

    logger.info("Starting RLM Bloodstream context injection...")
    logger.info(f"Task context: {args.task or 'None (general session)'}")
    logger.info(f"Max entities: {args.max_entities}")

    # Load entities and axioms
    entities = load_jsonl_files(ENTITIES_DIR)
    axioms = load_jsonl_files(AXIOMS_DIR)

    logger.info(f"Total entities: {len(entities)}")
    logger.info(f"Total axioms: {len(axioms)}")

    # Generate context block
    context_md = generate_bloodstream_context(
        entities,
        axioms,
        task=args.task,
        max_entities=args.max_entities
    )

    # Write output
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as f:
        f.write(context_md)

    logger.info(f"Bloodstream context written to: {args.output}")
    logger.info(f"Context size: {len(context_md)} characters (~{len(context_md) // 4} tokens)")

    # Token budget check
    estimated_tokens = len(context_md) // 4
    if estimated_tokens > MAX_CONTEXT_TOKENS:
        logger.warning(
            f"Context exceeds target token budget: {estimated_tokens} > {MAX_CONTEXT_TOKENS}. "
            f"Consider reducing --max-entities."
        )
    else:
        logger.info(f"Within token budget: {estimated_tokens}/{MAX_CONTEXT_TOKENS} tokens")

    print(f"\n✅ Bloodstream context generated: {args.output}")
    print(f"📊 Entities: {args.max_entities} | Tokens: ~{estimated_tokens}")

if __name__ == "__main__":
    main()