import sys
import io
import argparse
import requests
from PIL import Image, ImageDraw, ImageFont

def generate_test_image() -> bytes:
    """Generates a simple test image containing clear text in memory."""
    print("Generating a test image in memory...")
    # Create a 400x100 white image
    img = Image.new("RGB", (400, 100), color="white")
    d = ImageDraw.Draw(img)
    
    # Write a simple English sentence on it
    d.text((20, 40), "HELLO WORLD", fill="black")
    d.text((20, 60), "PADDLEOCR TEST", fill="black")
    
    # Save to a byte buffer as PNG
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, format="PNG")
    return img_byte_arr.getvalue()

def run_test(api_url: str, api_key: str, image_path: str = None, lang: str = "en"):
    """Sends an image to the PaddleOCR API and prints the response."""
    # Build endpoint URL (ensure correct trailing slash structure)
    endpoint = api_url.rstrip("/")
    if not endpoint.endswith("/api/v1/image/extract-text"):
        endpoint = f"{endpoint}/api/v1/image/extract-text"
        
    print(f"Target Endpoint: {endpoint}")
    
    # Get image bytes
    if image_path:
        print(f"Reading image from path: {image_path}")
        with open(image_path, "rb") as f:
            image_bytes = f.read()
            filename = image_path
    else:
        # Dynamically generate test image if no path is provided
        image_bytes = generate_test_image()
        filename = "test_generated_image.png"

    # Prepare multipart files and headers
    files = {
        "image": (filename, image_bytes, "image/png")
    }
    data = {
        "lang": lang
    }
    headers = {
        "x-api-key": api_key,
        "Accept": "application/json"
    }

    try:
        print("Sending POST request to API...")
        response = requests.post(endpoint, files=files, data=data, headers=headers, timeout=30)
        
        print(f"HTTP Status Code: {response.status_code}")
        response_json = response.json()
        
        print("\n--- API Response ---")
        import json
        print(json.dumps(response_json, indent=4))
        print("--------------------\n")
        
        if response.status_code == 200 and response_json.get("success"):
            print("SUCCESS: OCR text extracted successfully!")
            print("\n--- Plain Text ---")
            print(response_json["data"]["result"])
            print("\n--- Structured Markdown ---")
            print(response_json["data"].get("markdown", "(no markdown field — redeploy the service)"))
        else:
            print(f"FAILED: API returned an error: {response_json.get('error', 'Unknown error')}")
            
    except requests.exceptions.RequestException as re:
        print(f"ERROR: Connection request failed: {re}")
    except Exception as e:
        print(f"ERROR: An unexpected error occurred: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test script for PaddleOCR API server")
    parser.add_argument("--url", default="https://ashu-choudhury-ocr-backend.hf.space", help="Base URL of your Hugging Face Space (e.g. https://username-space.hf.space)")
    parser.add_argument("--key", default="BLINDbtcTECh2221", help="API key (defaults to standard project key)")
    parser.add_argument("--image", default=None, help="Optional path to a local image file. If omitted, a test image is generated.")
    parser.add_argument("--lang", default="en", help="Language code parameter (default: en)")
    
    args = parser.parse_args()
    
    run_test(
        api_url=args.url,
        api_key=args.key,
        image_path=args.image,
        lang=args.lang
    )