#!/usr/bin/env python3
"""
Check workspace selections in LightRAG UI using Selenium.
"""

import os
import sys
import time
import subprocess
import requests
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException

SERVER_URL = "http://localhost:3015"
USERNAME = "jleu3482"
PASSWORD = "jleu1212"

def start_server():
    """Start LightRAG server"""
    print("Starting server...")
    
    # Kill any existing server
    try:
        subprocess.run(["taskkill", "/F", "/IM", "python.exe"], capture_output=True)
    except:
        pass
    
    # Start server
    cmd = [sys.executable, "start_server_fixed.py"]
    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding='utf-8'
    )
    
    # Wait for server to start
    for i in range(30):
        try:
            requests.get(f"{SERVER_URL}/health", timeout=2)
            print(f"Server started (attempt {i+1})")
            return process
        except:
            time.sleep(1)
    
    print("Server failed to start")
    if process:
        process.terminate()
    return None

def check_workspace_ui():
    """Check workspace selections in UI using Selenium"""
    print("Setting up Selenium...")
    
    # Setup Chrome options
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')  # Run in headless mode
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    
    try:
        driver = webdriver.Chrome(options=options)
    except Exception as e:
        print(f"Chrome WebDriver failed: {e}")
        # Try Firefox
        try:
            options = webdriver.FirefoxOptions()
            options.add_argument('--headless')
            driver = webdriver.Firefox(options=options)
        except Exception as e2:
            print(f"Firefox WebDriver failed: {e2}")
            return False, "Could not start WebDriver"
    
    try:
        # Go to web UI
        print("Navigating to web UI...")
        driver.get(f"{SERVER_URL}/webui")
        time.sleep(3)
        
        # Take screenshot of initial page
        driver.save_screenshot("workspace_ui_initial.png")
        print("Screenshot saved: workspace_ui_initial.png")
        
        # Get page source for analysis
        page_source = driver.page_source
        with open("workspace_ui_page.html", "w", encoding="utf-8") as f:
            f.write(page_source)
        print("Page source saved: workspace_ui_page.html")
        
        # Look for workspace-related elements
        print("\nSearching for workspace UI elements...")
        
        # Check for workspace dropdown/selector
        workspace_selectors = [
            "workspace", "Workspace", "workspace-select", "workspace-selector",
            "select-workspace", "workspace-dropdown", "dropdown-workspace"
        ]
        
        found_elements = []
        
        # Search by various methods
        for selector in workspace_selectors:
            # By ID
            try:
                elements = driver.find_elements(By.ID, selector)
                if elements:
                    found_elements.append(f"ID: {selector} - Found {len(elements)} elements")
            except:
                pass
            
            # By class name
            try:
                elements = driver.find_elements(By.CLASS_NAME, selector)
                if elements:
                    found_elements.append(f"Class: {selector} - Found {len(elements)} elements")
            except:
                pass
            
            # By CSS selector
            try:
                elements = driver.find_elements(By.CSS_SELECTOR, f'[class*="{selector}"]')
                if elements:
                    found_elements.append(f"CSS class*={selector} - Found {len(elements)} elements")
            except:
                pass
            
            # By XPath containing text
            try:
                elements = driver.find_elements(By.XPATH, f'//*[contains(text(), "{selector}")]')
                if elements:
                    for elem in elements:
                        text = elem.text[:50] if elem.text else "no text"
                        found_elements.append(f"Text contains '{selector}': {text}")
            except:
                pass
        
        # Also check for common UI patterns
        common_patterns = [
            ("dropdown", "//select", "Select dropdown"),
            ("button with workspace", "//button[contains(text(), 'Workspace')]", "Workspace button"),
            ("menu item", "//*[contains(@class, 'menu')]//*[contains(text(), 'Workspace')]", "Workspace menu"),
            ("tab", "//*[contains(@class, 'tab')]//*[contains(text(), 'Workspace')]", "Workspace tab"),
        ]
        
        for pattern_name, xpath, description in common_patterns:
            try:
                elements = driver.find_elements(By.XPATH, xpath)
                if elements:
                    found_elements.append(f"{description}: Found {len(elements)} elements")
            except:
                pass
        
        # Check API for workspaces
        print("\nChecking workspace API...")
        try:
            response = requests.get(f"{SERVER_URL}/api/workspaces", timeout=10)
            if response.status_code == 200:
                workspaces = response.json()
                print(f"API returned {len(workspaces)} workspaces: {workspaces}")
                
                # If workspaces exist but UI doesn't show them, there might be a UI issue
                if workspaces and len(workspaces) > 0 and not found_elements:
                    print("WARNING: Workspaces exist in API but not visible in UI")
            else:
                print(f"Workspace API returned {response.status_code}")
        except Exception as e:
            print(f"Error checking workspace API: {e}")
        
        # Also check the workspace manager endpoint
        try:
            response = requests.get(f"{SERVER_URL}/api/workspace-manager", timeout=10)
            print(f"Workspace manager endpoint: {response.status_code}")
        except:
            print("Workspace manager endpoint not accessible")
        
        # Print findings
        print("\n" + "="*60)
        print("WORKSPACE UI CHECK RESULTS")
        print("="*60)
        
        if found_elements:
            print("Found workspace-related UI elements:")
            for elem in found_elements:
                print(f"  - {elem}")
            
            # Take screenshot of found elements
            for i, elem_desc in enumerate(found_elements[:5]):  # Limit to first 5
                try:
                    # Extract selector type and try to highlight
                    if "ID:" in elem_desc:
                        selector = elem_desc.split("ID: ")[1].split(" -")[0]
                        elements = driver.find_elements(By.ID, selector)
                    elif "Class:" in elem_desc:
                        selector = elem_desc.split("Class: ")[1].split(" -")[0]
                        elements = driver.find_elements(By.CLASS_NAME, selector)
                    
                    if elements and len(elements) > 0:
                        # Highlight element with JavaScript
                        driver.execute_script("arguments[0].style.border='3px solid red'", elements[0])
                        time.sleep(0.5)
                except:
                    pass
            
            driver.save_screenshot("workspace_ui_highlighted.png")
            print("\nHighlighted screenshot saved: workspace_ui_highlighted.png")
            
            result = True
            message = f"Found {len(found_elements)} workspace UI elements"
        else:
            print("NO workspace UI elements found!")
            print("\nPossible issues:")
            print("1. Workspace feature might be disabled in configuration")
            print("2. UI might need recompilation/rebuild")
            print("3. Workspace elements might have different CSS classes/IDs")
            print("4. JavaScript might not be loading properly")
            
            # Check if JavaScript is working
            js_check = driver.execute_script("return typeof window !== 'undefined'")
            print(f"\nJavaScript loaded: {js_check}")
            
            # Check console errors
            logs = driver.get_log('browser')
            if logs:
                print(f"\nBrowser console errors ({len(logs)}):")
                for log in logs[:10]:  # Show first 10
                    print(f"  {log.get('level', 'UNKNOWN')}: {log.get('message', '')[:200]}")
            
            result = False
            message = "No workspace UI elements found"
        
        # Check page title and structure
        print(f"\nPage title: {driver.title}")
        print(f"Current URL: {driver.current_url}")
        
        # Look for any React/Vue/Angular indicators
        body_html = driver.find_element(By.TAG_NAME, "body").get_attribute("innerHTML")
        if "react" in body_html.lower():
            print("React detected in page")
        if "vue" in body_html.lower():
            print("Vue detected in page")
        if "angular" in body_html.lower():
            print("Angular detected in page")
        
        driver.quit()
        return result, message
        
    except Exception as e:
        print(f"Error during Selenium test: {e}")
        import traceback
        traceback.print_exc()
        try:
            driver.quit()
        except:
            pass
        return False, f"Selenium error: {str(e)}"

def main():
    """Main function"""
    print("="*60)
    print("LightRAG Workspace UI Check")
    print("="*60)
    
    # Start server
    server_process = start_server()
    if not server_process:
        print("Failed to start server")
        return False
    
    time.sleep(3)  # Give server time to initialize
    
    # Check workspace UI
    success, message = check_workspace_ui()
    
    # Cleanup
    print("\nCleaning up...")
    if server_process:
        server_process.terminate()
        try:
            server_process.wait(timeout=5)
        except:
            pass
    
    print("\n" + "="*60)
    print("SUMMARY")
    print("="*60)
    print(f"Result: {'SUCCESS' if success else 'FAILURE'}")
    print(f"Message: {message}")
    
    if not success:
        print("\nRECOMMENDED ACTIONS:")
        print("1. Check if workspace feature is enabled in server configuration")
        print("2. Verify the web UI has been built with workspace support")
        print("3. Check browser console for JavaScript errors")
        print("4. Inspect the saved HTML file: workspace_ui_page.html")
        print("5. Check screenshots: workspace_ui_initial.png, workspace_ui_highlighted.png")
    
    return success

if __name__ == "__main__":
    result = main()
    sys.exit(0 if result else 1)