import base64
import json
import re
import requests
import os

# مفتاح Gemini API
API_KEY = "AIzaSyBr2-dUqHDZkk20hlWeEcpWnVVdkq9fqyE"

# المجلد الذي يحتوي صور cr1
cr1_images_folder = r"C:\Users\ASUS\OneDrive - Binder\Desktop\test.orch\classified\CR1"
# مجلد إخراج ملفات JSON
output_json_folder = r"C:\Users\ASUS\OneDrive - Binder\Desktop\test.orch\classified\CR1\cr1_json"

# Ensure output folder exists
os.makedirs(output_json_folder, exist_ok=True)

# Exact same prompt
prompt = """
Please extract the following fields in Arabic and English from the tax registration document image:
Taxpayer Name
VAT Registration Number
Effective Registration Date
Taxpayer Address

Return the result in a JSON format with these keys:
en_taxpayer_name, en_vat_registration_number, en_effective_registration_date, en_taxpayer_address,
ar_taxpayer_name, ar_vat_registration_number, ar_effective_registration_date, ar_taxpayer_address

If a field is missing, return null.
"""


url = f"https://generativelanguage.googleapis.com/v1/models/gemini-1.5-flash:generateContent?key={API_KEY}"
headers = {"Content-Type": "application/json"}

# Iterate over all images
for image_name in os.listdir(cr1_images_folder):
    if not image_name.lower().endswith(('.jpg', '.jpeg', '.png')):
        continue

    image_path = os.path.join(cr1_images_folder, image_name)
    base_name = os.path.splitext(image_name)[0]
    output_file = os.path.join(output_json_folder, base_name + ".json")

    # Skip if JSON already exists
    if os.path.exists(output_file):
        print(f"Skipped {image_name} (JSON file already exists)")
        continue

    # Read image and convert to base64
    with open(image_path, "rb") as f:
        image_b64 = base64.b64encode(f.read()).decode()

    # Send request to Gemini API
    data = {
        "contents": [
            {
                "role": "user",
                "parts": [
                    {"text": prompt},
                    {
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": image_b64
                        }
                    }
                ]
            }
        ]
    }

    try:
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()
        response_text = response.json()['candidates'][0]['content']['parts'][0]['text']

        match = re.search(r"```json\s*(\{.*\})\s*```", response_text, re.DOTALL)
        if match:
            json_text = match.group(1)
            result = json.loads(json_text)

            with open(output_file, "w", encoding="utf-8") as f:
                json.dump(result, f, ensure_ascii=False, indent=2)

            print(f"✅ Processed: {image_name}")
        else:
            print(f"❌ Failed to extract JSON from: {image_name}")
            print(response_text)

    except Exception as e:
        print(f"❌ Error processing image {image_name}: {e}")