#!/usr/bin/env python3
"""
OCR với Google Cloud Vision API (nhanh hơn Gemini cho OCR)
"""

import base64
import os
import json
import time
from PIL import Image
import io
import urllib.request

# Service account path
SA_PATH = "/opt/openclaw/.openclaw/workspace/google_service_account.json"
WORKDIR = "/opt/openclaw/.openclaw/workspace/Chuan Muc Sadi"

print("📖 OCR với Google Cloud Vision API...", flush=True)
start_time = time.time()

# Đọc service account
with open(SA_PATH) as f:
    sa = json.load(f)

# Tạo JWT để lấy access token
import jwt
import time as time_mod

now = int(time_mod.time())
payload = {
    "iss": sa["client_email"],
    "scope": "https://www.googleapis.com/auth/cloud-platform",
    "aud": "https://oauth2.googleapis.com/token",
    "exp": now + 3600,
    "iat": now
}

token = jwt.encode(payload, sa["private_key"], algorithm="RS256")

# Lấy access token
req = urllib.request.Request(
    "https://oauth2.googleapis.com/token",
    data=urllib.parse.urlencode({
        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
        "assertion": token
    }).encode(),
    headers={"Content-Type": "application/x-www-form-urlencoded"}
)

with urllib.request.urlopen(req, timeout=30) as resp:
    token_resp = json.loads(resp.read())
    access_token = token_resp["access_token"]

print(f"   ✅ Got access token", flush=True)

# OCR từng trang
output_lines = ["# Sadi-176-180 (Raw OCR - Google Cloud Vision)", ""]

for i in range(1, 6):
    print(f"\n🔄 Trang {i}/5...", flush=True)
    
    img_path = f"{WORKDIR}/sadi_page-{i}.png"
    
    with open(img_path, "rb") as f:
        img_data = base64.b64encode(f.read()).decode("utf-8")
    
    # Vision API request
    vision_req = {
        "requests": [{
            "image": {"content": img_data},
            "features": [{"type": "TEXT_DETECTION"}]
        }]
    }
    
    req = urllib.request.Request(
        f"https://vision.googleapis.com/v1/images:annotate?key={sa['client_email']}",
        data=json.dumps(vision_req).encode(),
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {access_token}"
        },
        method="POST"
    )
    
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            result = json.loads(resp.read())
            if "responses" in result and len(result["responses"]) > 0:
                text = result["responses"][0].get("fullTextAnnotation", {}).get("text", "[NO TEXT]")
                output_lines.append(f"--- Trang {i} ---")
                output_lines.append(text)
                output_lines.append("")
                print(f"   ✅ {len(text)} chars", flush=True)
            else:
                print(f"   ❌ No response", flush=True)
    except Exception as e:
        print(f"   ❌ {e}", flush=True)
    
    time.sleep(1)

output_path = f"{WORKDIR}/extracted/Sadi-176-180-Vision.md"
with open(output_path, "w", encoding="utf-8") as f:
    f.write("\n".join(output_lines))

print(f"\n✅ Done! {output_path}", flush=True)