#!/usr/bin/env python3
"""
Migration script to convert existing tech-stack-selector data to TSS namespace
This ensures data isolation between template-manager (TM) and tech-stack-selector (TSS)
"""

import os
import sys
from typing import Dict, Any, Optional, List
from neo4j import GraphDatabase
from loguru import logger

class TSSNamespaceMigration:
    """
    Migrates existing tech-stack-selector data to use TSS namespace
    """
    
    def __init__(self):
        self.neo4j_uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
        self.neo4j_user = os.getenv("NEO4J_USER", "neo4j")
        self.neo4j_password = os.getenv("NEO4J_PASSWORD", "password")
        self.namespace = "TSS"
        
        self.driver = GraphDatabase.driver(
            self.neo4j_uri,
            auth=(self.neo4j_user, self.neo4j_password),
            connection_timeout=10
        )
        
        self.migration_stats = {
            "nodes_migrated": 0,
            "relationships_migrated": 0,
            "errors": 0,
            "skipped": 0
        }
    
    def close(self):
        if self.driver:
            self.driver.close()
    
    def run_query(self, query: str, parameters: Optional[Dict[str, Any]] = None):
        """Execute a Neo4j query"""
        try:
            with self.driver.session() as session:
                result = session.run(query, parameters or {})
                return [record.data() for record in result]
        except Exception as e:
            logger.error(f"❌ Query failed: {e}")
            self.migration_stats["errors"] += 1
            raise e
    
    def check_existing_data(self):
        """Check what data exists before migration"""
        logger.info("🔍 Checking existing data...")
        
        # Check for existing TSS namespaced data
        tss_nodes_query = f"""
            MATCH (n)
            WHERE '{self.namespace}' IN labels(n)
            RETURN labels(n) as labels, count(n) as count
        """
        tss_results = self.run_query(tss_nodes_query)
        
        if tss_results:
            logger.info("✅ Found existing TSS namespaced data:")
            for record in tss_results:
                logger.info(f"   - {record['labels']}: {record['count']} nodes")
        else:
            logger.info("ℹ️ No existing TSS namespaced data found")
        
        # Check for non-namespaced tech-stack-selector data
        non_namespaced_query = """
            MATCH (n)
            WHERE (n:TechStack OR n:Technology OR n:PriceTier OR n:Tool OR n:Domain)
            AND NOT 'TM' IN labels(n) AND NOT 'TSS' IN labels(n)
            RETURN labels(n) as labels, count(n) as count
        """
        non_namespaced_results = self.run_query(non_namespaced_query)
        
        if non_namespaced_results:
            logger.info("🎯 Found non-namespaced data to migrate:")
            for record in non_namespaced_results:
                logger.info(f"   - {record['labels']}: {record['count']} nodes")
            return True
        else:
            logger.info("ℹ️ No non-namespaced data found to migrate")
            return False
    
    def migrate_nodes(self):
        """Migrate nodes to TSS namespace"""
        logger.info("🔄 Migrating nodes to TSS namespace...")
        
        # Define node types to migrate
        node_types = [
            "TechStack",
            "Technology", 
            "PriceTier",
            "Tool",
            "Domain"
        ]
        
        for node_type in node_types:
            try:
                # Add TSS label to existing nodes that don't have TM or TSS namespace
                query = f"""
                    MATCH (n:{node_type})
                    WHERE NOT 'TM' IN labels(n) AND NOT 'TSS' IN labels(n)
                    SET n:{node_type}:TSS
                    RETURN count(n) as migrated_count
                """
                
                result = self.run_query(query)
                migrated_count = result[0]['migrated_count'] if result else 0
                
                if migrated_count > 0:
                    logger.info(f"✅ Migrated {migrated_count} {node_type} nodes to TSS namespace")
                    self.migration_stats["nodes_migrated"] += migrated_count
                else:
                    logger.info(f"ℹ️ No {node_type} nodes to migrate")
                    
            except Exception as e:
                logger.error(f"❌ Failed to migrate {node_type} nodes: {e}")
                self.migration_stats["errors"] += 1
    
    def migrate_relationships(self):
        """Migrate relationships to TSS namespace"""
        logger.info("🔄 Migrating relationships to TSS namespace...")
        
        # Define relationship types to migrate
        relationship_mappings = {
            "BELONGS_TO_TIER": "BELONGS_TO_TIER_TSS",
            "USES_FRONTEND": "USES_FRONTEND_TSS",
            "USES_BACKEND": "USES_BACKEND_TSS",
            "USES_DATABASE": "USES_DATABASE_TSS",
            "USES_CLOUD": "USES_CLOUD_TSS",
            "USES_TESTING": "USES_TESTING_TSS",
            "USES_MOBILE": "USES_MOBILE_TSS",
            "USES_DEVOPS": "USES_DEVOPS_TSS",
            "USES_AI_ML": "USES_AI_ML_TSS",
            "RECOMMENDS": "RECOMMENDS_TSS",
            "COMPATIBLE_WITH": "COMPATIBLE_WITH_TSS",
            "HAS_CLAUDE_RECOMMENDATION": "HAS_CLAUDE_RECOMMENDATION_TSS"
        }
        
        for old_rel, new_rel in relationship_mappings.items():
            try:
                # Find relationships between TSS nodes that need to be updated
                query = f"""
                    MATCH (a)-[r:{old_rel}]->(b)
                    WHERE 'TSS' IN labels(a) AND 'TSS' IN labels(b)
                    AND NOT type(r) CONTAINS 'TSS'
                    AND NOT type(r) CONTAINS 'TM'
                    WITH a, b, r, properties(r) as props
                    DELETE r
                    CREATE (a)-[new_r:{new_rel}]->(b)
                    SET new_r = props
                    RETURN count(new_r) as migrated_count
                """
                
                result = self.run_query(query)
                migrated_count = result[0]['migrated_count'] if result else 0
                
                if migrated_count > 0:
                    logger.info(f"✅ Migrated {migrated_count} {old_rel} relationships to {new_rel}")
                    self.migration_stats["relationships_migrated"] += migrated_count
                else:
                    logger.info(f"ℹ️ No {old_rel} relationships to migrate")
                    
            except Exception as e:
                logger.error(f"❌ Failed to migrate {old_rel} relationships: {e}")
                self.migration_stats["errors"] += 1
    
    def verify_migration(self):
        """Verify the migration was successful"""
        logger.info("🔍 Verifying migration...")
        
        # Check TSS namespaced data
        tss_query = f"""
            MATCH (n)
            WHERE '{self.namespace}' IN labels(n)
            RETURN labels(n) as labels, count(n) as count
        """
        tss_results = self.run_query(tss_query)
        
        if tss_results:
            logger.info("✅ TSS namespaced nodes after migration:")
            for record in tss_results:
                logger.info(f"   - {record['labels']}: {record['count']} nodes")
        
        # Check TSS namespaced relationships
        tss_rel_query = f"""
            MATCH ()-[r]->()
            WHERE type(r) CONTAINS '{self.namespace}'
            RETURN type(r) as rel_type, count(r) as count
        """
        tss_rel_results = self.run_query(tss_rel_query)
        
        if tss_rel_results:
            logger.info("✅ TSS namespaced relationships after migration:")
            for record in tss_rel_results:
                logger.info(f"   - {record['rel_type']}: {record['count']} relationships")
        
        # Check for remaining non-namespaced data
        remaining_query = """
            MATCH (n)
            WHERE (n:TechStack OR n:Technology OR n:PriceTier OR n:Tool OR n:Domain)
            AND NOT 'TM' IN labels(n) AND NOT 'TSS' IN labels(n)
            RETURN labels(n) as labels, count(n) as count
        """
        remaining_results = self.run_query(remaining_query)
        
        if remaining_results:
            logger.warning("⚠️ Remaining non-namespaced data:")
            for record in remaining_results:
                logger.warning(f"   - {record['labels']}: {record['count']} nodes")
        else:
            logger.info("✅ All data has been properly namespaced")
    
    def run_migration(self):
        """Run the complete migration process"""
        logger.info("🚀 Starting TSS namespace migration...")
        logger.info("="*60)
        
        try:
            # Check connection
            with self.driver.session() as session:
                session.run("RETURN 1")
            logger.info("✅ Neo4j connection established")
            
            # Check existing data
            has_data_to_migrate = self.check_existing_data()
            
            if not has_data_to_migrate:
                logger.info("ℹ️ No non-namespaced data to migrate.")
                logger.info("✅ Either no data exists or data is already properly namespaced.")
                logger.info("✅ TSS namespace migration completed successfully.")
                return True
            
            # Migrate nodes
            self.migrate_nodes()
            
            # Migrate relationships
            self.migrate_relationships()
            
            # Verify migration
            self.verify_migration()
            
            # Print summary
            logger.info("="*60)
            logger.info("📊 Migration Summary:")
            logger.info(f"   - Nodes migrated: {self.migration_stats['nodes_migrated']}")
            logger.info(f"   - Relationships migrated: {self.migration_stats['relationships_migrated']}")
            logger.info(f"   - Errors: {self.migration_stats['errors']}")
            logger.info(f"   - Skipped: {self.migration_stats['skipped']}")
            
            if self.migration_stats["errors"] == 0:
                logger.info("✅ Migration completed successfully!")
                return True
            else:
                logger.error("❌ Migration completed with errors!")
                return False
                
        except Exception as e:
            logger.error(f"❌ Migration failed: {e}")
            return False
        finally:
            self.close()

def main():
    """Main function"""
    logger.remove()
    logger.add(sys.stdout, level="INFO", format="{time} | {level} | {message}")
    
    migration = TSSNamespaceMigration()
    success = migration.run_migration()
    
    if success:
        logger.info("🎉 TSS namespace migration completed successfully!")
        sys.exit(0)
    else:
        logger.error("💥 TSS namespace migration failed!")
        sys.exit(1)

if __name__ == "__main__":
    main()
