main.py

#!/usr/bin/env python3
"""
Interactive CLI Application
Main entry point for RAG Framework modules
"""
import os
import sys
import subprocess
import json
import yaml
from pathlib import Path
from typing import List, Dict, Optional

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from utils.console import (
    clear, print_box, print_error, print_success, 
    print_warning, print_info, wait_for_enter, interactive_menu, Colors
)

class InteractiveCLI:
    def __init__(self):
        self.data_input_path = Path("data/input")
        self.modules_path = Path("modules")
        self.selected_dataset = None
        self.selected_module = None
        
        # RAG Module configurations
        self.modules = {
            "chunking_engine": {
                "name": "Document Chunker",
                "script": "modules/chunking_engine/batch_processor_app.py",
                "description": "Chunk documents for RAG indexing and retrieval",
                "processes": 16  # Default number of processes
            },
            "embeddings_generator": {
                "name": "Embeddings Generator",
                "script": "modules/embeddings_generator/batch_processor_app.py",
                "description": "Generate embeddings from chunked data using vLLM server",
                "processes": 16  # Default number of processes
            },
            "elasticsearch_indexer": {
                "name": "Elasticsearch Indexer",
                "script": "modules/elasticsearch_indexer/app.py",
                "description": "Index chunked documents into Elasticsearch for search and retrieval",
                "processes": 1  # Elasticsearch indexing is single-threaded
            },
            "embeddings_indexer": {
                "name": "Embeddings Indexer",
                "script": "modules/embeddings_indexer/app.py",
                "description": "Index embeddings into Qdrant for vector search and retrieval",
                "processes": 1  # Qdrant indexing is single-threaded
            }
            # Add more modules here in the future
        }
        
        # Elasticsearch connection status
        self.es_connected = self.check_elasticsearch_connection()
    
    def parse_docker_compose(self) -> Dict:
        """Parse docker-compose.yml and extract service information"""
        try:
            with open("docker-compose.yml", "r") as file:
                compose_data = yaml.safe_load(file)
                return compose_data.get("services", {})
        except FileNotFoundError:
            print_error("docker-compose.yml file not found!")
            return {}
        except yaml.YAMLError as e:
            print_error(f"Error parsing docker-compose.yml: {str(e)}")
            return {}
    
    def format_service_name(self, service_key: str) -> str:
        """Format service name by removing prefix and converting to title case"""
        # Remove "digital_fortress_" prefix
        if service_key.startswith("digital_fortress_"):
            service_key = service_key[len("digital_fortress_"):]
        
        # Replace - and _ with spaces and title case
        formatted = service_key.replace("-", " ").replace("_", " ").title()
        return formatted
    
    def get_service_status(self, service_name: str) -> Dict:
        """Get Docker service status using docker compose commands"""
        try:
            # Get service status
            result = subprocess.run(
                ["docker", "compose", "ps", "--format", "json", service_name],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            if result.returncode == 0 and result.stdout.strip():
                # Parse JSON output
                service_info = json.loads(result.stdout.strip())
                if isinstance(service_info, list) and len(service_info) > 0:
                    service_info = service_info[0]
                
                status = service_info.get("State", "unknown")
                health = service_info.get("Health", "unknown")
                
                return {
                    "status": status,
                    "health": health
                }
            else:
                # Service not running
                return {
                    "status": "not_running",
                    "health": "not_available"
                }
        except subprocess.TimeoutExpired:
            return {
                "status": "error",
                "health": "timeout"
            }
        except Exception as e:
            return {
                "status": "error",
                "health": str(e)
            }
    
    def get_datasets(self) -> List[str]:
        """Get list of available document collections"""
        if not self.data_input_path.exists():
            return []
        
        datasets = [d.name for d in self.data_input_path.iterdir() if d.is_dir()]
        return sorted(datasets)
    
    def check_elasticsearch_connection(self) -> bool:
        """Check Elasticsearch connection status"""
        try:
            from elasticsearch import Elasticsearch
            from dotenv import load_dotenv
            import os
            
            # Load environment variables
            load_dotenv()
            
            es_host = os.getenv('ELASTIC_HOST', 'localhost')
            es_port = int(os.getenv('ELASTIC_PORT', 9200))
            es_scheme = os.getenv('ELASTIC_SCHEME', 'http')
            es_security = os.getenv('ELASTIC_SECURITY_ENABLED', 'false').lower() == 'true'
            
            if es_security:
                es_user = os.getenv('ELASTIC_USER', 'elastic')
                es_pass = os.getenv('ELASTIC_PASSWORD', 'changeme')
                es = Elasticsearch(
                    hosts=[{'host': es_host, 'port': es_port, 'scheme': es_scheme}],
                    basic_auth=(es_user, es_pass),
                    verify_certs=False
                )
            else:
                es = Elasticsearch(
                    hosts=[{'host': es_host, 'port': es_port, 'scheme': es_scheme}],
                    verify_certs=False
                )
            
            # Test connection
            info = es.info()
            return True
            
        except Exception as e:
            return False
    
    def display_menu(self, title: str, options: List[str], allow_back: bool = True) -> Optional[int]:
        """Display interactive menu and get user selection using arrow keys"""
        return interactive_menu(title, options, allow_back)
    
    def select_dataset(self) -> bool:
        """Document collection selection menu"""
        datasets = self.get_datasets()
        
        if not datasets:
            clear()
            print_error(f"No document collections found in {self.data_input_path}")
            print_info("Please add document collections to /data/input/<collection_name>")
            wait_for_enter()
            return False
        
        choice = self.display_menu(
            "SELECT DOCUMENT COLLECTION",
            datasets,
            allow_back=True
        )
        
        if choice is None:
            return False
        
        self.selected_dataset = datasets[choice]
        print_success(f"Selected document collection: {self.selected_dataset}")
        return True
    
    def select_module(self) -> bool:
        """Module selection menu"""
        if not self.selected_dataset:
            print_error("No document collection selected!")
            wait_for_enter()
            return False
        
        module_options = []
        module_keys = []
        
        for key, module in self.modules.items():
            module_options.append(f"{module['name']} - {module['description']}")
            module_keys.append(key)
        
        choice = self.display_menu(
            f"SELECT RAG MODULE (Collection: {self.selected_dataset})",
            module_options,
            allow_back=True
        )
        
        if choice is None:
            return False
        
        self.selected_module = module_keys[choice]
        print_success(f"Selected RAG module: {self.modules[self.selected_module]['name']}")
        return True
    
    def configure_module(self) -> Dict:
        """Configure RAG module parameters"""
        if self.selected_module == "chunking_engine":
            clear()
            print_box(f"CONFIGURE {self.modules[self.selected_module]['name'].upper()}")
            print()
            
            # Get number of processes
            while True:
                try:
                    print_info(f"Enter number of processes (1-{os.cpu_count()}) or press Enter for default [16]:")
                    num_processes = input("> ").strip()
                    
                    if not num_processes:
                        num_processes = 16
                    else:
                        num_processes = int(num_processes)
                    
                    if 1 <= num_processes <= os.cpu_count():
                        break
                    else:
                        print_error(f"Please enter a number between 1 and {os.cpu_count()}")
                except ValueError:
                    print_error("Please enter a valid number")
                except KeyboardInterrupt:
                    print("\n\nExiting...")
                    sys.exit(0)
            
            return {"processes": num_processes}
        
        elif self.selected_module == "embeddings_generator":
            clear()
            print_box(f"CONFIGURE {self.modules[self.selected_module]['name'].upper()}")
            print()
            
            # Get number of processes
            while True:
                try:
                    print_info(f"Enter number of processes (1-{os.cpu_count()}) or press Enter for default [16]:")
                    num_processes = input("> ").strip()
                    
                    if not num_processes:
                        num_processes = 16
                    else:
                        num_processes = int(num_processes)
                    
                    if 1 <= num_processes <= os.cpu_count():
                        break
                    else:
                        print_error(f"Please enter a number between 1 and {os.cpu_count()}")
                except ValueError:
                    print_error("Please enter a valid number")
                except KeyboardInterrupt:
                    print("\n\nExiting...")
                    sys.exit(0)
            
            return {"processes": num_processes}
        
        elif self.selected_module == "embeddings_generator":
            clear()
            print_box(f"CONFIGURE {self.modules[self.selected_module]['name'].upper()}")
            print()
            
            # Get number of processes
            while True:
                try:
                    print_info(f"Enter number of processes (1-{os.cpu_count()}) or press Enter for default [16]:")
                    num_processes = input("> ").strip()
                    
                    if not num_processes:
                        num_processes = 16
                    else:
                        num_processes = int(num_processes)
                    
                    if 1 <= num_processes <= os.cpu_count():
                        break
                    else:
                        print_error(f"Please enter a number between 1 and {os.cpu_count()}")
                except ValueError:
                    print_error("Please enter a valid number")
                except KeyboardInterrupt:
                    print("\n\nExiting...")
                    sys.exit(0)
            
            return {"processes": num_processes}
        
        elif self.selected_module == "embeddings_generator":
            clear()
            print_box(f"CONFIGURE {self.modules[self.selected_module]['name'].upper()}")
            print()
            
            # Get number of processes
            while True:
                try:
                    print_info(f"Enter number of processes (1-{os.cpu_count()}) or press Enter for default [16]:")
                    num_processes = input("> ").strip()
                    
                    if not num_processes:
                        num_processes = 16
                    else:
                        num_processes = int(num_processes)
                    
                    if 1 <= num_processes <= os.cpu_count():
                        break
                    else:
                        print_error(f"Please enter a number between 1 and {os.cpu_count()}")
                except ValueError:
                    print_error("Please enter a valid number")
                except KeyboardInterrupt:
                    print("\n\nExiting...")
                    sys.exit(0)
            
            return {"processes": num_processes}
        
        elif self.selected_module == "elasticsearch_indexer":
            clear()
            print_box(f"CONFIGURE {self.modules[self.selected_module]['name'].upper()}")
            print()
            
            # Elasticsearch indexing is single-threaded
            print_info("Elasticsearch bulk indexing is optimized for single-threaded operations")
            print_info("This ensures data consistency and optimal performance")
            print_info("Press Enter to continue with 1 process:")
            
            while True:
                try:
                    num_processes = input("> ").strip()
                    
                    if not num_processes:
                        num_processes = 1
                    else:
                        num_processes = int(num_processes)
                        if num_processes != 1:
                            print_warning("Elasticsearch indexing will use 1 process regardless of input")
                            num_processes = 1
                    
                    break
                except ValueError:
                    print_error("Please enter a valid number or press Enter for default")
                except KeyboardInterrupt:
                    print("\n\nExiting...")
                    sys.exit(0)
            
            return {"processes": num_processes}
        
        elif self.selected_module == "embeddings_indexer":
            clear()
            print_box(f"CONFIGURE {self.modules[self.selected_module]['name'].upper()}")
            print()
            
            # Qdrant indexing is single-threaded
            print_info("Qdrant bulk indexing is optimized for single-threaded operations")
            print_info("This ensures data consistency and optimal performance")
            print_info("Press Enter to continue with 1 process:")
            
            while True:
                try:
                    num_processes = input("> ").strip()
                    
                    if not num_processes:
                        num_processes = 1
                    else:
                        num_processes = int(num_processes)
                        if num_processes != 1:
                            print_warning("Qdrant indexing will use 1 process regardless of input")
                            num_processes = 1
                    
                    break
                except ValueError:
                    print_error("Please enter a valid number or press Enter for default")
                except KeyboardInterrupt:
                    print("\n\nExiting...")
                    sys.exit(0)
            
            return {"processes": num_processes}
        
        return {}
    
    def run_module(self):
        """Execute the selected module"""
        if not self.selected_dataset or not self.selected_module:
            print_error("Document collection and RAG module must be selected!")
            wait_for_enter()
            return
        
        # Configure module
        config = self.configure_module()
        
        # Build command
        dataset_path = self.data_input_path / self.selected_dataset
        module_info = self.modules[self.selected_module]
        
        cmd = [
            sys.executable,
            module_info["script"],
            str(dataset_path),
            str(config.get("processes", 16))
        ]
        
        # Set environment variables for the subprocess
        env = os.environ.copy()
        env['DATASET_NAME'] = self.selected_dataset
        env['MODULE_NAME'] = module_info['name']
        
        clear()
        print_info(f"Starting {module_info['name']}...")
        print_info(f"Document Collection: {self.selected_dataset}")
        print_info(f"Processes: {config.get('processes', 16)}")
        print()
        
        try:
            # Run the module
            result = subprocess.run(
                cmd,
                env=env,
                capture_output=False,
                text=True
            )
            
            if result.returncode == 0:
                print_success("Module completed successfully!")
            else:
                print_error(f"Module exited with code {result.returncode}")
                
        except FileNotFoundError:
            print_error(f"Module script not found: {module_info['script']}")
        except KeyboardInterrupt:
            print_warning("Module execution interrupted by user")
        except Exception as e:
            print_error(f"Failed to run module: {str(e)}")
        
        wait_for_enter()
    
    def main_menu(self):
        """Main menu loop"""
        while True:
            options = ["Select Document Collection", "Select RAG Module", "Run Module", "Service", "Exit"]
            
            # Add current selections and ES status to menu title
            menu_title = "RAG FRAMEWORK MAIN MENU"
            
            # Add Elasticsearch status
            es_status = f"{Colors.GREEN}Connected{Colors.END}" if self.es_connected else f"{Colors.RED}Not Connected{Colors.END}"
            menu_title += f"\n\nElasticsearch: {es_status}"
            
            if self.selected_dataset:
                menu_title += f"\nCurrent Collection: {self.selected_dataset}"
            if self.selected_module:
                menu_title += f"\nCurrent Module: {self.modules[self.selected_module]['name']}"
            
            choice = self.display_menu(menu_title, options, allow_back=False)
            
            if choice is None or choice == 4:  # Exit or None (escape/quit)
                clear()

                # Beautiful farewell message
                print(f"\n{Colors.BOLD}{Colors.CYAN}{'=' * 60}{Colors.END}")
                print(f"{Colors.BOLD}{Colors.GREEN}{' ' * 20}👋 See You Later! {' ' * 20}{Colors.END}")
                print(f"{Colors.BOLD}{Colors.CYAN}{'=' * 60}{Colors.END}")
                print()
                print(f"{Colors.YELLOW}{' ' * 15}Thanks for using our RAG Framework!{Colors.END}")
                print(f"{Colors.CYAN}{' ' * 10}Your data processing adventures await! 🚀{Colors.END}")
                print(f"{Colors.GREEN}{' ' * 18}Have a wonderful day! 😊{Colors.END}")
                print()
                print(f"{Colors.BOLD}{Colors.BLUE}{'=' * 60}{Colors.END}")
                print()

                sys.exit(0)
            elif choice == 0:  # Select Document Collection
                self.select_dataset()
            elif choice == 1:  # Select RAG Module
                self.select_module()
            elif choice == 2:  # Run Module
                self.run_module()
            elif choice == 3:  # Service
                self.service_management()
    
    def service_management(self):
        """Service management menu"""
        # Parse docker-compose.yml
        services = self.parse_docker_compose()
        
        if not services:
            print_error("No services found in docker-compose.yml")
            wait_for_enter()
            return
        
        # Create service options with health status
        service_options = []
        service_keys = list(services.keys())
        
        for service_key in service_keys:
            formatted_name = self.format_service_name(service_key)
            status_info = self.get_service_status(service_key)
            
            # Determine color based on health status
            if status_info["health"] == "healthy":
                health_display = f"{Colors.GREEN}Healthy{Colors.END}"
            elif status_info["health"] == "not_available":
                health_display = f"{Colors.YELLOW}Not Available{Colors.END}"
            elif status_info["health"] == "starting":
                health_display = f"{Colors.YELLOW}Starting{Colors.END}"
            elif status_info["health"] == "unhealthy":
                health_display = f"{Colors.RED}Unhealthy{Colors.END}"
            else:
                health_display = f"{Colors.RED}Critical{Colors.END}"
            
            service_options.append(f"{formatted_name} ({health_display})")
        
        # Display menu
        choice = self.display_menu(
            "SERVICE MANAGEMENT",
            service_options,
            allow_back=True
        )
        
        if choice is None:
            return
        
        # Show service details before starting
        selected_service = service_keys[choice]
        self.show_service_details(selected_service, services[selected_service])
        
        # Confirm start
        print()
        print_info("Do you want to start this service? (y/N):")
        response = input("> ").strip().lower()
        
        if response in ['y', 'yes']:
            self.start_service(selected_service)
    
    def show_service_details(self, service_key: str, service_config: Dict):
        """Display detailed information about a service"""
        clear()
        formatted_name = self.format_service_name(service_key)
        status_info = self.get_service_status(service_key)
        
        print_box(f"SERVICE DETAILS: {formatted_name}")
        print()
        
        # Display current status
        print(f"{Colors.BOLD}Current Status:{Colors.END}")
        if status_info["health"] == "healthy":
            status_display = f"{Colors.GREEN}Healthy{Colors.END}"
        elif status_info["health"] == "not_available":
            status_display = f"{Colors.YELLOW}Not Available{Colors.END}"
        elif status_info["health"] == "starting":
            status_display = f"{Colors.YELLOW}Starting{Colors.END}"
        elif status_info["health"] == "unhealthy":
            status_display = f"{Colors.RED}Unhealthy{Colors.END}"
        else:
            status_display = f"{Colors.RED}Critical{Colors.END}"
        
        print(f"  Status: {status_display}")
        print()
        
        # Display configuration details
        print(f"{Colors.BOLD}Configuration:{Colors.END}")
        
        # Image or build info
        if "image" in service_config:
            print(f"  Image: {service_config['image']}")
        elif "build" in service_config:
            build_info = service_config["build"]
            if isinstance(build_info, dict):
                print(f"  Build Context: {build_info.get('context', 'N/A')}")
                print(f"  Dockerfile: {build_info.get('dockerfile', 'Dockerfile')}")
            else:
                print(f"  Build: {build_info}")
        
        # Ports
        if "ports" in service_config:
            print(f"  Ports: {', '.join(service_config['ports'])}")
        
        # Environment variables (just count for brevity)
        if "environment" in service_config:
            env_count = len(service_config["environment"]) if isinstance(service_config["environment"], list) else "N/A"
            print(f"  Environment Variables: {env_count}")
        
        # Volumes
        if "volumes" in service_config:
            print(f"  Volumes: {len(service_config['volumes'])} mounted")
        
        # Healthcheck
        if "healthcheck" in service_config:
            healthcheck = service_config["healthcheck"]
            if isinstance(healthcheck, dict) and "test" in healthcheck:
                test_cmd = healthcheck["test"]
                if isinstance(test_cmd, list):
                    test_cmd = " ".join(test_cmd)
                print(f"  Healthcheck: {test_cmd}")
    
    def start_service(self, service_name: str):
        """Start a Docker service using docker compose up"""
        clear()
        print_info(f"Starting service: {self.format_service_name(service_name)}")
        print_info("This may take a few moments...")
        print()
        
        try:
            # Start service in detached mode
            result = subprocess.run(
                ["docker", "compose", "up", "-d", service_name],
                capture_output=True,
                text=True
            )
            
            if result.returncode == 0:
                print_success(f"Service '{service_name}' started successfully!")
                print_info("You can check the service status in the Service Management menu.")
            else:
                print_error(f"Failed to start service '{service_name}'")
                print_error(f"Error: {result.stderr}")
                
        except Exception as e:
            print_error(f"Error starting service: {str(e)}")
        
        wait_for_enter()
    
    def run(self):
        """Start the RAG Framework CLI application"""
        try:
            # Check if data directory exists
            if not self.data_input_path.exists():
                print_error(f"Data directory not found: {self.data_input_path}")
                print_info("Creating directory structure...")
                self.data_input_path.mkdir(parents=True, exist_ok=True)
                print_success("Directory created. Please add document collections and restart.")
                wait_for_enter()
                return
            
            # Start main menu
            self.main_menu()
            
        except KeyboardInterrupt:
            clear()
            print("\nExiting...")
            sys.exit(0)
        except Exception as e:
            print_error(f"Unexpected error: {str(e)}")
            wait_for_enter()
            sys.exit(1)

if __name__ == "__main__":
    cli = InteractiveCLI()
    cli.run()