#!/usr/bin/env python3
"""
Run Selenium workspace isolation test on port 3015.
"""
import os
import sys
import time
import urllib.request
import urllib.error

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

def is_server_running(url='http://localhost:3015', timeout=5):
    """Check if the server is reachable."""
    try:
        response = urllib.request.urlopen(url, timeout=timeout)
        return response.status < 500
    except urllib.error.URLError:
        return False
    except Exception:
        return False

def test_workspace_ui_isolation():
    """Test workspace isolation via UI."""
    if not is_server_running():
        print("LightRAG server not running on http://localhost:3015. Skipping Selenium test.")
        return
    
    driver = None
    try:
        # Initialize Chrome driver
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')  # Run in headless mode for CI
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        driver = webdriver.Chrome(options=options)
        driver.implicitly_wait(5)
        
        # Open LightRAG UI
        driver.get('http://localhost:3015')
        
        # Wait for page to load
        wait = WebDriverWait(driver, 10)
        
        # Check if workspace selector is present (look for SelectTrigger)
        try:
            workspace_selector = wait.until(
                EC.presence_of_element_located((By.CSS_SELECTOR, '.SelectTrigger'))
            )
            print("Workspace selector found.")
        except TimeoutException:
            print("Workspace selector NOT found. UI may not have workspace components.")
            # Take a screenshot for debugging
            driver.save_screenshot('workspace_ui_missing.png')
            print("Screenshot saved to workspace_ui_missing.png")
            raise
        
        # Click to open dropdown
        workspace_selector.click()
        
        # Wait for dropdown menu (SelectContent)
        dropdown = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, '.SelectContent'))
        )
        
        # Find the create workspace button (DialogTrigger)
        create_btn = dropdown.find_element(By.CSS_SELECTOR, 'button[aria-label="Create workspace"]')
        create_btn.click()
        
        # Wait for dialog
        dialog = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, '.DialogContent'))
        )
        
        # Fill workspace name
        name_input = dialog.find_element(By.CSS_SELECTOR, 'input[placeholder="Workspace name"]')
        workspace1_name = "test_workspace_1"
        name_input.send_keys(workspace1_name)
        
        # Submit (click Create button)
        submit_btn = dialog.find_element(By.XPATH, '//button[contains(text(), "Create")]')
        submit_btn.click()
        
        # Wait for workspace to be selected (UI updates)
        time.sleep(2)
        
        # Verify workspace is selected (optional)
        selected_workspace = driver.find_element(By.CSS_SELECTOR, '.SelectTrigger span')
        assert workspace1_name in selected_workspace.text, f"Workspace {workspace1_name} not selected"
        
        # Now create a second workspace
        workspace_selector.click()
        dropdown = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, '.SelectContent'))
        )
        create_btn = dropdown.find_element(By.CSS_SELECTOR, 'button[aria-label="Create workspace"]')
        create_btn.click()
        
        dialog = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, '.DialogContent'))
        )
        name_input = dialog.find_element(By.CSS_SELECTOR, 'input[placeholder="Workspace name"]')
        workspace2_name = "test_workspace_2"
        name_input.send_keys(workspace2_name)
        submit_btn = dialog.find_element(By.XPATH, '//button[contains(text(), "Create")]')
        submit_btn.click()
        
        time.sleep(2)
        
        # Verify workspace 2 is selected
        selected_workspace = driver.find_element(By.CSS_SELECTOR, '.SelectTrigger span')
        assert workspace2_name in selected_workspace.text, f"Workspace {workspace2_name} not selected"
        
        # Switch back to workspace 1
        workspace_selector.click()
        dropdown = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, '.SelectContent'))
        )
        # Find workspace 1 in list (SelectItem)
        workspace_items = dropdown.find_elements(By.CSS_SELECTOR, '.SelectItem')
        for item in workspace_items:
            if workspace1_name in item.text:
                item.click()
                break
        
        time.sleep(2)
        
        # Verify workspace 1 is selected again
        selected_workspace = driver.find_element(By.CSS_SELECTOR, '.SelectTrigger span')
        assert workspace1_name in selected_workspace.text, f"Workspace {workspace1_name} not selected after switch"
        
        print("✓ UI workspace isolation test passed!")
        
    except Exception as e:
        print(f"Test failed with error: {e}")
        raise
    finally:
        if driver:
            driver.quit()

if __name__ == "__main__":
    test_workspace_ui_isolation()