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

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_contours(image_path, colors_dict, output_dir, 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)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    combined_mask_output_path = os.path.join(output_dir, f'combined_mask_{timestamp}.png')
    cv2.imwrite(combined_mask_output_path, combined_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)
            bounding_rectangles.append([x, y, x + w, y + h])
            cv2.rectangle(source_image, (x, y), (x + w, y + h), (0, 255, 0), 2)

    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, f'floor_plan_with_bbox_{timestamp}.png')
    cv2.imwrite(output_path, source_image)

    result = {
        'coordinates': bounding_rectangles,
        'output_image_path': output_path,
        'combined_mask_path': combined_mask_output_path
    }

    print(json.dumps(result))

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

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

    image_file = sys.argv[1]
    output_dir = sys.argv[2]
    draw_contours(image_file, colors, output_dir, min_contour_size=150, max_contour_size=10000)