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

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"
    }
    
    payload = {
        "model": MODEL_ID,
        "messages": [
            {
                "role": "system",
                "content": "You are an OCR text extractor. Extract all text from the image exactly as it appears, preserving all characters, diacritics, punctuation, and line breaks. Do not add any commentary, explanations, or formatting instructions. Output only the extracted text."
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Extract all text from this image."
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{b64}"
                        }
                    }
                ]
            }
        ],
        "max_tokens": 4096,
        "temperature": 0.0
    }
    
    response = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload)
    if response.status_code != 200:
        print(f"Error: {response.status_code}", file=sys.stderr)
        print(response.text, file=sys.stderr)
        return None
    
    result = response.json()
    text = result["choices"][0]["message"]["content"]
    return text.strip()

def main():
    if len(sys.argv) < 2:
        print("Usage: python3 ocr_pages.py <start> <end>")
        sys.exit(1)
    
    start = int(sys.argv[1])
    end = int(sys.argv[2])
    
    output_lines = []
    for i in range(start, end + 1):
        img_name = f"{i:03d}.jpg"
        if not os.path.exists(img_name):
            print(f"Image {img_name} not found", file=sys.stderr)
            continue
        
        print(f"Processing {img_name}...", file=sys.stderr)
        text = ocr_image(img_name)
        if text:
            output_lines.append(f"=== Page {i} ===")
            output_lines.append(text)
            output_lines.append("")
        else:
            output_lines.append(f"=== Page {i} ===")
            output_lines.append("[OCR failed]")
            output_lines.append("")
        
        time.sleep(0.5)  # rate limit
    
    output_text = "\n".join(output_lines)
    print(output_text)

if __name__ == "__main__":
    main()