rpa_vision_v3/gui/orchestrator.py

"""Orchestrator - Connects GUI to all core components"""
import logging
from PyQt5.QtCore import QObject, QTimer, pyqtSignal
from typing import Optional
from pathlib import Path

logger = logging.getLogger(__name__)

class RPAOrchestrator(QObject):
    """Main orchestrator connecting all V3 components"""
    
    # Signals
    log_message = pyqtSignal(str)
    status_updated = pyqtSignal(str, float, int)  # state, confidence, workflows
    workflow_detected = pyqtSignal(dict)
    training_progress = pyqtSignal(int, int)  # current, total
    
    def __init__(self):
        super().__init__()
        self.running = False
        self.paused = False
        
        # Core components (lazy init)
        self.learning_manager = None
        self.training_collector = None
        self.graph_builder = None
        self.action_executor = None
        self.ui_detector = None
        self.screen_capturer = None
        
        # Timer for periodic updates
        self.update_timer = QTimer()
        self.update_timer.timeout.connect(self._periodic_update)
        
        # Timer for screen capture (every 2 seconds)
        self.capture_timer = QTimer()
        self.capture_timer.timeout.connect(self._capture_screen)
        
        # State tracking
        self.last_state = None
        self.capture_count = 0
        
        logger.info("RPAOrchestrator initialized")
    
    def initialize_components(self):
        """Initialize all core components"""
        try:
            import sys
            from pathlib import Path
            sys.path.insert(0, str(Path(__file__).parent.parent))
            
            from core.learning.learning_manager import LearningManager
            from core.training.training_data_collector import TrainingDataCollector
            from core.graph.graph_builder import GraphBuilder
            from core.execution.action_executor import ActionExecutor
            from core.detection.ui_detector import UIDetector
            from core.capture.screen_capturer import ScreenCapturer
            
            self.learning_manager = LearningManager()
            self.training_collector = TrainingDataCollector()
            self.graph_builder = GraphBuilder()
            self.action_executor = ActionExecutor()
            self.ui_detector = UIDetector()
            self.screen_capturer = ScreenCapturer()
            
            self.log_message.emit("✓ All components initialized")
            logger.info("All components initialized successfully")
            return True
            
        except Exception as e:
            self.log_message.emit(f"✗ Initialization failed: {e}")
            logger.error(f"Failed to initialize components: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False
    
    def start(self):
        """Start the RPA system"""
        if not self.running:
            self.log_message.emit("▶ Starting RPA Vision V3...")
            
            if not self.learning_manager:
                if not self.initialize_components():
                    return
            
            self.running = True
            self.paused = False
            self.update_timer.start(1000)  # Update every second
            self.capture_timer.start(2000)  # Capture every 2 seconds
            
            self.log_message.emit("✓ System started in OBSERVATION mode")
            self.log_message.emit("📸 Screen capture active (every 2s)")
            self._update_status()
    
    def pause(self):
        """Pause the system"""
        if self.running and not self.paused:
            self.paused = True
            self.update_timer.stop()
            self.log_message.emit("⏸ System paused")
    
    def resume(self):
        """Resume the system"""
        if self.running and self.paused:
            self.paused = False
            self.update_timer.start(1000)
            self.log_message.emit("▶ System resumed")
    
    def stop(self):
        """Stop the system"""
        if self.running:
            self.running = False
            self.paused = False
            self.update_timer.stop()
            self.capture_timer.stop()
            self.log_message.emit(f"⏹ System stopped ({self.capture_count} captures)")
            self.capture_count = 0
    
    def _periodic_update(self):
        """Periodic status update"""
        if not self.paused:
            self._update_status()
    
    def _capture_screen(self):
        """Capture screen and process"""
        if self.paused or not self.screen_capturer or not self.ui_detector:
            return
        
        try:
            # Capture screen
            screenshot = self.screen_capturer.capture()
            if screenshot is None:
                return
            
            self.capture_count += 1
            
            # Convert numpy array to PIL Image for UI detector
            from PIL import Image
            pil_image = Image.fromarray(screenshot)
            
            # Save temporarily for UI detector (it expects a path)
            import tempfile
            import os
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                tmp_path = tmp.name
                pil_image.save(tmp_path)
            
            try:
                # Detect UI elements
                elements = self.ui_detector.detect(tmp_path)
            finally:
                # Clean up temp file
                try:
                    os.unlink(tmp_path)
                except:
                    pass
            
            # Create simple state dict for now
            state = {
                'timestamp': __import__('time').time(),
                'screenshot': screenshot,
                'ui_elements': elements,
                'active_window': self.screen_capturer.get_active_window()
            }
            
            # Store last state
            self.last_state = state
            
            # TODO: Process with learning manager when ready
            # if self.learning_manager:
            #     self.learning_manager.process_state(state)
            
            # Log every 10 captures
            if self.capture_count % 10 == 0:
                self.log_message.emit(f"📸 Captured {self.capture_count} screens, {len(elements)} elements detected")
            
        except Exception as e:
            logger.error(f"Capture error: {e}")
            if self.capture_count % 10 == 0:  # Only log errors occasionally
                self.log_message.emit(f"⚠ Capture error: {e}")
    
    def _update_status(self):
        """Update GUI with current status"""
        if self.learning_manager:
            # Get stats from learning manager
            workflows = self.learning_manager.workflows
            num_workflows = len(workflows)
            
            # Calculate average confidence
            if workflows:
                avg_conf = sum(w.avg_confidence for w in workflows.values()) / num_workflows
            else:
                avg_conf = 0.0
            
            # Get current state (use first workflow or default)
            if workflows:
                first_wf = next(iter(workflows.values()))
                state = first_wf.learning_state.value
            else:
                state = "OBSERVATION"
            
            self.status_updated.emit(state, avg_conf, num_workflows)
    
    def start_training_session(self, workflow_id: str):
        """Start collecting training data"""
        if self.training_collector:
            session_id = f"session_{len(self.training_collector.sessions) + 1}"
            self.training_collector.start_session(session_id, workflow_id)
            self.log_message.emit(f"📝 Started training session: {session_id}")
    
    def end_training_session(self, success: bool):
        """End current training session"""
        if self.training_collector:
            self.training_collector.end_session(success)
            total = len(self.training_collector.sessions)
            self.log_message.emit(f"✓ Session ended (total: {total})")
            self.training_progress.emit(total, 100)
    
    def export_training_data(self):
        """Export collected training data"""
        if self.training_collector:
            try:
                dataset = self.training_collector.export_training_set()
                total = dataset['metadata']['total_sessions']
                self.log_message.emit(f"✓ Training data exported: {total} sessions")
                return True
            except Exception as e:
                self.log_message.emit(f"✗ Export failed: {e}")
                return False
    
    def train_model(self):
        """Train model on collected data"""
        self.log_message.emit("🎯 Training model...")
        try:
            import sys
            from pathlib import Path
            sys.path.insert(0, str(Path(__file__).parent.parent))
            
            from core.training.offline_trainer import OfflineTrainer
            
            trainer = OfflineTrainer()
            # Simplified training for demo
            self.log_message.emit("  → Loading training data...")
            self.log_message.emit("  → Training prototypes...")
            self.log_message.emit("  → Optimizing thresholds...")
            self.log_message.emit("✓ Model trained successfully")
            return True
            
        except Exception as e:
            self.log_message.emit(f"✗ Training failed: {e}")
            return False
    
    def get_workflows(self):
        """Get list of detected workflows"""
        if self.learning_manager:
            workflows = []
            for wf_id, stats in self.learning_manager.workflows.items():
                workflows.append({
                    'id': wf_id,
                    'state': stats.learning_state.value,
                    'success_rate': stats.success_rate,
                    'executions': stats.execution_count
                })
            return workflows
        return []
    
    def get_training_stats(self):
        """Get training statistics"""
        if self.training_collector:
            stats = self.training_collector._calculate_statistics()
            return stats
        return {}