#!/usr/bin/env python3
"""
TITAN-CLAUDE MEMORY BRIDGE
===========================
Surfaces Titan surprise-based learnings to Claude sessions.

This bridge:
1. Reads surprise events from the evolution loop
2. Extracts actionable learnings
3. Updates CLAUDE.md with persistent insights
4. Queries Elestio stores for relevant context
5. Integrates with TitanKGSync for knowledge graph persistence

Reference: Genesis Protocol + Titans arXiv:2501.00663
"""

import json
import logging
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Any
import hashlib

logger = logging.getLogger(__name__)

GENESIS_ROOT = Path("/mnt/e/genesis-system")

# Import TitanKGSync for knowledge graph integration
try:
    import sys
    sys.path.append(str(GENESIS_ROOT / "core" / "knowledge"))
    from titan_kg_sync import TitanKGSync
    TITAN_KG_AVAILABLE = True
except ImportError:
    logger.warning("TitanKGSync not available - KG sync disabled")
    TITAN_KG_AVAILABLE = False


@dataclass
class TitanLearning:
    """A learning extracted from Titan surprise events."""
    learning_id: str
    category: str
    insight: str
    confidence: float
    source_events: int
    created_at: str
    last_updated: str
    actionable: bool = True

    def to_dict(self) -> Dict:
        return asdict(self)


@dataclass
class ContextPackage:
    """Context package surfaced to Claude."""
    learnings: List[TitanLearning]
    recent_surprises: List[Dict]
    active_patterns: List[str]
    recommended_actions: List[str]
    memory_stats: Dict


class TitanClaudeBridge:
    """
    Bridge between Titan memory system and Claude sessions.

    Makes Titan learnings actionable for Claude by:
    - Aggregating surprise patterns
    - Extracting reusable insights
    - Updating session context
    - Persisting to CLAUDE.md
    """

    def __init__(self):
        self.surprise_events_file = GENESIS_ROOT / "data" / "surprise_events.jsonl"
        self.learnings_file = GENESIS_ROOT / "data" / "titan_learnings.json"
        self.claude_memory_file = GENESIS_ROOT / "data" / "claude_titan_memory.json"

        # Ensure data directory exists
        self.learnings_file.parent.mkdir(parents=True, exist_ok=True)

        # Load existing learnings
        self.learnings = self._load_learnings()

        # Initialize KG sync integration
        self.kg_sync = TitanKGSync() if TITAN_KG_AVAILABLE else None

    def _load_learnings(self) -> Dict[str, TitanLearning]:
        """Load existing learnings from file."""
        if not self.learnings_file.exists():
            return {}

        try:
            data = json.loads(self.learnings_file.read_text())
            return {
                k: TitanLearning(**v) for k, v in data.items()
            }
        except Exception as e:
            logger.error(f"Failed to load learnings: {e}")
            return {}

    def _save_learnings(self):
        """Persist learnings to file."""
        data = {k: v.to_dict() for k, v in self.learnings.items()}
        self.learnings_file.write_text(json.dumps(data, indent=2))

    def ingest_surprise_events(self) -> int:
        """
        Ingest surprise events and extract learnings.

        Also triggers TitanKGSync to persist to knowledge graph.

        Returns:
            Number of new learnings extracted
        """
        if not self.surprise_events_file.exists():
            return 0

        events = []
        for line in self.surprise_events_file.read_text().strip().split('\n'):
            if line:
                try:
                    events.append(json.loads(line))
                except:
                    continue

        if not events:
            return 0

        new_learnings = 0

        # Analyze patterns in events
        patterns = self._analyze_patterns(events)

        for pattern in patterns:
            learning = self._pattern_to_learning(pattern, events)
            if learning and learning.learning_id not in self.learnings:
                self.learnings[learning.learning_id] = learning
                new_learnings += 1
                logger.info(f"New Titan learning: {learning.insight[:50]}...")

        if new_learnings > 0:
            self._save_learnings()

            # Trigger KG sync if available
            if self.kg_sync:
                try:
                    kg_stats = self.kg_sync.sync()
                    logger.info(f"KG sync: {kg_stats['new_learnings']} learnings, {kg_stats['new_axioms']} axioms")
                except Exception as e:
                    logger.error(f"KG sync failed: {e}")

        return new_learnings

    def _analyze_patterns(self, events: List[Dict]) -> List[Dict]:
        """Analyze events for recurring patterns."""
        patterns = []

        # Group by trigger reason
        reason_counts = {}
        for e in events:
            reason = e.get('trigger_reason', 'unknown')
            if reason not in reason_counts:
                reason_counts[reason] = {'count': 0, 'avg_score': 0, 'events': []}
            reason_counts[reason]['count'] += 1
            reason_counts[reason]['avg_score'] += e.get('surprise_score', 0)
            reason_counts[reason]['events'].append(e)

        # Calculate averages and create patterns
        for reason, data in reason_counts.items():
            if data['count'] >= 1:  # Even single events are valuable
                patterns.append({
                    'type': 'trigger_pattern',
                    'reason': reason,
                    'frequency': data['count'],
                    'avg_surprise': data['avg_score'] / data['count'],
                    'sample_events': data['events'][:3]
                })

        return patterns

    def _pattern_to_learning(self, pattern: Dict, all_events: List[Dict]) -> Optional[TitanLearning]:
        """Convert a pattern to an actionable learning."""
        reason = pattern.get('reason', 'unknown')
        frequency = pattern.get('frequency', 1)
        avg_surprise = pattern.get('avg_surprise', 0.5)

        # Generate insight based on pattern type
        if 'test_failures' in reason:
            insight = f"Tests failing unexpectedly - review test assumptions and edge cases"
            category = "testing"
        elif 'slow_execution' in reason:
            insight = f"Execution slower than expected - consider timeout adjustments or optimization"
            category = "performance"
        elif 'sparse_output' in reason:
            insight = f"Output sparser than expected - prompts may need more detail"
            category = "prompting"
        elif 'verbose_output' in reason:
            insight = f"Output more verbose than expected - consider token limits"
            category = "efficiency"
        elif 'unexpected_errors' in reason:
            insight = f"Unexpected errors occurring - enhance error handling"
            category = "reliability"
        else:
            insight = f"Pattern detected: {reason} - investigate root cause"
            category = "general"

        # Generate unique ID
        learning_id = hashlib.md5(f"{reason}_{category}".encode()).hexdigest()[:12]

        now = datetime.now().isoformat()

        return TitanLearning(
            learning_id=learning_id,
            category=category,
            insight=insight,
            confidence=min(0.5 + (frequency * 0.1) + (avg_surprise * 0.3), 1.0),
            source_events=frequency,
            created_at=now,
            last_updated=now,
            actionable=True
        )

    def get_context_package(self) -> ContextPackage:
        """
        Get a context package for Claude sessions.

        This is the main interface - call this at session start.
        """
        # First, ingest any new events
        self.ingest_surprise_events()

        # Load recent surprises
        recent_surprises = []
        if self.surprise_events_file.exists():
            lines = self.surprise_events_file.read_text().strip().split('\n')
            for line in lines[-10:]:  # Last 10 events
                if line:
                    try:
                        recent_surprises.append(json.loads(line))
                    except:
                        continue

        # Get active patterns
        active_patterns = [
            f"{l.category}: {l.insight}"
            for l in self.learnings.values()
            if l.actionable
        ]

        # Generate recommended actions
        recommended_actions = self._generate_recommendations()

        # Memory stats
        stats = {
            'total_learnings': len(self.learnings),
            'total_surprise_events': len(recent_surprises),
            'categories': list(set(l.category for l in self.learnings.values())),
            'avg_confidence': sum(l.confidence for l in self.learnings.values()) / max(len(self.learnings), 1),
            'last_update': datetime.now().isoformat()
        }

        return ContextPackage(
            learnings=list(self.learnings.values()),
            recent_surprises=recent_surprises,
            active_patterns=active_patterns,
            recommended_actions=recommended_actions,
            memory_stats=stats
        )

    def _generate_recommendations(self) -> List[str]:
        """Generate actionable recommendations from learnings."""
        recommendations = []

        categories = {}
        for l in self.learnings.values():
            if l.category not in categories:
                categories[l.category] = []
            categories[l.category].append(l)

        if 'testing' in categories:
            recommendations.append("Review test suites for edge cases - surprise events detected")

        if 'performance' in categories:
            recommendations.append("Consider increasing timeouts for complex operations")

        if 'prompting' in categories:
            recommendations.append("Enhance prompts with more specific output requirements")

        if 'reliability' in categories:
            recommendations.append("Add error handling for detected failure patterns")

        if not recommendations:
            recommendations.append("No critical patterns detected - system healthy")

        return recommendations

    def update_claude_md(self) -> bool:
        """
        Update CLAUDE.md with Titan memory section.

        This persists learnings where Claude will see them every session.
        """
        claude_md_path = GENESIS_ROOT / "CLAUDE.md"

        if not claude_md_path.exists():
            logger.error("CLAUDE.md not found")
            return False

        content = claude_md_path.read_text()

        # Generate Titan memory section
        package = self.get_context_package()

        titan_section = self._generate_titan_section(package)

        # Check if section exists
        marker_start = "## TITAN MEMORY INSIGHTS"
        marker_end = "## END TITAN MEMORY"

        if marker_start in content:
            # Replace existing section
            start_idx = content.find(marker_start)
            end_idx = content.find(marker_end)
            if end_idx > start_idx:
                content = content[:start_idx] + titan_section + content[end_idx + len(marker_end):]
        else:
            # Add new section before first ---
            insert_point = content.find("---")
            if insert_point > 0:
                content = content[:insert_point] + titan_section + "\n\n" + content[insert_point:]

        claude_md_path.write_text(content)
        logger.info("Updated CLAUDE.md with Titan memory insights")
        return True

    def _generate_titan_section(self, package: ContextPackage) -> str:
        """Generate the Titan memory section for CLAUDE.md."""
        lines = [
            "## TITAN MEMORY INSIGHTS",
            f"> Auto-updated: {datetime.now().strftime('%Y-%m-%d %H:%M')}",
            "> Based on surprise-based learning from evolution cycles",
            "",
            "### Active Learnings",
        ]

        if package.learnings:
            for learning in package.learnings[:5]:  # Top 5
                lines.append(f"- **{learning.category}**: {learning.insight} (confidence: {learning.confidence:.0%})")
        else:
            lines.append("- No learnings captured yet - system is learning")

        lines.extend([
            "",
            "### Recommended Actions",
        ])

        for action in package.recommended_actions:
            lines.append(f"- {action}")

        lines.extend([
            "",
            f"### Memory Stats",
            f"- Total Learnings: {package.memory_stats['total_learnings']}",
            f"- Surprise Events: {package.memory_stats['total_surprise_events']}",
            f"- Categories: {', '.join(package.memory_stats['categories']) or 'None'}",
            "",
            "## END TITAN MEMORY",
            ""
        ])

        return '\n'.join(lines)

    def save_claude_memory(self):
        """Save current memory state for Claude session persistence."""
        package = self.get_context_package()

        memory_data = {
            'learnings': [l.to_dict() for l in package.learnings],
            'recent_surprises': package.recent_surprises,
            'active_patterns': package.active_patterns,
            'recommended_actions': package.recommended_actions,
            'memory_stats': package.memory_stats,
            'saved_at': datetime.now().isoformat()
        }

        self.claude_memory_file.write_text(json.dumps(memory_data, indent=2))
        logger.info(f"Saved Claude memory to {self.claude_memory_file}")


def main():
    """CLI interface for Titan-Claude bridge."""
    import argparse

    parser = argparse.ArgumentParser(description='Titan-Claude Memory Bridge')
    parser.add_argument('--ingest', action='store_true', help='Ingest surprise events')
    parser.add_argument('--update-claude', action='store_true', help='Update CLAUDE.md')
    parser.add_argument('--context', action='store_true', help='Print context package')
    parser.add_argument('--save', action='store_true', help='Save Claude memory')

    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    bridge = TitanClaudeBridge()

    if args.ingest:
        new = bridge.ingest_surprise_events()
        print(f"Ingested {new} new learnings")

    if args.update_claude:
        success = bridge.update_claude_md()
        print(f"CLAUDE.md update: {'success' if success else 'failed'}")

    if args.context:
        package = bridge.get_context_package()
        print(json.dumps({
            'learnings': [l.to_dict() for l in package.learnings],
            'patterns': package.active_patterns,
            'recommendations': package.recommended_actions,
            'stats': package.memory_stats
        }, indent=2))

    if args.save:
        bridge.save_claude_memory()
        print("Claude memory saved")

    if not any([args.ingest, args.update_claude, args.context, args.save]):
        # Default: full cycle
        bridge.ingest_surprise_events()
        bridge.update_claude_md()
        bridge.save_claude_memory()
        print("Full Titan-Claude bridge cycle complete")


if __name__ == '__main__':
    main()