import sys
import cv2
import json
import numpy as np
import os

def hex_to_hsv(hex_code):
    rgb_color = np.array([int(hex_code[i:i+2], 16) for i in (1, 3, 5)])
    bgr_color = np.flip(rgb_color)
    hsv_color = cv2.cvtColor(np.uint8([[bgr_color]]), cv2.COLOR_BGR2HSV)[0][0]
    return hsv_color

def draw_and_get_contours(image_path, colors_dict, output_path, min_contour_size=150, max_contour_size=10000):
    source_image = cv2.imread(image_path)

    if source_image is None:
        print(f"Error: Unable to load image at {image_path}")
        sys.exit(1)

    hsv_image = cv2.cvtColor(source_image, cv2.COLOR_BGR2HSV)
    combined_mask = np.zeros_like(hsv_image[:, :, 0])

    for color_hex_code, bounds in colors_dict.items():
        mask = cv2.inRange(hsv_image, bounds['lower'], bounds['upper'])
        combined_mask = cv2.bitwise_or(combined_mask, mask)

    contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    bounding_rectangles = []

    for cnt, contour in enumerate(contours):
        contour_area = cv2.contourArea(contour)
        if min_contour_size <= contour_area <= max_contour_size:
            x, y, w, h = cv2.boundingRect(contour)
            x1, y1 = x, y
            x2, y2 = x + w, y + h
            bounding_rectangles.append((x1, y1, x2, y2))

    merged_rectangles = merge_adjacent_rectangles(bounding_rectangles)

    for rect in merged_rectangles:
        x1, y1, x2, y2 = rect
        cv2.rectangle(source_image, (x1, y1), (x2, y2), (0, 255, 0), 2)

    cv2.imwrite(output_path, source_image)
    json_string = json.dumps(merged_rectangles)
    return json_string

def merge_adjacent_rectangles(rectangles):
    if not rectangles:
        return []

    merged_rectangles = []
    while rectangles:
        rect1 = rectangles.pop(0)
        merged = False

        for i in range(len(merged_rectangles)):
            rect2 = merged_rectangles[i]
            if is_adjacent_or_overlapping(rect1, rect2):
                new_rect = merge_group([rect1, rect2])
                merged_rectangles[i] = new_rect
                merged = True
                break

        if not merged:
            merged_rectangles.append(rect1)

    final_rectangles = []
    while merged_rectangles:
        rect1 = merged_rectangles.pop(0)
        merged = False

        for i in range(len(final_rectangles)):
            rect2 = final_rectangles[i]
            if is_adjacent_or_overlapping(rect1, rect2):
                new_rect = merge_group([rect1, rect2])
                final_rectangles[i] = new_rect
                merged = True
                break

        if not merged:
            final_rectangles.append(rect1)

    return final_rectangles

def is_adjacent_or_overlapping(rect1, rect2):
    x1, y1, x2, y2 = rect1
    x1_, y1_, x2_, y2_ = rect2

    horizontal_adjacent = (x1 <= x2_ + 20 and x2 >= x1_ - 20)
    vertical_adjacent = (y1 <= y2_ + 20 and y2 >= y1_ - 20)

    return horizontal_adjacent and vertical_adjacent

def merge_group(group):
    x1 = min(rect[0] for rect in group)
    y1 = min(rect[1] for rect in group)
    x2 = max(rect[2] for rect in group)
    y2 = max(rect[3] for rect in group)
    return (x1, y1, x2, y2)

if __name__ == "__main__":
    colors = {
        '#FFFBDB': {'lower': np.array([20, 10, 200]), 'upper': np.array([30, 100, 255])},  # Light Yellow
        '#AFCB04': {'lower': np.array([24, 200, 150]), 'upper': np.array([44, 255, 255])}  # Lime Green
    }

    output_dir = 'D:/crud-operations/public/storage/master_plans'
    output_contour_path = os.path.join(output_dir, 'floor_plan_with_bbox.png')

    if len(sys.argv) != 2:
        print("Usage: python script.py <image_file>")
        sys.exit(1)

    image_file = sys.argv[1]
    json_output = draw_and_get_contours(image_file, colors, output_contour_path, min_contour_size=150, max_contour_size=10000)
    print(json_output)