import json
import os
import shlex

# This script generates a shell script to migrate the knowledge graph.

def get_rlm_type(legacy_type):
    # Mapping from RLM_SCHEMA_MAPPING.md
    mapping = {
        "technology_enabler": "RESOURCE",
        "protocol": "PATTERN",
        "strategy_node": "LESSON",
        "pricing_model": "FEATURE",
        "service_architecture": "PATTERN",
        "market_segment": "RESOURCE",
        "premortem_risk": "BLOCKER",
        "infrastructure": "INTEGRATION",
        "value_proposition": "FEATURE",
        "retention_protocol": "PATTERN",
        "decision_point": "DECISION",
        "brand_strategy": "PATTERN",
        "sales_framework": "PATTERN",
        "success_criteria": "TASK",
        "milestone": "TASK",
        "localization_rule": "PATTERN",
        "quality_control": "PATTERN",
        "market_intelligence": "RESOURCE",
        "pricing_intelligence": "RESOURCE",
        "strategic_axiom": "LESSON",
        "decision_record": "DECISION",
        "research_request": "TASK",
        "pricing_recommendation": "FEATURE",
        "performance_data": "RESOURCE",
        "market_data": "RESOURCE",
        "strategic_plan": "TASK",
        "competitive_intelligence": "RESOURCE",
        "mandatory_protocol": "PATTERN",
        "critical_learning": "LESSON",
        "strategic_insight": "LESSON",
        "youtube_research": "RESOURCE",
        "strategic_intelligence": "RESOURCE",
        "knowledge_mastery": "RESOURCE",
    }
    return mapping.get(legacy_type, "RESOURCE") # Default to RESOURCE if not found


def migrate():
    entities_file = '/mnt/e/genesis-system/KNOWLEDGE_GRAPH/entities.jsonl'
    if not os.path.exists(entities_file):
        print(f"File not found: {entities_file}")
        return

    commands = []
    existing_entities = set()

    # First, create all entities
    with open(entities_file, 'r') as f:
        for line in f:
            try:
                entity = json.loads(line)
                entity_id = entity.get("id")
                legacy_type = entity.get("type")

                if not entity_id or not legacy_type:
                    continue
                
                if entity_id in existing_entities:
                    continue

                rlm_type = get_rlm_type(legacy_type)
                name = entity_id
                observations = [json.dumps(entity)]
                
                tool_code = f"print(genesis_voice_bridge.store_memory(name={shlex.quote(name)}, entity_type={shlex.quote(rlm_type)}, observations={shlex.quote(json.dumps(observations))}))"
                commands.append(f"gemini -p {shlex.quote(tool_code)}")
                existing_entities.add(entity_id)

            except json.JSONDecodeError:
                continue

    # Then, create all relationships
    with open(entities_file, 'r') as f:
        for line in f:
            try:
                entity = json.loads(line)
                entity_id = entity.get("id")

                if not entity_id:
                    continue

                if 'axiom_id' in entity:
                    to_entity = entity['axiom_id']
                    if to_entity not in existing_entities:
                        # Create the missing entity
                        tool_code = f"print(genesis_voice_bridge.store_memory(name={shlex.quote(to_entity)}, entity_type='RESOURCE', observations={shlex.quote(json.dumps([{'description': 'A generic entity created during migration.'}]))}))"
                        commands.insert(0, f"gemini -p {shlex.quote(tool_code)}")
                        existing_entities.add(to_entity)
                    
                    tool_code = f"print(genesis_voice_bridge.create_relation(from_entity={shlex.quote(entity_id)}, to_entity={shlex.quote(to_entity)}, relation_type='RELATED_TO'))"
                    commands.append(f"gemini -p {shlex.quote(tool_code)}")

            except json.JSONDecodeError:
                continue

    with open('migration_script.sh', 'w') as f:
        f.write("#!/bin/bash\\n")
        f.write("\\n".join(commands))

if __name__ == "__main__":
    migrate()