#!/opt/homebrew/bin/python3.11
"""
Batch Scan Keyframes for SMALL red stamps
"""

import cv2
import numpy as np
import os

UUID = "384b0ff44aaaa1f1"
BASE_DIR = f"output/{UUID}/florence2_results"

# The keyframes we extracted
FRAMES = [
    "scan_6751.jpg",  # 112:31
    "scan_6755.jpg",  # 112:35
    "scan_6759.jpg",  # 112:39
]


def find_small_stamps_in_frame(img_path):
    if not os.path.exists(img_path):
        return []

    img = cv2.imread(img_path)
    if img is None:
        return []

    h, w, _ = img.shape
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # Red Mask
    mask1 = cv2.inRange(hsv, np.array([0, 70, 50]), np.array([10, 255, 255]))
    mask2 = cv2.inRange(hsv, np.array([170, 70, 50]), np.array([180, 255, 255]))
    mask = mask1 + mask2

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    found = []
    # Constraint: Stamp must be small.
    # If the image is 1080p, a stamp held in hand is likely < 200x200 pixels?
    # Let's say Area < 30,000 pixels (approx 5% of width * 5% of height).
    MAX_STAMP_AREA = 30000

    for cnt in contours:
        area = cv2.contourArea(cnt)

        if area > MAX_STAMP_AREA:
            continue

        # Shape Constraint: Triangle or Rectangle
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.04 * peri, True)
        vertices = len(approx)

        # Check for Triangle (3) or Rectangle (4)
        if vertices in [3, 4]:
            x, y, w_box, h_box = cv2.boundingRect(approx)

            # Aspect ratio check? Stamps are usually rectangular or square.
            # Triangles are rare for stamps unless folded.
            # Let's just return everything small and triangular/rectangular for now.

            found.append((x, y, w_box, h_box, area, approx))

    return found, img


print("🔍 Scanning Keyframes for Small Stamps...")

for frame_name in FRAMES:
    frame_path = os.path.join(BASE_DIR, frame_name)
    result, img = find_small_stamps_in_frame(frame_path)

    print(f"\n🎞️ Analyzing {frame_name}...")

    if result:
        print(f"   ✅ Found {len(result)} potential small stamps!")

        # Draw them
        for x, y, w_box, h_box, area, approx in result:
            print(f"      📍 Box: ({x},{y}), Size: {w_box}x{h_box}, Area: {area}")
            cv2.rectangle(img, (x, y), (x + w_box, y + h_box), (0, 255, 0), 2)

            # Save crop
            crop_name = f"crop_{frame_name}_{x}_{y}.jpg"
            crop_path = os.path.join(BASE_DIR, crop_name)
            crop = img[y : y + h_box, x : x + w_box]
            cv2.imwrite(crop_path, crop)

        # Save result image
        res_name = f"result_opencv_{frame_name}"
        cv2.imwrite(os.path.join(BASE_DIR, res_name), img)
    else:
        print("   ❌ No small stamps found.")
