#!/usr/bin/env python3
"""
IDP System Project Runner
=========================
Unified script to run the complete IDP (Intelligent Document Processing) system.

This script:
1. Checks prerequisites (Python, Node.js, dependencies)
2. Starts the FastAPI backend server
3. Starts the Next.js frontend development server
4. Monitors both processes
5. Handles graceful shutdown

Usage:
    python run_project.py [options]

Options:
    --backend-only      Run only the backend server
    --frontend-only     Run only the frontend server
    --port-backend      Backend port (default: 7860)
    --port-frontend     Frontend port (default: 3000)
    --no-check          Skip dependency checks
    --help              Show this help message
"""

import os
import sys
import subprocess
import time
import signal
import argparse
import shutil
from pathlib import Path
from typing import Optional, List

# Color codes for terminal output
class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

class ProjectRunner:
    def __init__(self, args):
        self.args = args
        self.backend_process: Optional[subprocess.Popen] = None
        self.frontend_process: Optional[subprocess.Popen] = None
        self.project_root = Path(__file__).parent.absolute()
        self.frontend_dir = self.project_root / "frontend"
        
        # Environment variables for backend
        self.backend_env = os.environ.copy()
        self.backend_env.update({
            'KMP_DUPLICATE_LIB_OK': 'TRUE',
            'OMP_NUM_THREADS': '1',
            'OPENBLAS_NUM_THREADS': '1',
            'MKL_NUM_THREADS': '1',
            'VECLIB_MAXIMUM_THREADS': '1',
            'NUMEXPR_NUM_THREADS': '1',
            'TF_CPP_MIN_LOG_LEVEL': '2'
        })
        
    def print_header(self, message: str):
        """Print a formatted header message."""
        print(f"\n{Colors.HEADER}{Colors.BOLD}{'=' * 70}{Colors.ENDC}")
        print(f"{Colors.HEADER}{Colors.BOLD}{message:^70}{Colors.ENDC}")
        print(f"{Colors.HEADER}{Colors.BOLD}{'=' * 70}{Colors.ENDC}\n")
    
    def print_success(self, message: str):
        """Print a success message."""
        print(f"{Colors.OKGREEN}✓ {message}{Colors.ENDC}")
    
    def print_info(self, message: str):
        """Print an info message."""
        print(f"{Colors.OKCYAN}ℹ {message}{Colors.ENDC}")
    
    def print_warning(self, message: str):
        """Print a warning message."""
        print(f"{Colors.WARNING}⚠ {message}{Colors.ENDC}")
    
    def print_error(self, message: str):
        """Print an error message."""
        print(f"{Colors.FAIL}✗ {message}{Colors.ENDC}")
    
    def check_command(self, command: str) -> bool:
        """Check if a command is available in PATH."""
        return shutil.which(command) is not None
    
    def check_prerequisites(self):
        """Check if all prerequisites are installed."""
        if self.args.no_check:
            self.print_info("Skipping prerequisite checks (--no-check flag)")
            return True
        
        self.print_header("Checking Prerequisites")
        
        all_ok = True
        
        # Check Python
        if sys.version_info < (3, 10):
            self.print_error(f"Python 3.10+ required (found {sys.version_info.major}.{sys.version_info.minor})")
            all_ok = False
        else:
            self.print_success(f"Python {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}")
        
        # Check Node.js (only if frontend needed)
        if not self.args.backend_only:
            if self.check_command("node"):
                result = subprocess.run(["node", "--version"], capture_output=True, text=True)
                version = result.stdout.strip()
                self.print_success(f"Node.js {version}")
            else:
                self.print_error("Node.js not found (required for frontend)")
                all_ok = False
        
        # Check npm (only if frontend needed)
        if not self.args.backend_only:
            if self.check_command("npm"):
                result = subprocess.run(["npm", "--version"], capture_output=True, text=True)
                version = result.stdout.strip()
                self.print_success(f"npm {version}")
            else:
                self.print_error("npm not found (required for frontend)")
                all_ok = False
        
        # Check Python dependencies
        if not self.args.frontend_only:
            try:
                import fastapi
                self.print_success(f"FastAPI {fastapi.__version__}")
            except ImportError:
                self.print_warning("FastAPI not installed - run: pip install -r requirements.txt")
                all_ok = False
            
            try:
                import torch
                self.print_success(f"PyTorch {torch.__version__}")
            except ImportError:
                self.print_warning("PyTorch not installed - run: pip install -r requirements.txt")
                all_ok = False
        
        # Check frontend dependencies
        if not self.args.backend_only:
            node_modules = self.frontend_dir / "node_modules"
            if not node_modules.exists():
                self.print_warning("Frontend dependencies not installed - run: cd frontend && npm install")
                all_ok = False
            else:
                self.print_success("Frontend dependencies installed")
        
        if not all_ok:
            self.print_error("\nSome prerequisites are missing. Please install them and try again.")
            return False
        
        self.print_success("\nAll prerequisites check passed!")
        return True
    
    def start_backend(self):
        """Start the FastAPI backend server."""
        if self.args.frontend_only:
            return
        
        self.print_header("Starting Backend Server")
        self.print_info(f"Port: {self.args.port_backend}")
        self.print_info(f"Working directory: {self.project_root}")
        
        api_server_path = self.project_root / "api_server.py"
        if not api_server_path.exists():
            self.print_error(f"api_server.py not found at {api_server_path}")
            sys.exit(1)
        
        cmd = [
            sys.executable,
            "api_server.py"
        ]
        
        try:
            self.backend_process = subprocess.Popen(
                cmd,
                cwd=self.project_root,
                env=self.backend_env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True,
                bufsize=1
            )
            
            self.print_success("Backend server started!")
            self.print_info(f"API available at: http://localhost:{self.args.port_backend}")
            self.print_info(f"API docs at: http://localhost:{self.args.port_backend}/docs")
            
            # Give it a moment to start
            time.sleep(2)
            
            # Check if it's still running
            if self.backend_process.poll() is not None:
                self.print_error("Backend server failed to start")
                return False
            
            return True
            
        except Exception as e:
            self.print_error(f"Failed to start backend: {e}")
            return False
    
    def start_frontend(self):
        """Start the Next.js frontend development server."""
        if self.args.backend_only:
            return
        
        self.print_header("Starting Frontend Server")
        self.print_info(f"Port: {self.args.port_frontend}")
        self.print_info(f"Working directory: {self.frontend_dir}")
        
        if not self.frontend_dir.exists():
            self.print_error(f"Frontend directory not found at {self.frontend_dir}")
            sys.exit(1)
        
        # Set the API URL environment variable
        frontend_env = os.environ.copy()
        frontend_env['NEXT_PUBLIC_API_URL'] = f"http://localhost:{self.args.port_backend}"
        frontend_env['PORT'] = str(self.args.port_frontend)
        
        cmd = ["npm", "run", "dev"]
        
        try:
            self.frontend_process = subprocess.Popen(
                cmd,
                cwd=self.frontend_dir,
                env=frontend_env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True,
                bufsize=1
            )
            
            self.print_success("Frontend server started!")
            self.print_info(f"Application available at: http://localhost:{self.args.port_frontend}")
            
            # Give it a moment to start
            time.sleep(3)
            
            # Check if it's still running
            if self.frontend_process.poll() is not None:
                self.print_error("Frontend server failed to start")
                return False
            
            return True
            
        except Exception as e:
            self.print_error(f"Failed to start frontend: {e}")
            return False
    
    def monitor_processes(self):
        """Monitor running processes and display their output."""
        self.print_header("System Running")
        
        if not self.args.frontend_only:
            self.print_info(f"Backend: http://localhost:{self.args.port_backend}")
        if not self.args.backend_only:
            self.print_info(f"Frontend: http://localhost:{self.args.port_frontend}")
        
        print(f"\n{Colors.WARNING}Press Ctrl+C to stop all servers{Colors.ENDC}\n")
        
        try:
            while True:
                # Check backend
                if self.backend_process and self.backend_process.poll() is not None:
                    self.print_error("Backend process died unexpectedly!")
                    break
                
                # Check frontend
                if self.frontend_process and self.frontend_process.poll() is not None:
                    self.print_error("Frontend process died unexpectedly!")
                    break
                
                # Read and display output from backend
                if self.backend_process and self.backend_process.stdout:
                    line = self.backend_process.stdout.readline()
                    if line:
                        # Filter out mutex warnings
                        if "mutex.cc" not in line and "RAW: Lock blocking" not in line:
                            print(f"{Colors.OKBLUE}[Backend]{Colors.ENDC} {line.rstrip()}")
                
                # Read and display output from frontend
                if self.frontend_process and self.frontend_process.stdout:
                    line = self.frontend_process.stdout.readline()
                    if line:
                        print(f"{Colors.OKCYAN}[Frontend]{Colors.ENDC} {line.rstrip()}")
                
                time.sleep(0.1)
                
        except KeyboardInterrupt:
            self.print_info("\nReceived shutdown signal...")
            self.shutdown()
    
    def shutdown(self):
        """Gracefully shutdown all processes."""
        self.print_header("Shutting Down")
        
        if self.backend_process:
            self.print_info("Stopping backend server...")
            self.backend_process.terminate()
            try:
                self.backend_process.wait(timeout=5)
                self.print_success("Backend stopped")
            except subprocess.TimeoutExpired:
                self.print_warning("Backend didn't stop gracefully, forcing...")
                self.backend_process.kill()
        
        if self.frontend_process:
            self.print_info("Stopping frontend server...")
            self.frontend_process.terminate()
            try:
                self.frontend_process.wait(timeout=5)
                self.print_success("Frontend stopped")
            except subprocess.TimeoutExpired:
                self.print_warning("Frontend didn't stop gracefully, forcing...")
                self.frontend_process.kill()
        
        self.print_success("\nAll services stopped. Goodbye!")
    
    def run(self):
        """Main execution method."""
        # Register signal handlers
        signal.signal(signal.SIGINT, lambda s, f: None)  # Handle in monitor_processes
        signal.signal(signal.SIGTERM, lambda s, f: self.shutdown())
        
        # Print banner
        print(f"\n{Colors.BOLD}{Colors.HEADER}")
        print("╔════════════════════════════════════════════════════════════════════╗")
        print("║                  IDP System - Project Runner                      ║")
        print("║            Intelligent Document Processing Platform               ║")
        print("╚════════════════════════════════════════════════════════════════════╝")
        print(f"{Colors.ENDC}\n")
        
        # Check prerequisites
        if not self.check_prerequisites():
            sys.exit(1)
        
        # Start services
        success = True
        
        if not self.args.frontend_only:
            success = success and self.start_backend()
        
        if not self.args.backend_only:
            success = success and self.start_frontend()
        
        if not success:
            self.print_error("\nFailed to start all services")
            self.shutdown()
            sys.exit(1)
        
        # Monitor processes
        self.monitor_processes()

def main():
    parser = argparse.ArgumentParser(
        description='Run the IDP System (Backend + Frontend)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )
    
    parser.add_argument(
        '--backend-only',
        action='store_true',
        help='Run only the backend server'
    )
    
    parser.add_argument(
        '--frontend-only',
        action='store_true',
        help='Run only the frontend server'
    )
    
    parser.add_argument(
        '--port-backend',
        type=int,
        default=7860,
        help='Backend server port (default: 7860)'
    )
    
    parser.add_argument(
        '--port-frontend',
        type=int,
        default=3000,
        help='Frontend server port (default: 3000)'
    )
    
    parser.add_argument(
        '--no-check',
        action='store_true',
        help='Skip prerequisite checks'
    )
    
    args = parser.parse_args()
    
    # Validate arguments
    if args.backend_only and args.frontend_only:
        print(f"{Colors.FAIL}Error: Cannot use --backend-only and --frontend-only together{Colors.ENDC}")
        sys.exit(1)
    
    runner = ProjectRunner(args)
    runner.run()

if __name__ == "__main__":
    main()