import requests
import json

# Test the query endpoint for search functionality
base_url = 'http://localhost:3015'

# Login first
login_data = {'username': 'jleu3482', 'password': 'jleu1212'}
login_response = requests.post(f'{base_url}/login', data=login_data)

if login_response.status_code == 200:
    token = login_response.json().get('access_token')
    headers = {'Authorization': f'Bearer {token}'}
    print('✅ Login successful')
    
    # Test the query endpoint
    print('🔍 Testing /query endpoint...')
    query_data = {'query': 'test document', 'top_k': 5}
    try:
        response = requests.post(f'{base_url}/query', json=query_data, headers=headers)
        print(f'  Status: {response.status_code}')
        if response.status_code == 200:
            results = response.json()
            print(f'  ✅ Success! Found results (type: {type(results)})')
            print(f'  Raw response: {results}')
            
            # Handle different response formats
            if isinstance(results, list):
                print(f'  Found {len(results)} results')
                for i, result in enumerate(results):
                    if isinstance(result, dict):
                        content = result.get('content', '')
                        score = result.get('score', 0)
                        print(f'  Result {i+1} (score: {score:.3f}): {content[:150]}...')
                    else:
                        print(f'  Result {i+1}: {str(result)[:150]}...')
            elif isinstance(results, dict):
                print(f'  Dictionary result: {results}')
            else:
                print(f'  String result: {str(results)[:200]}...')
        else:
            print(f'  ❌ Failed: {response.text}')
    except Exception as e:
        print(f'  ❌ Error: {e}')
        
    # Test with different queries
    queries = ['PDF', 'sample', 'OCR', 'processing', 'GPU']
    for query in queries:
        print(f'\n🔎 Querying: "{query}"')
        query_data = {'query': query, 'top_k': 3}
        response = requests.post(f'{base_url}/query', json=query_data, headers=headers)
        if response.status_code == 200:
            results = response.json()
            print(f'  ✅ Found results (type: {type(results)})')
            if isinstance(results, list):
                print(f'  Found {len(results)} results')
                for i, result in enumerate(results):
                    if isinstance(result, dict):
                        content = result.get('content', '')
                        score = result.get('score', 0)
                        print(f'    {i+1}. (score: {score:.3f}) {content[:100]}...')
                    else:
                        print(f'    {i+1}. {str(result)[:100]}...')
            else:
                print(f'  Result: {str(results)[:200]}...')
        else:
            print(f'  ❌ Failed: {response.text}')
else:
    print('❌ Login failed')