import requests
import time
import sys

SERVER_URL = "http://localhost:3015"
UPLOAD_ENDPOINT = f"{SERVER_URL}/documents/upload"
SEARCH_ENDPOINT = f"{SERVER_URL}/search"
STATUS_ENDPOINT = f"{SERVER_URL}/documents/status"
API_KEY = "jleu1212"

def get_headers(workspace=None):
    headers = {"X-API-Key": API_KEY}
    if workspace:
        headers["X-Workspace"] = workspace
    return headers

def test_upload(file_path, workspace=None):
    """Test upload with optional workspace header"""
    headers = get_headers(workspace)
    with open(file_path, 'rb') as f:
        files = {'file': (file_path.name, f)}
        resp = requests.post(UPLOAD_ENDPOINT, files=files, headers=headers)
    print(f"Upload response: {resp.status_code} {resp.text}")
    return resp.json() if resp.status_code == 200 else None

def test_search(query, workspace=None):
    headers = get_headers(workspace)
    resp = requests.post(SEARCH_ENDPOINT, json={"query": query}, headers=headers)
    print(f"Search response: {resp.status_code} {resp.text}")
    return resp.json() if resp.status_code == 200 else None

def test_status(workspace=None):
    headers = get_headers(workspace)
    resp = requests.get(STATUS_ENDPOINT, headers=headers)
    print(f"Status response: {resp.status_code} {resp.text}")
    return resp.json() if resp.status_code == 200 else None

if __name__ == "__main__":
    from pathlib import Path
    test_file = Path("test/test.docx")
    if not test_file.exists():
        print("Test file not found")
        sys.exit(1)
    
    # Test upload to workspace "test1"
    print("=== Upload to workspace test1 ===")
    result = test_upload(test_file, workspace="test1")
    if result:
        track_id = result.get("track_id")
        print(f"Track ID: {track_id}")
        # Wait a bit for processing
        time.sleep(5)
        # Check status
        print("=== Status for workspace test1 ===")
        test_status("test1")
        # Search
        print("=== Search in workspace test1 ===")
        test_search("test", workspace="test1")
    
    # Test upload to workspace "test2"
    print("\n=== Upload to workspace test2 ===")
    result = test_upload(test_file, workspace="test2")
    if result:
        track_id = result.get("track_id")
        print(f"Track ID: {track_id}")
        time.sleep(5)
        print("=== Status for workspace test2 ===")
        test_status("test2")
        print("=== Search in workspace test2 ===")
        test_search("test", workspace="test2")
    
    # Test search without workspace (should default to default?)
    print("\n=== Search without workspace (should be default?) ===")
    test_search("test")