import requests
import json

base_url = 'http://localhost:3015'
api_key = 'jleu1212'

def test_search_with_workspace(workspace, query):
    headers = {'X-API-Key': api_key, 'X-Workspace': workspace, 'Content-Type': 'application/json'}
    payload = {'query': query, 'top_k': 5}
    resp = requests.post(f'{base_url}/api/search', headers=headers, json=payload, timeout=30)
    print(f'Workspace: {workspace}, Query: {query}, Status: {resp.status_code}')
    if resp.status_code == 200:
        data = resp.json()
        print(f'  Entities: {len(data.get("entities", []))}, Chunks: {len(data.get("chunks", []))}')
        if data.get('chunks'):
            print(f'  First chunk preview: {data["chunks"][0].get("text", "")[:100]}...')
        return True
    else:
        print(f'  Error: {resp.text}')
        return False

def test_query_with_workspace(workspace, query):
    headers = {'X-API-Key': api_key, 'X-Workspace': workspace, 'Content-Type': 'application/json'}
    payload = {'query': query, 'top_k': 5}
    resp = requests.post(f'{base_url}/query/stream', headers=headers, json=payload, timeout=30)
    print(f'Workspace: {workspace}, Query: {query}, Status: {resp.status_code}')
    if resp.status_code == 200:
        # stream response lines
        lines = resp.text.strip().split('\n')
        for line in lines:
            if line:
                try:
                    obj = json.loads(line)
                    if 'response' in obj:
                        print(f'  Response: {obj["response"][:100]}...')
                        break
                except:
                    pass
        return True
    else:
        print(f'  Error: {resp.text}')
        return False

def main():
    print('Testing workspace isolation...')
    # Workspace test1 has tir.docx (contains "minimum safe working distance"?)
    # Workspace test2 may have other documents (maybe ocr.pdf)
    # Workspace '' (default) may have some data.
    
    # Search for "minimum safe working distance" with workspace test1
    print('\n--- Search with workspace test1 ---')
    test_search_with_workspace('test1', 'minimum safe working distance')
    
    # Search for "ODDS" with workspace test1
    print('\n--- Search with workspace test1 for ODDS ---')
    test_search_with_workspace('test1', 'ODDS')
    
    # Search for "minimum safe working distance" with workspace '' (default)
    print('\n--- Search with workspace "" (default) ---')
    test_search_with_workspace('', 'minimum safe working distance')
    
    # Search for "ODDS" with workspace ''
    print('\n--- Search with workspace "" for ODDS ---')
    test_search_with_workspace('', 'ODDS')
    
    # Query stream test
    print('\n--- Query stream with workspace test1 ---')
    test_query_with_workspace('test1', 'what is the minimum safe working distance')
    
    print('\n--- Query stream with workspace "" ---')
    test_query_with_workspace('', 'what is the minimum safe working distance')

if __name__ == '__main__':
    main()