import requests
import json
import time
import sys

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

headers = {
    'X-API-Key': api_key,
    'Content-Type': 'application/json',
    'X-Workspace': workspace
}

def upload_document():
    """Upload a test document to workspace test1"""
    files = {'file': open('test/tir.docx', 'rb')}
    upload_url = f'{base_url}/documents/upload'
    headers = {'X-API-Key': api_key, 'X-Workspace': workspace}
    data = {'workspace': workspace, 'chunk_size': 1200, 'chunk_overlap': 100}
    resp = requests.post(upload_url, headers=headers, files=files, data=data)
    print(f'Upload status: {resp.status_code}')
    if resp.status_code == 200:
        print('Upload successful')
        # Response may be a list of document IDs
        result = resp.json()
        print(f'Upload response: {result}')
        # Assume first document ID
        if isinstance(result, list) and len(result) > 0:
            doc_id = result[0].get('document_id')
            return {'document_id': doc_id}
        elif isinstance(result, dict) and 'document_id' in result:
            return result
        else:
            # fallback
            return {'document_id': 'unknown'}
    else:
        print(f'Upload error: {resp.text}')
        return None

def wait_for_indexing(doc_id):
    """Poll document status until processed"""
    if doc_id == 'unknown':
        # Cannot poll, just wait a bit
        print('Cannot poll status, waiting 10 seconds...')
        time.sleep(10)
        return True
    status_url = f'{base_url}/documents/{doc_id}/status'
    for _ in range(30):
        resp = requests.get(status_url, headers={'X-API-Key': api_key, 'X-Workspace': workspace})
        if resp.status_code == 200:
            data = resp.json()
            status = data.get('status')
            print(f'Document status: {status}')
            if status == 'processed':
                return True
            elif status == 'error':
                print(f'Document processing error: {data}')
                return False
        time.sleep(2)
    print('Timeout waiting for indexing')
    return False

def test_search():
    """Test search with keyword extraction"""
    search_url = f'{base_url}/api/search'
    payload = {'query': 'minimum safe working distance', 'top_k': 5}
    resp = requests.post(search_url, headers=headers, json=payload, timeout=30)
    print(f'Search status: {resp.status_code}')
    if resp.status_code == 200:
        data = resp.json()
        print(f'Search successful. Found {len(data.get("chunks", []))} chunks')
        print(f'Entities: {len(data.get("entities", []))}')
        return True
    else:
        print(f'Search error: {resp.text}')
        return False

def main():
    print('Testing LLM region restriction...')
    # Step 1: Upload document
    print('Uploading document...')
    result = upload_document()
    if not result:
        sys.exit(1)
    doc_id = result.get('document_id')
    print(f'Document ID: {doc_id}')
    
    # Step 2: Wait for indexing
    print('Waiting for indexing...')
    if not wait_for_indexing(doc_id):
        sys.exit(1)
    
    # Step 3: Test search
    print('Testing search with keyword extraction...')
    if test_search():
        print('✅ LLM region restriction passed!')
    else:
        print('❌ LLM region restriction may still be an issue')
        sys.exit(1)

if __name__ == '__main__':
    main()