#!/usr/bin/env python3
"""
Genesis Graphiti Integration
=============================
Temporal Knowledge Graph for Genesis Ambient Memory

Integrates with Graphiti to provide:
- Real-time incremental knowledge graph updates
- Bi-temporal data model (event time + ingestion time)
- Entity and relationship extraction from events
- Time-travel queries (what did we know at time X?)

Usage:
    from graphiti_integration import GenesisGraphiti

    graphiti = GenesisGraphiti()
    await graphiti.ingest_event(event)
    results = await graphiti.query("What decisions were made about memory?")
"""

import asyncio
import json
import os
from datetime import datetime
from typing import Dict, List, Optional, Any
from dataclasses import dataclass

# Try to import graphiti
try:
    from graphiti_core import Graphiti
    from graphiti_core.nodes import EpisodeType
    GRAPHITI_AVAILABLE = True
except ImportError:
    GRAPHITI_AVAILABLE = False
    print("Warning: graphiti-core not available, using fallback mode")

# Import our event types
try:
    from genesis_ambient_memory import AmbientEvent, EventType
except ImportError:
    from .genesis_ambient_memory import AmbientEvent, EventType


@dataclass
class TemporalFact:
    """A fact with temporal awareness."""
    content: str
    entity_from: str
    entity_to: str
    relation: str
    valid_at: datetime
    invalid_at: Optional[datetime] = None
    confidence: float = 1.0
    source: str = "ambient"


class GenesisGraphiti:
    """
    Temporal Knowledge Graph integration for Genesis.

    Wraps Graphiti to provide:
    - Automatic entity/relation extraction from ambient events
    - Bi-temporal tracking (when it happened vs when we learned it)
    - Conflict resolution for contradicting facts
    """

    def __init__(
        self,
        neo4j_uri: str = None,
        neo4j_user: str = None,
        neo4j_password: str = None
    ):
        self.neo4j_uri = neo4j_uri or os.getenv(
            "GENESIS_NEO4J_URI",
            "bolt://localhost:7687"
        )
        self.neo4j_user = neo4j_user or os.getenv("GENESIS_NEO4J_USER", "neo4j")
        self.neo4j_password = neo4j_password or os.getenv("GENESIS_NEO4J_PASSWORD", "genesis")

        self._graphiti = None
        self._initialized = False

        # Fallback storage when Graphiti/Neo4j unavailable
        self._fallback_facts: List[TemporalFact] = []

    async def initialize(self) -> bool:
        """Initialize Graphiti connection."""
        if not GRAPHITI_AVAILABLE:
            print("Graphiti not available, using fallback storage")
            return False

        try:
            self._graphiti = Graphiti(
                uri=self.neo4j_uri,
                user=self.neo4j_user,
                password=self.neo4j_password
            )
            await self._graphiti.build_indices_and_constraints()
            self._initialized = True
            return True
        except Exception as e:
            print(f"Could not connect to Neo4j: {e}")
            print("Using fallback storage mode")
            return False

    async def ingest_event(self, event: AmbientEvent) -> Dict[str, Any]:
        """
        Ingest an ambient event into the temporal knowledge graph.

        Extracts entities and relationships, adds with temporal metadata.
        """
        # Extract content for ingestion
        content = self._event_to_narrative(event)
        source = f"genesis:ambient:{event.event_type.value}"

        if self._initialized and self._graphiti:
            try:
                # Use Graphiti's episode ingestion
                await self._graphiti.add_episode(
                    name=f"event-{event.event_id}",
                    episode_body=content,
                    source=EpisodeType.text,
                    source_description=source,
                    reference_time=datetime.fromisoformat(event.timestamp)
                )
                return {
                    "status": "ingested",
                    "mode": "graphiti",
                    "event_id": event.event_id
                }
            except Exception as e:
                print(f"Graphiti ingestion error: {e}")
                # Fall through to fallback

        # Fallback: Extract and store locally
        facts = self._extract_facts(event, content)
        self._fallback_facts.extend(facts)

        return {
            "status": "ingested",
            "mode": "fallback",
            "event_id": event.event_id,
            "facts_extracted": len(facts)
        }

    async def query(
        self,
        query: str,
        as_of: datetime = None
    ) -> List[Dict[str, Any]]:
        """
        Query the knowledge graph.

        Args:
            query: Natural language query
            as_of: Point-in-time query (what did we know at this time?)
        """
        if self._initialized and self._graphiti:
            try:
                results = await self._graphiti.search(query)
                return [
                    {
                        "content": r.content if hasattr(r, 'content') else str(r),
                        "score": r.score if hasattr(r, 'score') else 1.0
                    }
                    for r in results
                ]
            except Exception as e:
                print(f"Graphiti query error: {e}")

        # Fallback: Simple keyword search
        query_lower = query.lower()
        matching = []
        for fact in self._fallback_facts:
            if as_of and fact.valid_at > as_of:
                continue
            if query_lower in fact.content.lower():
                matching.append({
                    "content": fact.content,
                    "entity_from": fact.entity_from,
                    "entity_to": fact.entity_to,
                    "relation": fact.relation,
                    "valid_at": fact.valid_at.isoformat()
                })

        return matching

    async def get_entity_history(
        self,
        entity_name: str
    ) -> List[Dict[str, Any]]:
        """Get the temporal history of an entity."""
        if self._initialized and self._graphiti:
            try:
                # Query for entity edges over time
                results = await self._graphiti.search(f"entity:{entity_name}")
                return [
                    {
                        "content": r.content if hasattr(r, 'content') else str(r),
                        "timestamp": r.created_at.isoformat() if hasattr(r, 'created_at') else None
                    }
                    for r in results
                ]
            except Exception as e:
                print(f"Entity history error: {e}")

        # Fallback
        entity_lower = entity_name.lower()
        return [
            {
                "content": f.content,
                "relation": f.relation,
                "valid_at": f.valid_at.isoformat()
            }
            for f in self._fallback_facts
            if entity_lower in f.entity_from.lower() or entity_lower in f.entity_to.lower()
        ]

    def _event_to_narrative(self, event: AmbientEvent) -> str:
        """Convert event to natural language narrative for Graphiti."""
        timestamp = event.timestamp[:19]  # Truncate to seconds

        if event.event_type == EventType.TOOL_CALL:
            tool = event.content.get("tool_name", "unknown")
            success = "successfully" if event.content.get("success", True) else "with errors"
            return f"At {timestamp}, Claude Code called the {tool} tool {success}."

        elif event.event_type == EventType.DECISION:
            decision = event.content.get("decision", "unknown decision")
            reasoning = event.content.get("reasoning", "")
            narrative = f"At {timestamp}, a decision was made: {decision}."
            if reasoning:
                narrative += f" The reasoning was: {reasoning}"
            return narrative

        elif event.event_type == EventType.PLAN:
            summary = event.content.get("summary", "unknown plan")
            steps = event.content.get("steps", [])
            narrative = f"At {timestamp}, a plan was created: {summary}."
            if steps:
                narrative += f" Steps: {', '.join(steps[:5])}"
            return narrative

        elif event.event_type == EventType.ERROR:
            error = event.content.get("error", "unknown error")
            return f"At {timestamp}, an error occurred: {error}"

        elif event.event_type == EventType.LEARNING:
            learning = event.content.get("learning", "unknown learning")
            return f"At {timestamp}, a learning was captured: {learning}"

        else:
            return f"At {timestamp}, event of type {event.event_type.value}: {json.dumps(event.content)[:200]}"

    def _extract_facts(self, event: AmbientEvent, narrative: str) -> List[TemporalFact]:
        """Extract facts from event (fallback when Graphiti unavailable)."""
        facts = []
        timestamp = datetime.fromisoformat(event.timestamp)

        if event.event_type == EventType.TOOL_CALL:
            tool = event.content.get("tool_name", "unknown")
            facts.append(TemporalFact(
                content=narrative,
                entity_from="Claude Code",
                entity_to=tool,
                relation="USED_TOOL",
                valid_at=timestamp,
                source=f"session:{event.session_id}"
            ))

        elif event.event_type == EventType.DECISION:
            decision = event.content.get("decision", "")
            facts.append(TemporalFact(
                content=narrative,
                entity_from="Agent",
                entity_to=decision[:50],
                relation="DECIDED",
                valid_at=timestamp,
                source=f"session:{event.session_id}"
            ))

        elif event.event_type == EventType.LEARNING:
            learning = event.content.get("learning", "")
            facts.append(TemporalFact(
                content=narrative,
                entity_from="Genesis",
                entity_to=learning[:50],
                relation="LEARNED",
                valid_at=timestamp,
                source=f"session:{event.session_id}"
            ))

        return facts

    async def export_fallback_facts(self) -> List[Dict]:
        """Export fallback facts for later Neo4j import."""
        return [
            {
                "content": f.content,
                "entity_from": f.entity_from,
                "entity_to": f.entity_to,
                "relation": f.relation,
                "valid_at": f.valid_at.isoformat(),
                "invalid_at": f.invalid_at.isoformat() if f.invalid_at else None,
                "confidence": f.confidence,
                "source": f.source
            }
            for f in self._fallback_facts
        ]

    async def close(self):
        """Close connections."""
        if self._graphiti:
            try:
                await self._graphiti.close()
            except:
                pass


class AmbientToGraphitiBridge:
    """
    Bridge that automatically syncs ambient events to Graphiti.

    Runs as a background process, consuming from the event stream
    and ingesting into the temporal knowledge graph.
    """

    def __init__(self, graphiti: GenesisGraphiti = None):
        self.graphiti = graphiti or GenesisGraphiti()
        self._running = False

    async def start(self):
        """Start the bridge (consumes events and ingests to Graphiti)."""
        await self.graphiti.initialize()
        self._running = True

        # In a real implementation, this would consume from Redis Stream
        print("AmbientToGraphitiBridge started (would consume from Redis Stream)")

    async def process_event(self, event: AmbientEvent) -> Dict:
        """Process a single event."""
        return await self.graphiti.ingest_event(event)

    async def stop(self):
        """Stop the bridge."""
        self._running = False
        await self.graphiti.close()


# CLI Interface
async def main():
    import sys

    if len(sys.argv) < 2:
        print("Genesis Graphiti Integration")
        print("\nUsage:")
        print("  python graphiti_integration.py query 'What decisions were made?'")
        print("  python graphiti_integration.py ingest <event_json>")
        print("  python graphiti_integration.py status")
        return

    command = sys.argv[1]
    graphiti = GenesisGraphiti()

    if command == "status":
        initialized = await graphiti.initialize()
        print(json.dumps({
            "graphiti_available": GRAPHITI_AVAILABLE,
            "neo4j_connected": initialized,
            "neo4j_uri": graphiti.neo4j_uri
        }, indent=2))

    elif command == "query" and len(sys.argv) > 2:
        await graphiti.initialize()
        query = " ".join(sys.argv[2:])
        results = await graphiti.query(query)
        print(json.dumps(results, indent=2))

    elif command == "ingest" and len(sys.argv) > 2:
        await graphiti.initialize()
        event_data = json.loads(sys.argv[2])
        event = AmbientEvent.from_dict(event_data)
        result = await graphiti.ingest_event(event)
        print(json.dumps(result, indent=2))

    else:
        print(f"Unknown command: {command}")

    await graphiti.close()


if __name__ == "__main__":
    asyncio.run(main())