#!/usr/bin/env python3
"""
Update database schema for enhanced automation models
"""
import sys
import os
from sqlalchemy import create_engine, text

# Add the project root to Python path
sys.path.insert(0, '/home/pi/Desktop/Server_Monitorizare_v2')

from config.database_config import DatabaseConfig

class DatabaseSchemaUpdater:
    def __init__(self):
        db_config = DatabaseConfig()
        self.engine = db_config.engine
    
    def update_playbook_executions_schema(self):
        """Add new columns to playbook_executions table"""
        print("📊 Updating playbook_executions table schema...")
        
        new_columns = [
            ("playbook_description", "TEXT"),
            ("estimated_duration", "INTEGER"),
            ("priority", "INTEGER DEFAULT 5"),
            ("retry_count", "INTEGER DEFAULT 0"),
            ("max_retries", "INTEGER DEFAULT 0"),
            ("summary_message", "TEXT"),
            ("skipped_hosts", "INTEGER DEFAULT 0"),
            ("changed_hosts", "INTEGER DEFAULT 0")
        ]
        
        with self.engine.connect() as conn:
            # Check if table exists
            result = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='playbook_executions'"))
            if not result.fetchone():
                print("   ⚠️  playbook_executions table doesn't exist. Creating tables...")
                self.create_tables()
                return
            
            for column_name, column_type in new_columns:
                try:
                    # Check if column already exists
                    check_sql = f"PRAGMA table_info(playbook_executions)"
                    result = conn.execute(text(check_sql))
                    existing_columns = [row[1] for row in result.fetchall()]
                    
                    if column_name not in existing_columns:
                        alter_sql = f"ALTER TABLE playbook_executions ADD COLUMN {column_name} {column_type}"
                        conn.execute(text(alter_sql))
                        conn.commit()
                        print(f"   ✅ Added column: {column_name}")
                    else:
                        print(f"   ⏭️  Column {column_name} already exists")
                        
                except Exception as e:
                    print(f"   ❌ Error adding {column_name}: {e}")
    
    def create_tables(self):
        """Create all tables using SQLAlchemy metadata"""
        print("🏗️  Creating all database tables...")
        
        try:
            from app.models import Base
            Base.metadata.create_all(self.engine)
            print("   ✅ All tables created successfully")
        except Exception as e:
            print(f"   ❌ Error creating tables: {e}")
    
    def verify_schema(self):
        """Verify the schema updates were successful"""
        print("\n🔍 Verifying schema updates...")
        
        with self.engine.connect() as conn:
            try:
                # Test with a simple query 
                result = conn.execute(text("SELECT COUNT(*) FROM playbook_executions"))
                count = result.fetchone()[0]
                print(f"   ✅ playbook_executions table accessible: {count} records")
                
                # Check for new columns
                result = conn.execute(text("PRAGMA table_info(playbook_executions)"))
                columns = [row[1] for row in result.fetchall()]
                
                expected_columns = [
                    'playbook_description', 'estimated_duration', 'priority',
                    'retry_count', 'max_retries', 'summary_message',
                    'skipped_hosts', 'changed_hosts'
                ]
                
                missing_columns = [col for col in expected_columns if col not in columns]
                if missing_columns:
                    print(f"   ⚠️  Missing columns: {missing_columns}")
                else:
                    print("   ✅ All expected columns present")
                
                return len(missing_columns) == 0
                
            except Exception as e:
                print(f"   ❌ Schema verification failed: {e}")
                return False

if __name__ == "__main__":
    updater = DatabaseSchemaUpdater()
    
    print("🔧 Database Schema Update for Enhanced Automation")
    print("=" * 60)
    
    # Update schema
    updater.update_playbook_executions_schema()
    
    # Verify 
    if updater.verify_schema():
        print("\n🎉 Schema update completed successfully!")
    else:
        print("\n❌ Schema update verification failed!")
        sys.exit(1)