#!/usr/bin/env python3
# Database initialization script
import os
import sys
import logging
from sqlalchemy import create_engine, text
from config import get_config

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

def create_database():
    """Create the database if it doesn't exist"""
    config = get_config()
    db_url = config.DATABASE_URL
    
    if db_url.startswith('postgresql'):
        # Parse database name from URL
        parts = db_url.split('/')
        db_name = parts[-1].split('?')[0]
        base_url = '/'.join(parts[:-1])
        
        # Connect to postgres database to create our database
        engine = create_engine(f"{base_url}/postgres")
        
        try:
            with engine.connect() as conn:
                # Check if database exists
                result = conn.execute(
                    text("SELECT 1 FROM pg_database WHERE datname = :dbname"),
                    {"dbname": db_name}
                )
                exists = result.fetchone() is not None
                
                if not exists:
                    # Create database
                    conn.execute(text(f"CREATE DATABASE {db_name}"))
                    logger.info(f"Database '{db_name}' created successfully")
                else:
                    logger.info(f"Database '{db_name}' already exists")
                    
        except Exception as e:
            logger.error(f"Error creating database: {e}")
            return False
        finally:
            engine.dispose()
    
    return True

def check_redis():
    """Check Redis connectivity"""
    config = get_config()
    
    try:
        import redis
        r = redis.from_url(config.REDIS_URL)
        r.ping()
        logger.info("Redis connection successful")
        return True
    except Exception as e:
        logger.error(f"Redis connection failed: {e}")
        return False

def init_database_extensions():
    """Initialize PostgreSQL extensions"""
    config = get_config()
    
    if not config.DATABASE_URL.startswith('postgresql'):
        return True
    
    engine = create_engine(config.DATABASE_URL)
    
    try:
        with engine.connect() as conn:
            # Enable UUID extension
            conn.execute(text("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\""))
            logger.info("PostgreSQL extensions initialized")
            
    except Exception as e:
        logger.error(f"Error initializing extensions: {e}")
        return False
    finally:
        engine.dispose()
    
    return True

def main():
    """Main initialization function"""
    logger.info("Starting database initialization...")
    
    # Create database
    if not create_database():
        logger.error("Failed to create database")
        sys.exit(1)
    
    # Initialize extensions
    if not init_database_extensions():
        logger.error("Failed to initialize database extensions")
        sys.exit(1)
    
    # Check Redis
    if not check_redis():
        logger.warning("Redis not available - caching will be disabled")
    
    logger.info("Database initialization completed successfully")
    
    # Create all tables using SQLAlchemy models
    logger.info("Creating database tables...")
    try:
        from flask import Flask
        from database import db, init_db
        from config import get_config
        
        # Import all models to ensure they're registered
        from auth_models import User, LoginHistory, UserSession, RevokedToken
        
        # Create Flask app context
        app = Flask(__name__)
        config = get_config()
        app.config.from_mapping(config.__dict__)
        
        # Initialize database
        init_db(app)
        
        with app.app_context():
            # Create all tables
            db.create_all()
            logger.info("Database tables created successfully")
            
    except Exception as e:
        logger.error(f"Failed to create database tables: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()