import fitz  # PyMuPDF
import argparse
import json  # Import json to return the output in JSON format

# Detected colors from the PDF: Pink and Orange
target_colors_rgb_normalized = [
    (1.0, 0.0, 0.498),         # Pink (#FF007E)
    (0.949, 0.404, 0.133)      # Orange (#F16721)
]

# Tolerance for color matching
def colors_are_similar(color1, color2, tolerance=0.05):
    return all(abs(a - b) <= tolerance for a, b in zip(color1, color2))

# Argument parsing to accept PDF path from command line
parser = argparse.ArgumentParser(description="Extract polygons from a PDF.")
parser.add_argument("pdf_path", help="Path to the PDF file")
args = parser.parse_args()

# Open the PDF file
doc = fitz.open(args.pdf_path)

# Initialize a dictionary to group polygons by color
polygon_groups = {}

# Iterate through all pages in the PDF
for page_num in range(doc.page_count):
    page = doc.load_page(page_num)
    
    # Extract all drawing commands on the page
    for item in page.get_drawings():
        # Get the bounding box of the drawing item
        bbox = item.get('rect', None)

        # Check for both 'fill' and 'stroke' colors
        fill_color = item.get('fill')
        stroke_color = item.get('stroke')

        # Check if the fill color matches pink or orange (with tolerance)
        if fill_color and any(colors_are_similar(tuple(fill_color[:3]), target_color) for target_color in target_colors_rgb_normalized):
            if bbox:
                color_key = tuple(fill_color[:3])
                if color_key not in polygon_groups:
                    polygon_groups[color_key] = [bbox]
                else:
                    polygon_groups[color_key].append(bbox)

        # Check if the stroke color matches pink or orange (with tolerance)
        if stroke_color and any(colors_are_similar(tuple(stroke_color[:3]), target_color) for target_color in target_colors_rgb_normalized):
            if bbox:
                color_key = tuple(stroke_color[:3])
                if color_key not in polygon_groups:
                    polygon_groups[color_key] = [bbox]
                else:
                    polygon_groups[color_key].append(bbox)

# Close the document after processing
doc.close()

# Function to calculate the combined bounding box for a group of polygons
def calculate_combined_bbox(bboxes):
    if not bboxes:
        return None
    # Start with the first bounding box
    x0, y0, x1, y1 = bboxes[0].x0, bboxes[0].y0, bboxes[0].x1, bboxes[0].y1
    # Iterate over all bounding boxes and expand the combined box
    for bbox in bboxes[1:]:
        x0 = min(x0, bbox.x0)
        y0 = min(y0, bbox.y0)
        x1 = max(x1, bbox.x1)
        y1 = max(y1, bbox.y1)
    return (x0, y0, x1, y1)

# Output the combined bounding box for each color group
combined_bboxes = {}
if polygon_groups:
    for color, bboxes in polygon_groups.items():
        combined_bbox = calculate_combined_bbox(bboxes)
        if combined_bbox:
            color_hex = '#FF007E' if colors_are_similar(color, (1.0, 0.0, 0.498)) else '#F16721'  # Set the color name
            combined_bboxes[color_hex] = combined_bbox

# Convert combined_bboxes to JSON format and print the result
print(json.dumps(combined_bboxes))