#!/usr/bin/env python3
"""
AI Analysis Service Database Migration Runner
Runs the database migration for AI Analysis Service during container startup.
"""

import os
import sys
import subprocess
import time
from pathlib import Path

def log(message):
    """Log with timestamp."""
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}")

def check_database_connection():
    """Check if database is available."""
    try:
        import psycopg2
        from dotenv import load_dotenv
        
        load_dotenv()
        
        conn = psycopg2.connect(
            host=os.getenv('POSTGRES_HOST', 'localhost'),
            port=os.getenv('POSTGRES_PORT', 5432),
            database=os.getenv('POSTGRES_DB', 'dev_pipeline'),
            user=os.getenv('POSTGRES_USER', 'pipeline_admin'),
            password=os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
        )
        conn.close()
        return True
    except Exception as e:
        log(f"Database connection failed: {e}")
        return False

def run_migration():
    """Run the database migration."""
    try:
        log("Starting AI Analysis Service database migration...")
        
        # Check if database is available
        max_retries = 30
        retry_count = 0
        
        while retry_count < max_retries:
            if check_database_connection():
                log("Database connection successful")
                break
            else:
                retry_count += 1
                log(f"Database not ready, retrying in 2 seconds... ({retry_count}/{max_retries})")
                time.sleep(2)
        else:
            log("ERROR: Could not connect to database after 60 seconds")
            return False
        
        # Run the migration script
        schema_file = Path(__file__).parent / "001-schema.sql"
        if not schema_file.exists():
            log("ERROR: Schema file not found")
            return False
        
        log(f"Running migration from {schema_file}")
        
        # Use psql to run the migration
        env = os.environ.copy()
        env['PGPASSWORD'] = os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
        
        result = subprocess.run([
            'psql',
            '-h', os.getenv('POSTGRES_HOST', 'localhost'),
            '-p', os.getenv('POSTGRES_PORT', '5432'),
            '-U', os.getenv('POSTGRES_USER', 'pipeline_admin'),
            '-d', os.getenv('POSTGRES_DB', 'dev_pipeline'),
            '-f', str(schema_file),
            '-v', 'ON_ERROR_STOP=1'
        ], env=env, capture_output=True, text=True)
        
        if result.returncode == 0:
            log("✅ AI Analysis Service database migration completed successfully")
            return True
        else:
            log(f"❌ Migration failed: {result.stderr}")
            return False
            
    except Exception as e:
        log(f"❌ Migration error: {e}")
        return False

if __name__ == "__main__":
    success = run_migration()
    sys.exit(0 if success else 1)