api/run/initialization/neo4j.py

"""
Neo4j database initialization module for ClassroomCopilot
Handles database creation, schema setup, and calendar structure
"""
import os
import json
import time
from typing import Dict, Any, Optional
from datetime import datetime, timedelta
from modules.database.services.neo4j_service import Neo4jService
from modules.database.init.init_calendar import create_calendar
from modules.logger_tool import initialise_logger

logger = initialise_logger(__name__, os.getenv("LOG_LEVEL"), os.getenv("LOG_PATH"), 'default', True)

class Neo4jInitializer:
    """Handles Neo4j database initialization including database, schema, and calendar structure"""
    
    def __init__(self):
        self.neo4j_service = Neo4jService()
        # Use lowercase to avoid case sensitivity issues
        self.db_name = os.getenv("NEO4J_CC_DB", "classroomcopilot")
    
    def initialize_database(self) -> Dict[str, Any]:
        """Initialize the single ClassroomCopilot database"""
        logger.info(f"Initializing Neo4j database: {self.db_name}")
        
        try:
            # Create the main database
            result = self.neo4j_service.create_database(self.db_name)
            if result["status"] != "success":
                return {
                    "success": False,
                    "message": f"Failed to create {self.db_name} database: {result['message']}"
                }
            
            logger.info(f"Successfully created {self.db_name} database")
            
            # Wait for database to be fully available
            logger.info("Waiting for database to be fully available...")
            time.sleep(5)  # Wait 5 seconds for database to be ready
            
            # Verify database exists and is accessible
            max_retries = 10
            retry_delay = 2
            for attempt in range(max_retries):
                try:
                    # Try to check if the database exists to verify it's ready
                    test_result = self.neo4j_service.check_database_exists(self.db_name)
                    if test_result["status"] == "success" and test_result["exists"]:
                        logger.info(f"Database {self.db_name} is ready and accessible")
                        break
                    else:
                        logger.info(f"Database not ready yet (attempt {attempt + 1}/{max_retries}), waiting {retry_delay}s...")
                        time.sleep(retry_delay)
                except Exception as e:
                    if attempt < max_retries - 1:
                        logger.info(f"Database existence check failed (attempt {attempt + 1}/{max_retries}), waiting {retry_delay}s...")
                        time.sleep(retry_delay)
                    else:
                        logger.warning(f"Database existence check failed after {max_retries} attempts: {str(e)}")
                        # Continue anyway as the database was created successfully
            
            return {
                "success": True,
                "message": f"Successfully created {self.db_name} database",
                "result": result
            }
            
        except Exception as e:
            logger.error(f"Error creating database: {str(e)}")
            return {
                "success": False,
                "message": f"Error creating database: {str(e)}"
            }
    
    def initialize_schema(self) -> Dict[str, Any]:
        """Initialize Neo4j schema on the ClassroomCopilot database"""
        logger.info(f"Initializing Neo4j schema on {self.db_name}...")
        
        try:
            result = self.neo4j_service.initialize_schema(self.db_name)
            if result["status"] != "success":
                return {
                    "success": False,
                    "message": f"Failed to initialize schema on {self.db_name}: {result['message']}"
                }
            
            logger.info(f"Successfully initialized schema on {self.db_name}")
            return {
                "success": True,
                "message": f"Successfully initialized schema on {self.db_name}",
                "result": result
            }
            
        except Exception as e:
            logger.error(f"Error initializing schema: {str(e)}")
            return {
                "success": False,
                "message": f"Error initializing schema: {str(e)}"
            }
    
    def initialize_calendar_structure(self) -> Dict[str, Any]:
        """Initialize the calendar structure with days, weeks, months, and years"""
        logger.info("Initializing calendar structure...")
        
        try:
            # Create calendar structure for the next 1 year
            start_date = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
            end_date = start_date + timedelta(days=365)  # 1 year
            
            # Create calendar structure directly (no calendar node needed)
            calendar_nodes = create_calendar(
                db_name=self.db_name,
                start_date=start_date,
                end_date=end_date,
                time_chunk_node_length=0,
                storage_tools=None
            )
            
            if calendar_nodes:
                logger.info("Calendar structure created successfully")
                return {
                    "success": True,
                    "message": "Calendar structure created successfully",
                    "calendar_nodes": calendar_nodes
                }
            else:
                return {
                    "success": False,
                    "message": "Failed to create calendar structure"
                }
            
        except Exception as e:
            logger.error(f"Error creating calendar structure: {str(e)}")
            return {
                "success": False,
                "message": f"Error creating calendar structure: {str(e)}"
            }

def initialize_neo4j() -> Dict[str, Any]:
    """Initialize Neo4j database setup"""
    logger.info("Starting Neo4j database initialization...")
    
    initializer = Neo4jInitializer()
    
    # 1. Create the single database
    db_result = initializer.initialize_database()
    if not db_result["success"]:
        return db_result
    
    # 2. Initialize schema
    schema_result = initializer.initialize_schema()
    if not schema_result["success"]:
        return schema_result
    
    # 3. Initialize calendar structure
    calendar_result = initializer.initialize_calendar_structure()
    if not calendar_result["success"]:
        return calendar_result
    
    logger.info("Neo4j database initialization completed successfully")
    return {
        "success": True,
        "message": "Neo4j database initialization completed successfully",
        "database": db_result,
        "schema": schema_result,
        "calendar": calendar_result
    }