#!/usr/bin/env python
"""
Fix notifications table - add missing action_required column
Run this in production to fix the OperationalError
"""
import os
import sys
import django

# Setup Django
try:
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'config.settings')
    django.setup()
except Exception:
    # Try alternative settings
    try:
        os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'settings')
        django.setup()
    except Exception:
        print("Error: Could not setup Django. Please set DJANGO_SETTINGS_MODULE")
        sys.exit(1)

from django.db import connection

print("="*70)
print("Fix Notifications Table - Add Missing action_required Column")
print("="*70)
print()

try:
    with connection.cursor() as cursor:
        # Check if column exists
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE()
            AND TABLE_NAME = 'utils_notification'
            AND COLUMN_NAME = 'action_required'
        """)
        
        exists = cursor.fetchone()[0] > 0
        
        if exists:
            print("✓ Column 'action_required' already exists")
        else:
            print("⚠ Column 'action_required' does not exist. Adding it...")
            
            # Add the column
            cursor.execute("""
                ALTER TABLE `utils_notification` 
                ADD COLUMN `action_required` tinyint(1) NOT NULL DEFAULT 0
                COMMENT 'Whether this notification requires immediate action'
            """)
            
            print("✓ Column 'action_required' added successfully")
        
        # Check for other missing columns
        columns_to_check = [
            ('alert_data', 'json', 'DEFAULT NULL COMMENT \'Additional alert data in JSON format\''),
            ('portfolio_snapshot_id', 'char(32)', 'NULL'),
        ]
        
        for col_name, col_type, extra in columns_to_check:
            cursor.execute("""
                SELECT COUNT(*) 
                FROM information_schema.COLUMNS 
                WHERE TABLE_SCHEMA = DATABASE()
                AND TABLE_NAME = 'utils_notification'
                AND COLUMN_NAME = %s
            """, [col_name])
            
            exists = cursor.fetchone()[0] > 0
            
            if not exists:
                print(f"⚠ Column '{col_name}' does not exist. Adding it...")
                
                if 'portfolio_snapshot_id' in col_name:
                    # Foreign key column
                    cursor.execute(f"""
                        ALTER TABLE `utils_notification` 
                        ADD COLUMN `{col_name}` {col_type} {extra}
                    """)
                else:
                    cursor.execute(f"""
                        ALTER TABLE `utils_notification` 
                        ADD COLUMN `{col_name}` {col_type} {extra}
                    """)
                
                print(f"✓ Column '{col_name}' added successfully")
            else:
                print(f"✓ Column '{col_name}' already exists")
        
        print()
        print("="*70)
        print("✓ All columns verified/added successfully!")
        print("="*70)
        
except Exception as e:
    print()
    print("="*70)
    print(f"✗ Error: {str(e)}")
    print("="*70)
    sys.exit(1)