#!/usr/bin/env python3
"""
Apply partition migration for license_heartbeats table.
This script creates missing partitions for the current and future months.
"""

import psycopg2
import os
import sys
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

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 create_partition(cursor, year, month):
    """Create a partition for the given year and month"""
    partition_name = f"license_heartbeats_{year}_{month:02d}"
    start_date = f"{year}-{month:02d}-01"
    
    # Calculate end date (first day of next month)
    if month == 12:
        end_date = f"{year + 1}-01-01"
    else:
        end_date = f"{year}-{month + 1:02d}-01"
    
    # Check if partition already exists
    cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM pg_tables 
            WHERE tablename = %s
        )
    """, (partition_name,))
    
    exists = cursor.fetchone()[0]
    
    if not exists:
        try:
            cursor.execute(f"""
                CREATE TABLE {partition_name} PARTITION OF license_heartbeats
                FOR VALUES FROM ('{start_date}') TO ('{end_date}')
            """)
            print(f"✓ Created partition {partition_name}")
            return True
        except Exception as e:
            print(f"✗ Error creating partition {partition_name}: {e}")
            return False
    else:
        print(f"- Partition {partition_name} already exists")
        return False

def main():
    """Main function"""
    print("Applying license_heartbeats partition migration...")
    print("-" * 50)
    
    try:
        # Connect to database
        conn = get_db_connection()
        cursor = conn.cursor()
        
        # Check if license_heartbeats table exists
        cursor.execute("""
            SELECT EXISTS (
                SELECT 1 
                FROM information_schema.tables 
                WHERE table_name = 'license_heartbeats'
            )
        """)
        
        if not cursor.fetchone()[0]:
            print("✗ Error: license_heartbeats table does not exist!")
            print("  Please run the init.sql script first.")
            return 1
        
        # Get current date
        current_date = datetime.now()
        partitions_created = 0
        
        # Create partitions for the next 6 months (including current month)
        for i in range(7):
            target_date = current_date + relativedelta(months=i)
            if create_partition(cursor, target_date.year, target_date.month):
                partitions_created += 1
        
        # Commit changes
        conn.commit()
        
        print("-" * 50)
        print(f"✓ Migration complete. Created {partitions_created} new partitions.")
        
        # List all partitions
        cursor.execute("""
            SELECT tablename 
            FROM pg_tables 
            WHERE tablename LIKE 'license_heartbeats_%'
            ORDER BY tablename
        """)
        
        partitions = cursor.fetchall()
        print(f"\nTotal partitions: {len(partitions)}")
        for partition in partitions:
            print(f"  - {partition[0]}")
        
        cursor.close()
        conn.close()
        
        return 0
        
    except Exception as e:
        print(f"✗ Error: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())