#!/usr/bin/env python3
"""
Enhanced Knowledge Graph v2 - FIXED
Implements Eden's vision for extensible, multi-source knowledge
"""
import json
import os
from datetime import datetime
from collections import defaultdict

class EnhancedKnowledgeGraph:
    """
    Extensible knowledge graph that ingests from multiple sources
    Part of Eden's General Intelligence enhancement
    """
    def __init__(self):
        self.graph_dir = '/Eden/KNOWLEDGE_GRAPH'
        os.makedirs(self.graph_dir, exist_ok=True)
        
        # Multi-domain knowledge storage
        self.domains = {
            'technology': defaultdict(list),
            'business': defaultdict(list),
            'science': defaultdict(list),
            'customers': defaultdict(list),
            'competitors': defaultdict(list)
        }
        
        self.entity_count = 0
        self.relationship_count = 0
        
        self.load_existing_knowledge()
    
    def load_existing_knowledge(self):
        """Load knowledge from existing Eden systems"""
        # Load from market research
        if os.path.exists('/Eden/MARKET_RESEARCH'):
            research_files = [f for f in os.listdir('/Eden/MARKET_RESEARCH') 
                            if f.endswith('.json')]
            for filename in research_files[:100]:  # Start with first 100
                filepath = f'/Eden/MARKET_RESEARCH/{filename}'
                try:
                    with open(filepath, 'r') as f:
                        data = json.load(f)
                        self.ingest_research_data(data)
                except:
                    pass
        
        # Load from customer leads
        if os.path.exists('/Eden/LEADS/leads_database.json'):
            with open('/Eden/LEADS/leads_database.json', 'r') as f:
                leads = json.load(f)
                self.ingest_customer_data(leads)
        
        print(f"✅ Loaded existing knowledge")
        print(f"   Entities: {self.entity_count}")
        print(f"   Relationships: {self.relationship_count}")
    
    def ingest_research_data(self, data):
        """Ingest competitive research into graph"""
        # Extract entities and relationships
        if 'competitors' in data:
            for competitor in data['competitors']:
                self.add_entity('competitors', competitor, 
                              {'source': 'market_research', 
                               'timestamp': data.get('timestamp')})
                self.entity_count += 1
        
        if 'technologies' in data:
            for tech in data['technologies']:
                self.add_entity('technology', tech,
                              {'source': 'market_research'})
                self.entity_count += 1
    
    def ingest_customer_data(self, leads):
        """Ingest customer leads into graph"""
        for lead in leads:
            customer_data = {
                'name': lead.get('author'),
                'platform': lead.get('source'),
                'quality': lead.get('quality_score'),
                'needs': lead.get('needs', []),
                'timestamp': lead.get('timestamp')
            }
            self.add_entity('customers', lead.get('author'), customer_data)
            self.entity_count += 1
            
            # Create relationships
            if 'technologies' in lead:
                for tech in lead['technologies']:
                    self.add_relationship('customers', lead.get('author'),
                                        'uses', 'technology', tech)
                    self.relationship_count += 1
    
    def add_entity(self, domain, entity_id, properties):
        """Add entity to knowledge graph"""
        self.domains[domain][entity_id].append(properties)
    
    def add_relationship(self, from_domain, from_id, relation, 
                        to_domain, to_id):
        """Add relationship between entities"""
        rel_data = {
            'from': f"{from_domain}:{from_id}",
            'relation': relation,
            'to': f"{to_domain}:{to_id}",
            'timestamp': datetime.now().isoformat()
        }
        # Store in both entities for bidirectional lookup
        self.domains[from_domain][from_id].append(
            {'relationship': rel_data}
        )
    
    def query(self, domain, entity_id):
        """Query knowledge about an entity"""
        return self.domains.get(domain, {}).get(entity_id, [])
    
    def get_related(self, domain, entity_id, relation_type=None):
        """Get entities related to this one"""
        entity_data = self.query(domain, entity_id)
        related = []
        
        for entry in entity_data:
            if 'relationship' in entry:
                rel = entry['relationship']
                if relation_type is None or rel['relation'] == relation_type:
                    related.append(rel)
        
        return related
    
    def save(self):
        """Save knowledge graph to disk"""
        filepath = f"{self.graph_dir}/knowledge_graph_v2.json"
        
        # Convert defaultdict to regular dict for JSON
        save_data = {
            'domains': {
                domain: dict(entities) 
                for domain, entities in self.domains.items()
            },
            'stats': {
                'entity_count': self.entity_count,
                'relationship_count': self.relationship_count,
                'last_updated': datetime.now().isoformat()
            }
        }
        
        with open(filepath, 'w') as f:
            json.dump(save_data, f, indent=2)
        
        print(f"\n✅ Knowledge Graph v2 saved to {filepath}")
        return filepath
    
    def get_stats(self):
        """Get knowledge graph statistics - FIXED"""
        return {
            'total_entities': self.entity_count,
            'total_relationships': self.relationship_count,  # FIXED: was missing this
            'relationship_count': self.relationship_count,  # Added for compatibility
            'domains': {
                domain: len(entities) 
                for domain, entities in self.domains.items()
            }
        }

def main():
    print("\n" + "="*70)
    print("🧠 ENHANCED KNOWLEDGE GRAPH v2")
    print("="*70)
    print("Part of Eden's General Intelligence Enhancement")
    print("="*70 + "\n")
    
    # Create and populate knowledge graph
    kg = EnhancedKnowledgeGraph()
    
    # Show stats
    stats = kg.get_stats()
    print("\n📊 Knowledge Graph Statistics:")
    print(f"   Total Entities: {stats['total_entities']}")
    print(f"   Total Relationships: {stats['relationship_count']}")
    print("\n   By Domain:")
    for domain, count in stats['domains'].items():
        print(f"      {domain}: {count} entities")
    
    # Save
    filepath = kg.save()
    
    print("\n✅ Enhanced Knowledge Graph v2 operational!")
    print("   Eden's general intelligence base improved")
    
    return kg

if __name__ == "__main__":
    kg = main()