#!/usr/bin/env python
"""
Fix missing database tables by running migrations and creating any missing tables.
This script should be run on the production server when database issues occur.

Usage:
    python fix_missing_tables.py
"""

import os
import sys
import django
import logging
from django.db import connection, ProgrammingError, OperationalError

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

# Setup Django environment
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'wtlms.settings')
django.setup()

from django.core.management import call_command
from django.db.migrations.recorder import MigrationRecorder

def check_table_exists(table_name):
    """Check if a table exists in the database."""
    with connection.cursor() as cursor:
        try:
            cursor.execute(f"SELECT 1 FROM {table_name} LIMIT 1;")
            return True
        except (ProgrammingError, OperationalError):
            return False

def fix_django_session_table():
    """Create the django_session table if it doesn't exist."""
    logger.info("Checking for django_session table...")
    
    if check_table_exists('django_session'):
        logger.info("django_session table already exists.")
        return
    
    logger.info("django_session table doesn't exist. Creating it...")
    
    # Create the django_session table
    with connection.cursor() as cursor:
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS "django_session" (
                "session_key" varchar(40) NOT NULL PRIMARY KEY,
                "session_data" text NOT NULL,
                "expire_date" datetime NOT NULL
            );
        """)
        
        # Create index
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS "django_session_expire_date_a5c62663" 
            ON "django_session" ("expire_date");
        """)
    
    logger.info("django_session table created successfully.")

def fix_django_migrations_table():
    """Create the django_migrations table if it doesn't exist."""
    logger.info("Checking for django_migrations table...")
    
    if check_table_exists('django_migrations'):
        logger.info("django_migrations table already exists.")
        return
    
    logger.info("django_migrations table doesn't exist. Creating it...")
    
    # Create the django_migrations table
    with connection.cursor() as cursor:
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS "django_migrations" (
                "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
                "app" varchar(255) NOT NULL,
                "name" varchar(255) NOT NULL,
                "applied" datetime NOT NULL
            );
        """)
    
    logger.info("django_migrations table created successfully.")

def run_migrations():
    """Run all migrations to ensure all tables are created."""
    logger.info("Running migrations...")
    try:
        call_command('migrate', '--noinput')
        logger.info("Migrations completed successfully.")
    except Exception as e:
        logger.error(f"Error running migrations: {e}")
        logger.info("Attempting to fix specific tables...")
        fix_django_session_table()
        fix_django_migrations_table()
        
        # Try migrations again
        try:
            logger.info("Running migrations again...")
            call_command('migrate', '--noinput')
            logger.info("Migrations completed successfully on second attempt.")
        except Exception as e:
            logger.error(f"Error running migrations on second attempt: {e}")
            logger.info("Manual intervention may be required.")

def check_essential_tables():
    """Check if essential Django tables exist."""
    essential_tables = [
        'django_session',
        'django_migrations',
        'auth_user',
        'auth_group',
        'django_content_type'
    ]
    
    missing_tables = []
    for table in essential_tables:
        if not check_table_exists(table):
            missing_tables.append(table)
    
    if missing_tables:
        logger.warning(f"Missing essential tables: {', '.join(missing_tables)}")
    else:
        logger.info("All essential tables exist.")
    
    return missing_tables

def main():
    """Main function to fix database issues."""
    logger.info("Starting database fix script...")
    
    # Check for essential tables
    missing_tables = check_essential_tables()
    
    if missing_tables:
        # Run migrations to create missing tables
        run_migrations()
        
        # Check again after migrations
        remaining_missing = check_essential_tables()
        if remaining_missing:
            logger.warning(f"Tables still missing after migrations: {', '.join(remaining_missing)}")
            logger.info("Attempting to fix specific tables...")
            
            if 'django_session' in remaining_missing:
                fix_django_session_table()
            
            if 'django_migrations' in remaining_missing:
                fix_django_migrations_table()
    else:
        logger.info("No missing essential tables detected.")
        
        # Run migrations anyway to ensure all tables are up to date
        logger.info("Running migrations to ensure all tables are up to date...")
        run_migrations()
    
    logger.info("Database fix script completed.")

if __name__ == "__main__":
    main()