import sys
import time
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QVBoxLayout, QLineEdit, QTextEdit, QLabel
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from uied_detector import capture_and_detect
from gpt_interaction import process_command
from action_performer import perform_action, switch_to_chrome
import pyautogui
from ui import CommandInput
import os

class CommandThread(QThread):
    update_output = pyqtSignal(str)
    finished = pyqtSignal()

    def __init__(self, command, template):
        super().__init__()
        self.command = command
        self.template = template
        self.stop_flag = False

    def run(self):
        step_count = 0
        current_state = "Initial state"
        action_history = []
        while step_count < 20 and not self.stop_flag:  # Increased max steps
            self.update_output.emit(f"Step {step_count + 1}:")
            
            # Switch to Chrome window at the beginning of each step
            if not switch_to_chrome():
                self.update_output.emit("Warning: Failed to switch to Chrome window. Continuing execution.")
                # We continue execution instead of breaking

            annotated_path, elements = capture_and_detect()
            response = process_command(self.command, annotated_path, elements, current_state, action_history, self.template)
            
            self.update_output.emit("AI Thoughts:")
            self.update_output.emit(response["thoughts"])
            self.update_output.emit("\nActions:")
            
            if response["actions"] and not self.stop_flag:
                self.update_output.emit("Performing actions:")
                for action in response["actions"]:
                    if self.stop_flag:
                        break
                    self.update_output.emit(str(action))
                
                if not self.stop_flag:
                    action_results = perform_action(response["actions"], elements, annotated_path)
                    action_history.extend(action_results)
                    current_state = f"Performed actions: {', '.join(action_results)}"
            else:
                self.update_output.emit("No actions to perform. Requesting more information.")
                current_state = "No actions performed in this step. More information needed."
            
            step_count += 1
            
            if response.get("task_complete", False) or self.stop_flag:
                self.update_output.emit("Task completed or stopped.")
                break
            
            time.sleep(2)  # Wait before next iteration
        
        if step_count >= 20:
            self.update_output.emit("Reached maximum number of steps. Stopping.")
        
        self.finished.emit()

    def stop(self):
        self.stop_flag = True

class CommandWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.command_input = CommandInput()
        self.command_input.command_entered.connect(self.process_command)
        self.initUI()
        self.command_thread = None
        self.current_task = ""

    def initUI(self):
        self.setWindowTitle('Command Input')
        self.setWindowFlags(Qt.Window | Qt.WindowStaysOnTopHint)
        layout = QVBoxLayout()
        layout.addWidget(self.command_input)
        
        self.current_task_label = QLabel("Current task: None", self)
        self.current_task_label.setFixedHeight(30)  # Set a fixed height for the label
        layout.addWidget(self.current_task_label)
        
        self.output = QTextEdit(self)
        self.output.setReadOnly(True)
        layout.addWidget(self.output)
        
        self.stop_button = QPushButton('STOP', self)
        self.stop_button.clicked.connect(self.stop_command)
        self.stop_button.setEnabled(False)
        layout.addWidget(self.stop_button)
        
        self.setLayout(layout)
        self.setGeometry(300, 300, 500, 400)

    def process_command(self, command, template):
        self.output.clear()
        self.output.append(f"Processing command: {command}")
        if template:
            with open(os.path.join("templates", template), "r") as f:
                template_content = f.read()
            self.output.append(f"Using template: {template}")
        else:
            template_content = None
        
        self.command_input.setEnabled(False)
        self.stop_button.setEnabled(True)
        self.current_task = command
        self.update_task_label()
        
        self.command_thread = CommandThread(command, template_content)
        self.command_thread.update_output.connect(self.update_output)
        self.command_thread.finished.connect(self.on_command_finished)
        self.command_thread.start()

    def update_task_label(self):
        max_length = 30  # Maximum number of characters to display
        if len(self.current_task) > max_length:
            truncated_task = self.current_task[:max_length - 3] + "..."
        else:
            truncated_task = self.current_task
        self.current_task_label.setText(f"Current task: {truncated_task}")

    def stop_command(self):
        if self.command_thread and self.command_thread.isRunning():
            self.command_thread.stop()
            self.output.append("Stopping command execution...")
            self.on_command_finished()

    def update_output(self, text):
        self.output.append(text)

    def on_command_finished(self):
        self.command_input.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.current_task = ""
        self.current_task_label.setText("Current task: None")

def main():
    app = QApplication(sys.argv)
    window = CommandWindow()
    window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()