#!/usr/bin/env python3
"""
Split a PNG file into multiple SVG files based on brightness/color clusters.
Useful for multi-color 3D printing when the source is a rasterized logo.
"""

import argparse
import numpy as np
from PIL import Image
from pathlib import Path
import cv2
import xml.etree.ElementTree as ET


def extract_color_regions(image_path: Path, dark_threshold: int = 30, mid_min: int = 108, mid_max: int = 118):
    """
    Extract different color regions from PNG based on brightness.
    Uses tighter thresholds to avoid antialiasing artifacts.
    Returns dict of {color_name: binary_mask}
    """
    img = Image.open(image_path).convert('RGBA')
    img_array = np.array(img)

    # Get alpha mask (non-transparent pixels)
    alpha = img_array[..., 3] > 128

    # Calculate brightness for each pixel
    rgb = img_array[..., :3]
    brightness = np.mean(rgb, axis=-1)

    # Create masks for different brightness levels
    masks = {}

    # Dark pixels (black text) - very tight threshold to avoid antialiasing
    dark_mask = alpha & (brightness < dark_threshold)
    if np.sum(dark_mask) > 100:
        masks['black'] = dark_mask

    # Mid-range pixels (grey icon) - tight range to get only the core grey
    mid_mask = alpha & (brightness >= mid_min) & (brightness <= mid_max)
    if np.sum(mid_mask) > 100:
        masks['grey'] = mid_mask

    return masks


def compute_contour_area(contour):
    """Compute the signed area of a contour (positive = clockwise, negative = counter-clockwise)"""
    if len(contour) < 3:
        return 0
    # Shoelace formula
    x = contour[:, 0]
    y = contour[:, 1]
    return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))


def mask_to_svg_paths_with_holes(mask: np.ndarray, simplify_epsilon: float = 1.0):
    """
    Convert a binary mask to SVG path data with proper hole handling.
    Returns list of path d attributes with fill-rule evenodd.
    """
    # Convert to uint8 for OpenCV
    mask_uint8 = (mask * 255).astype(np.uint8)

    # Find contours with hierarchy (to detect holes)
    contours, hierarchy = cv2.findContours(mask_uint8, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

    if hierarchy is None or len(contours) == 0:
        return []

    # Simplify contours
    simplified_contours = []
    for contour in contours:
        if simplify_epsilon > 0:
            simplified = cv2.approxPolyDP(contour, simplify_epsilon, True)
        else:
            simplified = contour
        simplified_contours.append(simplified)

    # Group contours by parent/child relationship
    # hierarchy format: [Next, Previous, First_Child, Parent]
    hierarchy = hierarchy[0]

    # Find all top-level contours (no parent)
    top_level_indices = [i for i in range(len(hierarchy)) if hierarchy[i][3] == -1]

    paths = []

    for top_idx in top_level_indices:
        contour = simplified_contours[top_idx]

        if len(contour) < 3:
            continue

        # Start with outer contour
        path_data = contour_to_svg_path(contour)

        # Find all children (holes) of this contour
        child_idx = hierarchy[top_idx][2]
        while child_idx != -1:
            child_contour = simplified_contours[child_idx]
            if len(child_contour) >= 3:
                # Add hole to the same path (evenodd fill-rule will handle it)
                path_data += " " + contour_to_svg_path(child_contour)

            # Move to next sibling
            child_idx = hierarchy[child_idx][0]

        paths.append(path_data)

    return paths


def contour_to_svg_path(contour):
    """Convert OpenCV contour to SVG path data"""
    # OpenCV contours are shape (N, 1, 2)
    points = contour.reshape(-1, 2)

    if len(points) < 2:
        return ""

    path_data = f"M {points[0, 0]},{points[0, 1]}"
    for point in points[1:]:
        path_data += f" L {point[0]},{point[1]}"
    path_data += " Z"

    return path_data


def get_mask_bounds(mask: np.ndarray):
    """Get the bounding box of a binary mask"""
    rows = np.any(mask, axis=1)
    cols = np.any(mask, axis=0)

    if not np.any(rows) or not np.any(cols):
        return None

    y_min, y_max = np.where(rows)[0][[0, -1]]
    x_min, x_max = np.where(cols)[0][[0, -1]]

    return (x_min, y_min, x_max + 1, y_max + 1)


def create_svg_from_mask(mask: np.ndarray, output_path: Path, color: str, image_size: tuple, crop_to_bounds: bool = True):
    """Create SVG file from binary mask with proper hole handling and optional cropping"""
    width, height = image_size

    # Get bounding box of the mask
    if crop_to_bounds:
        bounds = get_mask_bounds(mask)
        if bounds is None:
            print(f"WARNING: No content found for {color}")
            return

        x_min, y_min, x_max, y_max = bounds
        cropped_width = x_max - x_min
        cropped_height = y_max - y_min

        # Crop the mask
        cropped_mask = mask[y_min:y_max, x_min:x_max]
    else:
        x_min, y_min = 0, 0
        cropped_width, cropped_height = width, height
        cropped_mask = mask

    # Create SVG root with cropped dimensions
    svg = ET.Element('svg', {
        'xmlns': 'http://www.w3.org/2000/svg',
        'width': str(cropped_width),
        'height': str(cropped_height),
        'viewBox': f'0 0 {cropped_width} {cropped_height}'
    })

    # Convert mask to paths with hole detection
    paths = mask_to_svg_paths_with_holes(cropped_mask, simplify_epsilon=1.0)

    # Color mapping
    color_hex = {
        'black': '#000000',
        'grey': '#808080',
        'gray': '#808080',
        'white': '#FFFFFF',
    }.get(color, color)

    # Add paths to SVG (no translation needed since we cropped the mask)
    for path_data in paths:
        ET.SubElement(svg, 'path', {
            'd': path_data,
            'fill': color_hex,
            'fill-rule': 'evenodd'
        })

    # Write SVG
    tree = ET.ElementTree(svg)
    ET.indent(tree, space='  ')
    tree.write(output_path, encoding='utf-8', xml_declaration=True)

    if crop_to_bounds:
        print(f"Created {output_path} with color {color} ({len(paths)} shape groups, cropped to {cropped_width}x{cropped_height})")
    else:
        print(f"Created {output_path} with color {color} ({len(paths)} shape groups)")


def main():
    parser = argparse.ArgumentParser(
        description='Split PNG into separate SVG files by brightness/color'
    )
    parser.add_argument('input', help='Input PNG file')
    parser.add_argument('--output-dir', '-o', help='Output directory (default: same as input)')
    parser.add_argument('--prefix', '-p', help='Prefix for output files (default: input filename)')
    parser.add_argument('--dark-threshold', type=int, default=30,
                       help='Brightness threshold for dark/black elements (default: 30)')
    parser.add_argument('--mid-min', type=int, default=108,
                       help='Minimum brightness for mid/grey elements (default: 108)')
    parser.add_argument('--mid-max', type=int, default=118,
                       help='Maximum brightness for mid/grey elements (default: 118)')
    parser.add_argument('--no-crop', action='store_true',
                       help='Do not crop to bounding box (default: crop enabled)')
    args = parser.parse_args()

    input_path = Path(args.input)

    if not input_path.exists():
        print(f"ERROR: File not found: {input_path}")
        return 1

    # Setup output directory
    if args.output_dir:
        output_dir = Path(args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
    else:
        output_dir = input_path.parent

    # Setup prefix
    prefix = args.prefix if args.prefix else input_path.stem

    # Load image to get size
    img = Image.open(input_path)
    image_size = img.size

    # Extract color regions
    print(f"Analyzing colors in {input_path}...")
    masks = extract_color_regions(input_path, args.dark_threshold, args.mid_min, args.mid_max)

    if not masks:
        print("No color regions found!")
        return 1

    print(f"\nFound {len(masks)} color regions:")
    for color, mask in masks.items():
        pixel_count = np.sum(mask)
        print(f"  {color}: {pixel_count:,} pixels")

    # Create SVG for each color
    print(f"\nCreating separate SVG files...")
    crop_enabled = not args.no_crop
    for color, mask in masks.items():
        output_path = output_dir / f"{prefix}_{color}.svg"
        create_svg_from_mask(mask, output_path, color, image_size, crop_to_bounds=crop_enabled)

    print(f"\nDone! Created {len(masks)} SVG files in {output_dir}")


if __name__ == '__main__':
    main()