import requests
import json
import time

BASE_URL = "http://localhost:3015"
API_KEY = "jleu1212"

def make_headers(workspace=None):
    headers = {
        "X-API-Key": API_KEY,
        "Content-Type": "application/json",
    }
    if workspace is not None:
        headers["X-Workspace"] = workspace
    return headers

def upload_text(workspace, text, file_source):
    url = f"{BASE_URL}/documents/text"
    payload = {"text": text, "file_source": file_source}
    response = requests.post(url, headers=make_headers(workspace), json=payload)
    if response.status_code != 200:
        print(f"Upload failed: {response.text}")
        return None
    return response.json().get("track_id")

def get_document_count(workspace):
    url = f"{BASE_URL}/documents/paginated"
    payload = {"page": 1, "page_size": 10}
    response = requests.post(url, headers=make_headers(workspace), json=payload)
    if response.status_code != 200:
        print(f"Failed to get documents: {response.text}")
        return None
    data = response.json()
    return data["pagination"]["total_count"]

def clear_documents(workspace):
    url = f"{BASE_URL}/documents"
    response = requests.delete(url, headers=make_headers(workspace))
    if response.status_code != 200:
        print(f"Clear failed: {response.text}")
        return False
    return True

def delete_workspace(workspace):
    url = f"{BASE_URL}/workspaces/{workspace}"
    response = requests.delete(url, headers=make_headers(workspace))
    if response.status_code != 200:
        print(f"Delete workspace failed: {response.text}")
        return False
    return True

def main():
    workspace_a = "workspace_a"
    workspace_b = "workspace_b"
    
    print("Step 1: Upload a document to workspace A")
    track_a = upload_text(workspace_a, "Document for workspace A", "doc_a.txt")
    if track_a:
        print(f"  Track ID: {track_a}")
    else:
        print("  Upload failed (likely embedding error), but doc_status may still exist")
    
    print("Step 2: Upload a document to workspace B")
    track_b = upload_text(workspace_b, "Document for workspace B", "doc_b.txt")
    if track_b:
        print(f"  Track ID: {track_b}")
    
    # Wait a bit for processing (though it may fail)
    time.sleep(3)
    
    print("Step 3: Get document counts before clear")
    count_a = get_document_count(workspace_a)
    count_b = get_document_count(workspace_b)
    print(f"  Workspace A documents: {count_a}")
    print(f"  Workspace B documents: {count_b}")
    
    print("Step 4: Clear workspace A only")
    if clear_documents(workspace_a):
        print("  Workspace A cleared successfully")
    
    print("Step 5: Get document counts after clear")
    count_a2 = get_document_count(workspace_a)
    count_b2 = get_document_count(workspace_b)
    print(f"  Workspace A documents: {count_a2}")
    print(f"  Workspace B documents: {count_b2}")
    
    # Verify workspace A is empty, workspace B unchanged
    if count_a2 == 0 and count_b2 == count_b:
        print("SUCCESS: Clear button only cleared selected workspace.")
    else:
        print("FAILURE: Clear button affected other workspace.")
        print(f"  Expected workspace A 0, got {count_a2}")
        print(f"  Expected workspace B {count_b}, got {count_b2}")
    
    # Clean up
    print("Step 6: Delete workspaces")
    delete_workspace(workspace_a)
    delete_workspace(workspace_b)
    print("Test completed.")

if __name__ == "__main__":
    main()