import cv2
import numpy as np
import os
from PIL import Image
from google import genai
from google.genai import types

# 1. הגדרות (שים את המפתח החדש שלך!)
client = genai.Client(api_key="AIzaSyA0odvjKCTS-vZsEjMOlsVsJ4iGaYj5jf0")

def process_and_ocr(image_path):
    print(f"\n--- Processing: {image_path} ---")
    
    # 2. OpenCV - זיהוי שורות וחיתוך
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (7,7), 0)
    thresh = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
    
    # מחברים שורות (קרנל רחב)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 5))
    dilate = cv2.dilate(thresh, kernel, iterations=1)
    contours, _ = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # מיון מלמעלה למטה
    contours = sorted(contours, key=lambda c: cv2.boundingRect(c)[1])
    
    crops = []
    max_w = 0
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        if w > 60 and h > 20:
            crop = img[y:y+h, x:x+w]
            crops.append(crop)
            if w > max_w: max_w = w

    # 3. Stitching - הדבקת החיתוכים לסטריפ אחד (מנקה שטחים לבנים)
    # מוסיפים פדינג שיהיה רוחב אחיד להדבקה
    final_crops = []
    for c in crops:
        h, w, _ = c.shape
        padded = cv2.copyMakeBorder(c, 5, 5, 0, max_w - w, cv2.BORDER_CONSTANT, value=[255, 255, 255])
        final_crops.append(padded)
    
    if not final_crops:
        print("Error: No text blocks found!")
        return

    # הדבקה אנכית
    stitched_strip = cv2.vconcat(final_crops)
    strip_name = f"strip_{image_path}"
    cv2.imwrite(strip_name, stitched_strip)
    print(f"Stitched strip saved as {strip_name}")

    # 4. שליחה ל-Gemini Flash
    pil_img = Image.open(strip_name)
    
    prompt = """
    התמונה המצורפת היא "סטריפ" המורכב מחיתוכים של שאלת בגרות במתמטיקה.
    המשימה שלך: חלץ את כל הטקסט והמשוואות בדיוק מושלם של 100%.
    חוקים:
    1. קרא את העברית והמתמטיקה ברצף מלמעלה למטה.
    2. שים לב לפרטים הקטנים ביותר: שברים, חזקות, e, ln, וסימני מינוס.
    3. החזר LaTeX תקני לכל המתמטיקה ועברית נקייה.
    4. החזר רק את הטקסט המחולץ.
    """

    response = client.models.generate_content(
        model='gemini-2.0-flash',
        contents=[pil_img, prompt],
        config=types.GenerateContentConfig(temperature=0.0, max_output_tokens=4096)
    )
    
    print("\n--- OCR RESULT ---")
    print(response.text)
    print("------------------")

# הרצת הניסוי על 3 הקבצים (שנה את השמות אם צריך)
test_files = ["hardest_1.jpg", "ocr_only.jpeg", "simple.jpg"]
for file in test_files:
    if os.path.exists(file):
        process_and_ocr(file)
    else:
        print(f"File {file} not found, skipping...")