import cv2
import json
import os
import argparse

def draw_boxes_from_json(image_path: str, json_path: str, output_path: str):
    # 1. Load the image
    image_bgr = cv2.imread(image_path)
    if image_bgr is None:
        print(f"❌ Error: Cannot load image at {image_path}")
        return

    ih, iw = image_bgr.shape[:2]

    # 2. Load the JSON data
    if not os.path.exists(json_path):
        print(f"❌ Error: JSON file not found at {json_path}")
        return

    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # Color map for different region types (BGR format)
    COLOR_MAP = {
        "dialogue":  (0,   200,   0),   # Green
        "narration": (0,   165, 255),   # Orange
        "reaction":  (255, 200,   0),   # Cyan/Blue
        "sfx":       (0,     0, 220),   # Red
        "unknown":   (120, 120, 120),   # Gray
    }

    # 3. Iterate through the JSON and draw boxes
    # Sort by order to keep numbering consistent
    sorted_items = sorted(data.values(), key=lambda x: x.get("order", 0))

    for item in sorted_items:
        bid = item.get("order", "?")
        rtype = item.get("region_type", "unknown")
        box = item.get("box", {})
        text = item.get("corrected_ocr", "")

        if not box:
            continue

        # Extract xywh and convert to xyxy
        x1, y1 = int(box.get("x", 0)), int(box.get("y", 0))
        w, h = int(box.get("w", 0)), int(box.get("h", 0))
        x2, y2 = x1 + w, y1 + h

        color = COLOR_MAP.get(rtype, (120, 120, 120))

        # Draw the main bounding box
        cv2.rectangle(image_bgr, (x1, y1), (x2, y2), color, 2)

        # Prepare labels
        label = f"BOX#{bid} [{rtype}]"
        preview = (text[:40] + "...") if len(text) > 40 else text
        
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.38
        thickness = 1

        # Draw Label Background
        (lw, lh), _ = cv2.getTextSize(label, font, font_scale, thickness)
        cv2.rectangle(image_bgr,
                      (x1, max(0, y1 - lh - 6)),
                      (x1 + lw + 4, y1),
                      color, -1)
        
        # Draw Label Text (Box ID + Type)
        cv2.putText(image_bgr, label,
                    (x1 + 2, max(lh, y1 - 3)),
                    font, font_scale, (255, 255, 255), thickness,
                    cv2.LINE_AA)

        # Draw Preview Text below the box
        cv2.putText(image_bgr, preview,
                    (x1 + 2, min(ih - 5, y2 + 12)),
                    font, font_scale * 0.85, color, thickness,
                    cv2.LINE_AA)

    # 4. Save the final image
    cv2.imwrite(output_path, image_bgr)
    print(f"✅ Debug image successfully saved to: {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Draw bounding boxes from bubbles.json onto an image.")
    parser.add_argument("image", help="Path to the original manga page image")
    parser.add_argument("json", help="Path to the bubbles.json file")
    parser.add_argument("--output", "-o", default="debug_clusters_from_json.png", help="Output image path")
    
    args = parser.parse_args()
    
    draw_boxes_from_json(args.image, args.json, args.output)