#!/usr/bin/env python3
"""
Test Web UI authentication after fixing the configuration
"""

import requests
import time
import base64

# Configuration
BASE_URL = "http://localhost:3015"
WEBUI_URL = "http://localhost:3015/webui/"
USERNAME = "jleu3482"
PASSWORD = "jleu1212"

def wait_for_server():
    """Wait for server to be ready"""
    print("=== WAITING FOR SERVER ===")
    
    max_wait = 30
    for i in range(max_wait):
        try:
            response = requests.get(f"{BASE_URL}/health", timeout=5)
            if response.status_code == 200:
                print("✅ Server is running!")
                return True
        except:
            print(f"⏳ Waiting for server... ({i+1}/{max_wait})")
            time.sleep(1)
    
    print("❌ Server not ready within timeout")
    return False

def test_webui_auth():
    """Test Web UI authentication"""
    print("\n=== TESTING WEB UI AUTHENTICATION ===")
    
    # Test without authentication first
    try:
        response = requests.get(WEBUI_URL)
        print(f"Without auth: {response.status_code}")
        
        if response.status_code == 401:
            print("✅ Server now requires authentication!")
            print("✅ Authentication is working correctly!")
            
            # Check WWW-Authenticate header
            auth_header = response.headers.get('WWW-Authenticate', '')
            print(f"WWW-Authenticate: {auth_header}")
            
        elif response.status_code == 200:
            print("❌ Server still not requiring authentication")
            return False
            
    except Exception as e:
        print(f"Error testing without auth: {e}")
        return False
    
    # Test with correct credentials
    try:
        credentials = base64.b64encode(f"{USERNAME}:{PASSWORD}".encode()).decode()
        headers = {"Authorization": f"Basic {credentials}"}
        response = requests.get(WEBUI_URL, headers=headers)
        
        print(f"\nWith correct credentials: {response.status_code}")
        
        if response.status_code == 200:
            print("✅ Authentication successful with correct credentials!")
            
            # Check if we got the actual Web UI content
            if "Search" in response.text or "LightRAG" in response.text or "RailSeek" in response.text:
                print("✅ Web UI content loaded successfully")
                return True
            else:
                print("❌ Got 200 but no Web UI content")
                print(f"Response preview: {response.text[:500]}")
                return False
        else:
            print(f"❌ Authentication failed with correct credentials: {response.status_code}")
            return False
            
    except Exception as e:
        print(f"Error testing with auth: {e}")
        return False

def test_api_auth():
    """Test API authentication"""
    print("\n=== TESTING API AUTHENTICATION ===")
    
    headers = {"X-API-Key": "jleu1212"}
    
    try:
        response = requests.get(f"{BASE_URL}/documents", headers=headers)
        print(f"API with API Key: {response.status_code}")
        
        if response.status_code == 200:
            print("✅ API authentication working!")
            return True
        else:
            print(f"❌ API authentication failed: {response.status_code}")
            return False
            
    except Exception as e:
        print(f"API test error: {e}")
        return False

def test_wrong_credentials():
    """Test with wrong credentials"""
    print("\n=== TESTING WRONG CREDENTIALS ===")
    
    try:
        wrong_credentials = base64.b64encode(b"wrong:wrong").decode()
        headers = {"Authorization": f"Basic {wrong_credentials}"}
        response = requests.get(WEBUI_URL, headers=headers)
        
        print(f"With wrong credentials: {response.status_code}")
        
        if response.status_code == 401:
            print("✅ Wrong credentials correctly rejected!")
            return True
        else:
            print(f"❌ Wrong credentials not rejected: {response.status_code}")
            return False
            
    except Exception as e:
        print(f"Wrong credentials test error: {e}")
        return False

def main():
    print("Testing Fixed Web UI Authentication")
    print("=" * 60)
    
    # Wait for server to be ready
    if not wait_for_server():
        print("❌ Server not ready, cannot test authentication")
        return
    
    # Test Web UI authentication
    webui_ok = test_webui_auth()
    
    # Test API authentication
    api_ok = test_api_auth()
    
    # Test wrong credentials
    wrong_ok = test_wrong_credentials()
    
    print("\n" + "=" * 60)
    print("FINAL RESULTS")
    print("=" * 60)
    
    if webui_ok and api_ok and wrong_ok:
        print("🎉 ALL AUTHENTICATION TESTS PASSED! 🎉")
        print("\n✅ Web UI authentication is now working correctly")
        print("✅ API authentication is working correctly") 
        print("✅ Wrong credentials are properly rejected")
        
        print(f"\nYou can now access the Web UI at: {WEBUI_URL}")
        print("Use credentials: jleu3482 / jleu1212")
        print("\nThe browser should now properly prompt for authentication")
        
    else:
        print("❌ Some authentication tests failed")
        print(f"Web UI: {'✅' if webui_ok else '❌'}")
        print(f"API: {'✅' if api_ok else '❌'}")
        print(f"Wrong credentials: {'✅' if wrong_ok else '❌'}")
        
        print("\nTroubleshooting steps:")
        print("1. Check server logs for authentication errors")
        print("2. Verify AUTH_ACCOUNTS environment variable is set")
        print("3. Try restarting the server again")
        print("4. Check if there are multiple server instances running")

if __name__ == "__main__":
    main()