#!/root/venvs/itb/bin/python
import sys
import os
import time
import argparse
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains

class InputController:
    def __init__(self, debug_port: int, verbose: bool = False):
        """Initialize connection to Chrome instance"""
        self.verbose = verbose
        options = webdriver.ChromeOptions()
        options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}")
        self.driver = webdriver.Chrome(options=options)

    def log(self, msg: str):
        """Log message if verbose mode is enabled"""
        if self.verbose:
            print(f"DEBUG: {msg}", file=sys.stderr)

    def focus_element(self, element_id: str = None):
        """Focus on an input element"""
        try:
            if element_id:
                element = self.driver.find_element(By.ID, element_id)
            else:
                # Get currently focused element
                element = self.driver.switch_to.active_element
                
            element.click()
            return element
            
        except Exception as e:
            self.log(f"Failed to focus element: {str(e)}")
            return None

    def enter_text(self, text: str, element_id: str = None, append: bool = False):
        """Enter text into an input element"""
        try:
            element = self.focus_element(element_id)
            if not element:
                return False
                
            if not append:
                # Clear existing text first
                if element.tag_name.lower() in ['input', 'textarea']:
                    element.clear()
                else:
                    # For contenteditable elements
                    element.send_keys(Keys.CONTROL + 'a')
                    element.send_keys(Keys.DELETE)
            
            # Type the text
            element.send_keys(text)
            return True
            
        except Exception as e:
            self.log(f"Failed to enter text: {str(e)}")
            return False

    def set_caret_position(self, position: int, element_id: str = None):
        """Set the text cursor position"""
        try:
            element = self.focus_element(element_id)
            if not element:
                return False
                
            # Use JavaScript to set caret position
            self.driver.execute_script(
                """
                const el = arguments[0];
                const pos = arguments[1];
                if (el.setSelectionRange) {
                    el.setSelectionRange(pos, pos);
                } else if (el.createTextRange) {
                    const range = el.createTextRange();
                    range.collapse(true);
                    range.moveEnd('character', pos);
                    range.moveStart('character', pos);
                    range.select();
                }
                """,
                element, position
            )
            return True
            
        except Exception as e:
            self.log(f"Failed to set caret position: {str(e)}")
            return False

    def select_text(self, start: int, end: int, element_id: str = None):
        """Select text range in an input element"""
        try:
            element = self.focus_element(element_id)
            if not element:
                return False
                
            # Use JavaScript to set selection
            self.driver.execute_script(
                """
                const el = arguments[0];
                const start = arguments[1];
                const end = arguments[2];
                if (el.setSelectionRange) {
                    el.setSelectionRange(start, end);
                } else if (el.createTextRange) {
                    const range = el.createTextRange();
                    range.collapse(true);
                    range.moveStart('character', start);
                    range.moveEnd('character', end);
                    range.select();
                }
                """,
                element, start, end
            )
            return True
            
        except Exception as e:
            self.log(f"Failed to select text: {str(e)}")
            return False

def main():
    parser = argparse.ArgumentParser(description='Control text input in Chrome')
    parser.add_argument('debug_port', type=int, help='Chrome debug port')
    parser.add_argument('-i', '--id', help='Element ID')
    parser.add_argument('-t', '--text', help='Text to input')
    parser.add_argument('--append', action='store_true', help='Append text instead of replacing')
    parser.add_argument('--caret', type=int, help='Set caret position')
    parser.add_argument('--select-start', type=int, help='Text selection start')
    parser.add_argument('--select-end', type=int, help='Text selection end')
    parser.add_argument('-v', '--verbose', action='store_true', help='Enable debug output')
    
    args = parser.parse_args()
    
    # Validate args
    if args.select_start is not None and args.select_end is None:
        print("Must specify both --select-start and --select-end", file=sys.stderr)
        sys.exit(1)
        
    if args.select_end is not None and args.select_start is None:
        print("Must specify both --select-start and --select-end", file=sys.stderr)
        sys.exit(1)
        
    try:
        controller = InputController(args.debug_port, args.verbose)
        success = False
        
        if args.text is not None:
            success = controller.enter_text(args.text, args.id, args.append)
        elif args.caret is not None:
            success = controller.set_caret_position(args.caret, args.id)
        elif args.select_start is not None:
            success = controller.select_text(args.select_start, args.select_end, args.id)
        else:
            success = controller.focus_element(args.id) is not None
            
        sys.exit(0 if success else 1)
        
    except Exception as e:
        if args.verbose:
            print(f"Fatal error: {str(e)}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()