#!/usr/bin/env python3
import requests
import base64
import json
import sys
import os

# Clarifai config
API_KEY = "a1a09aad41aa43abbe23122fa1996544"
BASE_URL = "https://api.clarifai.com/v2/ext/openai/v1"
MODEL_ID = "https://clarifai.com/deepseek-ai/deepseek-ocr/models/DeepSeek-OCR/versions/c52cf7da9b1c4095b07e2a1ccb842811"

def encode_image(image_path):
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def ocr_image(image_path):
    b64 = encode_image(image_path)
    
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    # OpenAI-compatible payload
    payload = {
        "model": MODEL_ID,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Extract all text from this image. Keep original formatting, including line breaks. Preserve all special characters, diacritics, and Pali/Sanskrit letters exactly as they appear."
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{b64}"
                        }
                    }
                ]
            }
        ],
        "max_tokens": 4096
    }
    
    response = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload)
    if response.status_code != 200:
        print(f"Error: {response.status_code}")
        print(response.text)
        return None
    
    result = response.json()
    text = result["choices"][0]["message"]["content"]
    return text

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 ocr_test.py <image_path>")
        sys.exit(1)
    
    image_path = sys.argv[1]
    if not os.path.exists(image_path):
        print(f"File not found: {image_path}")
        sys.exit(1)
    
    print(f"Processing {image_path}...")
    text = ocr_image(image_path)
    if text:
        print("=== OCR Result ===")
        print(text)
    else:
        print("OCR failed.")