"""Composite utilities for combining node results."""

import numpy as np
from typing import List, Tuple
from PIL import Image


def composite_nodes(nodes: List[dict], bboxes: List[Tuple[float, float, float, float]],
                    container_bbox: Tuple[float, float, float, float]) -> Tuple[np.ndarray, np.ndarray]:
    """Composite multiple nodes into a single mask.
    
    Args:
        nodes: List of node dictionaries with "mask" key
        bboxes: List of (x, y, w, h) bounding boxes for each node
        container_bbox: Container bounding box (x, y, w, h)
    
    Returns:
        Tuple of (composite_mask, composite_sdf)
        - composite_mask: Combined mask (H, W)
        - composite_sdf: Combined SDF (H, W) - simplified version
    """
    # Handle container_bbox format
    if isinstance(container_bbox, (tuple, list)) and len(container_bbox) == 4:
        cx, cy, cw, ch = container_bbox
    elif isinstance(container_bbox, dict):
        cx = container_bbox.get("x", 0)
        cy = container_bbox.get("y", 0)
        cw = container_bbox.get("width", container_bbox.get("w", 0))
        ch = container_bbox.get("height", container_bbox.get("h", 0))
    else:
        raise ValueError(f"Invalid container_bbox format: {container_bbox}")
    
    cw_int = int(float(cw))
    ch_int = int(float(ch))
    
    # Initialize composite mask
    composite_mask = np.zeros((ch_int, cw_int), dtype=np.float32)
    
    # Place each node's mask at its bbox position
    for i, (node, bbox) in enumerate(zip(nodes, bboxes)):
        mask = node.get("mask")
        if mask is None:
            continue
        
        # Handle different bbox formats
        if isinstance(bbox, (tuple, list)) and len(bbox) == 4:
            x, y, w, h = bbox
        elif isinstance(bbox, dict):
            x = bbox.get("x", 0)
            y = bbox.get("y", 0)
            w = bbox.get("width", bbox.get("w", 0))
            h = bbox.get("height", bbox.get("h", 0))
        else:
            print(f"Warning: Invalid bbox format at index {i}: {bbox} (type: {type(bbox)})")
            continue
        
        x_int = int(float(x))
        y_int = int(float(y))
        w_int = int(float(w))
        h_int = int(float(h))
        
        # Resize mask to bbox size if needed
        if mask.shape != (h_int, w_int):
            from scipy.ndimage import zoom
            zoom_y = h_int / mask.shape[0]
            zoom_x = w_int / mask.shape[1]
            mask_resized = zoom(mask, (zoom_y, zoom_x), order=1)
        else:
            mask_resized = mask
        
        # Clip to valid range
        mask_resized = np.clip(mask_resized, 0, 1)
        
        # Place mask in composite
        y_end = min(y_int + h_int, ch_int)
        x_end = min(x_int + w_int, cw_int)
        y_start = max(0, y_int)
        x_start = max(0, x_int)
        
        if y_end > y_start and x_end > x_start:
            mask_crop = mask_resized[:y_end-y_start, :x_end-x_start]
            composite_mask[y_start:y_end, x_start:x_end] = np.maximum(
                composite_mask[y_start:y_end, x_start:x_end],
                mask_crop
            )
    
    # Generate simplified SDF from composite mask
    # This is a simplified version - full SDF generation would use distance_transform_edt
    composite_sdf = composite_mask - 0.5  # Simple approximation
    
    return composite_mask, composite_sdf