#!/usr/bin/env python3
"""
Tech Stack Selector Database Setup Script
Handles PostgreSQL migrations and Neo4j data migration
"""

import os
import sys
import subprocess
import psycopg2
from neo4j import GraphDatabase
from loguru import logger

def setup_environment():
    """Set up environment variables"""
    os.environ.setdefault("POSTGRES_HOST", "postgres")
    os.environ.setdefault("POSTGRES_PORT", "5432")
    os.environ.setdefault("POSTGRES_USER", "pipeline_admin")
    os.environ.setdefault("POSTGRES_PASSWORD", "secure_pipeline_2024")
    os.environ.setdefault("POSTGRES_DB", "dev_pipeline")
    os.environ.setdefault("NEO4J_URI", "bolt://neo4j:7687")
    os.environ.setdefault("NEO4J_USER", "neo4j")
    os.environ.setdefault("NEO4J_PASSWORD", "password")
    os.environ.setdefault("CLAUDE_API_KEY", "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")

def check_postgres_connection():
    """Check if PostgreSQL is accessible"""
    try:
        conn = psycopg2.connect(
            host=os.getenv('POSTGRES_HOST'),
            port=int(os.getenv('POSTGRES_PORT')),
            user=os.getenv('POSTGRES_USER'),
            password=os.getenv('POSTGRES_PASSWORD'),
            database='postgres'
        )
        conn.close()
        logger.info("✅ PostgreSQL connection successful")
        return True
    except Exception as e:
        logger.error(f"❌ PostgreSQL connection failed: {e}")
        return False

def check_neo4j_connection():
    """Check if Neo4j is accessible"""
    try:
        driver = GraphDatabase.driver(
            os.getenv('NEO4J_URI'),
            auth=(os.getenv('NEO4J_USER'), os.getenv('NEO4J_PASSWORD'))
        )
        driver.verify_connectivity()
        driver.close()
        logger.info("✅ Neo4j connection successful")
        return True
    except Exception as e:
        logger.error(f"❌ Neo4j connection failed: {e}")
        return False

def run_postgres_migrations():
    """Run PostgreSQL migrations"""
    logger.info("🔄 Running PostgreSQL migrations...")
    
    migration_files = [
        "db/001_schema.sql",
        "db/002_tools_migration.sql", 
        "db/003_tools_pricing_migration.sql",
        "db/004_comprehensive_stacks_migration.sql",
        "db/005_comprehensive_ecommerce_stacks.sql",
        "db/006_comprehensive_all_domains_stacks.sql"
    ]
    
    # Set PGPASSWORD to avoid password prompts
    os.environ["PGPASSWORD"] = os.getenv('POSTGRES_PASSWORD')
    
    for migration_file in migration_files:
        if not os.path.exists(migration_file):
            logger.warning(f"⚠️ Migration file not found: {migration_file}")
            continue
            
        logger.info(f"📄 Running migration: {migration_file}")
        
        try:
            result = subprocess.run([
                'psql', 
                '-h', os.getenv('POSTGRES_HOST'),
                '-p', os.getenv('POSTGRES_PORT'),
                '-U', os.getenv('POSTGRES_USER'),
                '-d', os.getenv('POSTGRES_DB'),
                '-f', migration_file,
                '-q'
            ], capture_output=True, text=True)
            
            if result.returncode == 0:
                logger.info(f"✅ Migration completed: {migration_file}")
            else:
                logger.error(f"❌ Migration failed: {migration_file}")
                logger.error(f"Error: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"❌ Migration error: {e}")
            return False
    
    # Unset password
    if 'PGPASSWORD' in os.environ:
        del os.environ['PGPASSWORD']
    
    logger.info("✅ All PostgreSQL migrations completed")
    return True

def check_postgres_data():
    """Check if PostgreSQL has the required data"""
    try:
        conn = psycopg2.connect(
            host=os.getenv('POSTGRES_HOST'),
            port=int(os.getenv('POSTGRES_PORT')),
            user=os.getenv('POSTGRES_USER'),
            password=os.getenv('POSTGRES_PASSWORD'),
            database=os.getenv('POSTGRES_DB')
        )
        cursor = conn.cursor()
        
        # Check if price_tiers table exists and has data
        cursor.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_schema = 'public' 
                AND table_name = 'price_tiers'
            );
        """)
        table_exists = cursor.fetchone()[0]
        
        if not table_exists:
            logger.warning("⚠️ price_tiers table does not exist")
            cursor.close()
            conn.close()
            return False
        
        # Check if price_tiers has data
        cursor.execute('SELECT COUNT(*) FROM price_tiers;')
        count = cursor.fetchone()[0]
        
        if count == 0:
            logger.warning("⚠️ price_tiers table is empty")
            cursor.close()
            conn.close()
            return False
        
        # Check stack_recommendations (but don't fail if empty due to foreign key constraints)
        cursor.execute('SELECT COUNT(*) FROM stack_recommendations;')
        rec_count = cursor.fetchone()[0]
        
        # Check price_based_stacks instead (this is what actually gets populated)
        cursor.execute('SELECT COUNT(*) FROM price_based_stacks;')
        stacks_count = cursor.fetchone()[0]
        
        if stacks_count < 10:
            logger.warning(f"⚠️ price_based_stacks has only {stacks_count} records")
            cursor.close()
            conn.close()
            return False
        
        logger.info(f"✅ Found {stacks_count} price-based stacks and {rec_count} stack recommendations")
        
        cursor.close()
        conn.close()
        logger.info("✅ PostgreSQL data validation passed")
        return True
        
    except Exception as e:
        logger.error(f"❌ PostgreSQL data check failed: {e}")
        return False

def run_neo4j_migration():
    """Run Neo4j migration"""
    logger.info("🔄 Running Neo4j migration...")
    
    try:
        # Add src to path
        sys.path.append('src')
        
        from postgres_to_neo4j_migration import PostgresToNeo4jMigration
        
        # Configuration
        postgres_config = {
            'host': os.getenv('POSTGRES_HOST'),
            'port': int(os.getenv('POSTGRES_PORT')),
            'user': os.getenv('POSTGRES_USER'),
            'password': os.getenv('POSTGRES_PASSWORD'),
            'database': os.getenv('POSTGRES_DB')
        }
        
        neo4j_config = {
            'uri': os.getenv('NEO4J_URI'),
            'user': os.getenv('NEO4J_USER'),
            'password': os.getenv('NEO4J_PASSWORD')
        }
        
        # Run migration with TSS namespace
        migration = PostgresToNeo4jMigration(postgres_config, neo4j_config, namespace='TSS')
        success = migration.run_full_migration()
        
        if success:
            logger.info("✅ Neo4j migration completed successfully")
            return True
        else:
            logger.error("❌ Neo4j migration failed")
            return False
            
    except Exception as e:
        logger.error(f"❌ Neo4j migration error: {e}")
        return False

def check_neo4j_data():
    """Check if Neo4j has the required data"""
    try:
        driver = GraphDatabase.driver(
            os.getenv('NEO4J_URI'),
            auth=(os.getenv('NEO4J_USER'), os.getenv('NEO4J_PASSWORD'))
        )
        
        with driver.session() as session:
            # Check for TSS namespaced data specifically
            result = session.run('MATCH (p:PriceTier:TSS) RETURN count(p) as tss_price_tiers')
            tss_price_tiers = result.single()['tss_price_tiers']
            
            result = session.run('MATCH (t:Technology:TSS) RETURN count(t) as tss_technologies')
            tss_technologies = result.single()['tss_technologies']
            
            result = session.run('MATCH ()-[r:TSS_BELONGS_TO_TIER]->() RETURN count(r) as tss_relationships')
            tss_relationships = result.single()['tss_relationships']
            
            # Check if we have sufficient data
            if tss_price_tiers == 0:
                logger.warning("⚠️ No TSS price tiers found in Neo4j")
                driver.close()
                return False
                
            if tss_technologies == 0:
                logger.warning("⚠️ No TSS technologies found in Neo4j")
                driver.close()
                return False
                
            if tss_relationships == 0:
                logger.warning("⚠️ No TSS price tier relationships found in Neo4j")
                driver.close()
                return False
            
            logger.info(f"✅ Found {tss_price_tiers} TSS price tiers, {tss_technologies} TSS technologies, {tss_relationships} TSS relationships")
            driver.close()
            return True
        
    except Exception as e:
        logger.error(f"❌ Neo4j data check failed: {e}")
        return False

def run_tss_namespace_migration():
    """Run TSS namespace migration"""
    logger.info("🔄 Running TSS namespace migration...")
    
    try:
        result = subprocess.run([
            sys.executable, 'src/migrate_to_tss_namespace.py'
        ], capture_output=True, text=True)
        
        if result.returncode == 0:
            logger.info("✅ TSS namespace migration completed")
            return True
        else:
            logger.error(f"❌ TSS namespace migration failed: {result.stderr}")
            return False
            
    except Exception as e:
        logger.error(f"❌ TSS namespace migration error: {e}")
        return False

def main():
    """Main setup function"""
    logger.info("🚀 Starting Tech Stack Selector database setup...")
    
    # Setup environment variables
    setup_environment()
    
    # Check connections
    if not check_postgres_connection():
        logger.error("❌ Cannot proceed without PostgreSQL connection")
        sys.exit(1)
    
    if not check_neo4j_connection():
        logger.error("❌ Cannot proceed without Neo4j connection")
        sys.exit(1)
    
    # Run PostgreSQL migrations
    if not run_postgres_migrations():
        logger.error("❌ PostgreSQL migrations failed")
        sys.exit(1)
    
    # Check PostgreSQL data
    if not check_postgres_data():
        logger.error("❌ PostgreSQL data validation failed")
        sys.exit(1)
    
    # Check if Neo4j migration is needed
    if not check_neo4j_data():
        logger.info("🔄 Neo4j data not found, running migration...")
        if not run_neo4j_migration():
            logger.error("❌ Neo4j migration failed")
            sys.exit(1)
    else:
        logger.info("✅ Neo4j data already exists")
    
    # Run TSS namespace migration
    if not run_tss_namespace_migration():
        logger.error("❌ TSS namespace migration failed")
        sys.exit(1)
    
    logger.info("✅ Database setup completed successfully!")
    logger.info("🚀 Ready to start Tech Stack Selector service")

if __name__ == "__main__":
    main()
