#!/usr/bin/env python3
"""Test Google Vision batch API"""

from google.cloud import vision
import cv2
import os

# Initialize client
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'google-credentials.json'
client = vision.ImageAnnotatorClient()

# Load a test image
test_image_path = 'googleocroutput/cells/page003_cell01.png'
if not os.path.exists(test_image_path):
    print(f"Test image not found: {test_image_path}")
    exit(1)

image_data = cv2.imread(test_image_path)
_, encoded = cv2.imencode('.png', image_data)
content = encoded.tobytes()

# Test single request first
print("Testing single request...")
image = vision.Image(content=content)
response = client.text_detection(image=image)
if response.text_annotations:
    print(f"Single request works: {response.text_annotations[0].description[:50]}")
else:
    print("Single request: No text detected")

# Test batch request
print("\nTesting batch request...")
requests = []
for i in range(2):
    image = vision.Image(content=content)
    requests.append({
        'image': image,
        'features': [{'type_': vision.Feature.Type.TEXT_DETECTION}]
    })

print(f"Sending batch of {len(requests)} requests...")
response = client.batch_annotate_images(requests=requests)
print(f"Received {len(response.responses)} responses")

for idx, image_response in enumerate(response.responses):
    if image_response.error.message:
        print(f"  Response {idx}: Error - {image_response.error.message}")
    elif image_response.text_annotations:
        print(f"  Response {idx}: {image_response.text_annotations[0].description[:50]}")
    else:
        print(f"  Response {idx}: No text detected")

print("\nBatch API test complete!")