import cv2
import numpy as np
from PIL import Image
import os

def run_final_assembly_test(image_path):
    # --- שלב 1: הזיהוי המנצח שלך (הריבועים הירוקים) ---
    image = cv2.imread(image_path)
    img_h, img_w = image.shape[:2]
    
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    thresh = cv2.adaptiveThreshold(cv2.GaussianBlur(gray, (7,7), 0), 255, 
                                 cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 10))
    dilate = cv2.dilate(thresh, kernel, iterations=1)
    contours, _ = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    blocks = sorted([cv2.boundingRect(c) for c in contours if cv2.boundingRect(c)[2] > 50], key=lambda b: b[1])

    # --- שלב 2: בניית ה-Payload (חיבור הריבועים לסלייסים) ---
    pil_img = Image.open(image_path).convert("RGB")
    slices = []
    
    # Header Lock
    slices.append(pil_img.crop((0, 0, img_w, 180)))
    
    remaining = [b for b in blocks if (b[1] + b[3]) > 180]
    while remaining:
        curr = remaining.pop(0)
        x, y, w, h = curr
        if h > 120: # גרף
            y_s, y_e = max(0, y-30), (img_h if not remaining else min(img_h, y+h+30))
            slices.append(pil_img.crop((0, y_s, img_w, y_e)))
            remaining = [b for b in remaining if b[1] > (y_e - 10)]
            continue
        
        c_min, c_max = y, y + h
        to_consume = []
        for nb in remaining:
            if nb[3] > 120 or (max(c_max, nb[1]+nb[3]) - c_min) > 750: break
            c_max = max(c_max, nb[1]+nb[3])
            to_consume.append(nb)
        for r in to_consume: remaining.remove(r)
        
        y_t, y_b = max(0, c_min-25), (img_h if not remaining else min(img_h, c_max+25))
        slices.append(pil_img.crop((0, y_t, img_w, y_b)))

    # --- שלב 3: ויזואליזציה של החיבור (ה"ניבוי") ---
    # אנחנו מחברים את כל הסלייסים לתמונה אחת ארוכה עם קו מפריד אדום
    assembly_parts = []
    separator = np.zeros((10, img_w, 3), dtype=np.uint8)
    separator[:] = [0, 0, 255] # קו אדום בולט
    
    for i, s in enumerate(slices):
        assembly_parts.append(np.array(s))
        if i < len(slices) - 1:
            assembly_parts.append(separator)
            
    final_assembly = np.vstack(assembly_parts)
    cv2.imwrite("final_assembly_debug.jpg", cv2.cvtColor(final_assembly, cv2.COLOR_RGB2BGR))
    
    print(f"✅ הצלחנו לחבר {len(slices)} חלקים.")
    print(f"📂 התוצאה הסופית מחכה לך ב: final_assembly_debug.jpg")
    print("\n--- מבנה ה-Payload שיישלח ל-LLM ---")
    print(f"Payload = [PROMPT, " + ", ".join([f"PAGE_{i+1}" for i in range(len(slices))]) + "]")

if __name__ == "__main__":
    target = r"C:\Projects\BuddyMath\buddy_math_server\hardest_1.jpg"
    run_final_assembly_test(target)