#!/usr/bin/env python3
"""
Apply the license_heartbeats table migration
"""

import os
import psycopg2
import logging
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_db_connection():
    """Get database connection"""
    return psycopg2.connect(
        host=os.environ.get('POSTGRES_HOST', 'postgres'),
        database=os.environ.get('POSTGRES_DB', 'v2_adminpanel'),
        user=os.environ.get('POSTGRES_USER', 'postgres'),
        password=os.environ.get('POSTGRES_PASSWORD', 'postgres')
    )

def apply_migration():
    """Apply the license_heartbeats migration"""
    conn = None
    try:
        logger.info("Connecting to database...")
        conn = get_db_connection()
        cur = conn.cursor()
        
        # Read migration file
        migration_file = os.path.join(os.path.dirname(__file__), 'migrations', 'create_license_heartbeats_table.sql')
        logger.info(f"Reading migration file: {migration_file}")
        
        with open(migration_file, 'r') as f:
            migration_sql = f.read()
        
        # Execute migration
        logger.info("Executing migration...")
        cur.execute(migration_sql)
        
        # Verify table was created
        cur.execute("""
            SELECT EXISTS (
                SELECT 1 
                FROM information_schema.tables 
                WHERE table_name = 'license_heartbeats'
            )
        """)
        
        if cur.fetchone()[0]:
            logger.info("✓ license_heartbeats table created successfully!")
            
            # Check partitions
            cur.execute("""
                SELECT tablename 
                FROM pg_tables 
                WHERE tablename LIKE 'license_heartbeats_%'
                ORDER BY tablename
            """)
            
            partitions = cur.fetchall()
            logger.info(f"✓ Created {len(partitions)} partitions:")
            for partition in partitions:
                logger.info(f"  - {partition[0]}")
        else:
            logger.error("✗ Failed to create license_heartbeats table")
            return False
        
        conn.commit()
        logger.info("✓ Migration completed successfully!")
        return True
        
    except Exception as e:
        logger.error(f"✗ Migration failed: {str(e)}")
        if conn:
            conn.rollback()
        return False
    finally:
        if conn:
            cur.close()
            conn.close()

if __name__ == "__main__":
    logger.info("=== Applying license_heartbeats migration ===")
    logger.info(f"Timestamp: {datetime.now()}")
    
    if apply_migration():
        logger.info("=== Migration successful! ===")
    else:
        logger.error("=== Migration failed! ===")
        exit(1)