import argparse
import io
import zipfile
import xml.etree.ElementTree as ET
from pathlib import Path

import cairosvg
import numpy as np
import trimesh
from PIL import Image
from shapely.geometry import LinearRing, Polygon
from shapely.ops import unary_union
from skimage import measure, morphology


def render_svg_to_image(svg_path: Path, pixel_width: int = 2048) -> Image.Image:
    png_output = io.BytesIO()
    cairosvg.svg2png(url=str(svg_path), write_to=png_output, output_width=pixel_width)
    png_output.seek(0)
    return Image.open(png_output).convert("RGBA")


def clean_mask(mask: np.ndarray, min_size: int = 10) -> np.ndarray:
    # Use smaller min_size to preserve small features like parts of the icon
    if min_size > 0:
        cleaned = morphology.remove_small_objects(mask, min_size=min_size)
    else:
        cleaned = mask
    structure = np.ones((3, 3), dtype=bool)
    cleaned = morphology.closing(cleaned, structure)
    return cleaned


def color_mask(image: np.ndarray, target_rgb: tuple[int, int, int], tolerance: int = 64) -> np.ndarray:
    alpha = image[..., 3] > 32
    diff = np.linalg.norm(image[..., :3].astype(np.int16) - np.array(target_rgb, dtype=np.int16), axis=-1)
    return alpha & (diff <= tolerance)


def detect_color_clusters(image: np.ndarray, tolerance: int = 80, min_pixels: int = 100):
    """
    Automatically detect distinct color regions in an image.
    Returns list of (color_name, target_rgb, mask) tuples.
    """
    img_array = np.array(image)
    alpha = img_array[..., 3] > 32
    rgb_pixels = img_array[alpha][:, :3]

    if len(rgb_pixels) == 0:
        return []

    # Get unique colors and their counts
    unique_colors, counts = np.unique(rgb_pixels, axis=0, return_counts=True)

    # Sort by count (most common first)
    sorted_indices = np.argsort(-counts)
    unique_colors = unique_colors[sorted_indices]
    counts = counts[sorted_indices]

    # Cluster similar colors together
    color_clusters = []
    used = set()

    for i, color in enumerate(unique_colors):
        if i in used:
            continue

        # Find all colors within tolerance of this one
        cluster_mask = np.zeros(img_array.shape[:2], dtype=bool)
        cluster_pixels = 0

        for j, other_color in enumerate(unique_colors):
            if j in used:
                continue
            diff = np.linalg.norm(color.astype(np.int16) - other_color.astype(np.int16))
            if diff <= tolerance:
                mask = color_mask(img_array, tuple(other_color), tolerance=tolerance)
                cluster_mask |= mask
                cluster_pixels += counts[j]
                used.add(j)

        if cluster_pixels >= min_pixels:
            # Generate color name
            r, g, b = color
            if r > 200 and g > 200 and b > 200:
                color_name = 'white'
            elif r < 50 and g < 50 and b < 50:
                color_name = 'black'
            elif r > max(g, b):
                color_name = 'red'
            elif g > max(r, b):
                color_name = 'green'
            elif b > max(r, g):
                color_name = 'blue'
            elif r > 200 and g > 200:
                color_name = 'yellow'
            elif r > 200 and b > 200:
                color_name = 'magenta'
            elif g > 200 and b > 200:
                color_name = 'cyan'
            else:
                color_name = f'color_{len(color_clusters)+1}'

            color_clusters.append((color_name, tuple(color), cluster_mask, cluster_pixels))

    # Sort by pixel count (largest first)
    color_clusters.sort(key=lambda x: x[3], reverse=True)

    return [(name, rgb, mask) for name, rgb, mask, _ in color_clusters]


def signed_area(coords: np.ndarray) -> float:
    x = coords[:, 0]
    y = coords[:, 1]
    return 0.5 * np.sum(x[:-1] * y[1:] - x[1:] * y[:-1])


def mask_to_polygons(mask: np.ndarray, min_area: float = 1.0, simplify_tolerance: float = 0.25, debug=False):
    contours = measure.find_contours(mask.astype(np.uint8), 0.5)
    if debug:
        print(f"    Found {len(contours)} contours")
    shapes: list[tuple[Polygon, float]] = []
    for contour in contours:
        if contour.shape[0] < 4:
            continue
        coords = np.column_stack((contour[:, 1], contour[:, 0]))
        if not np.allclose(coords[0], coords[-1]):
            coords = np.vstack([coords, coords[0]])
        area = signed_area(coords)
        ring = LinearRing(coords)
        if not ring.is_valid or ring.length == 0:
            continue
        poly = Polygon(ring)
        if not poly.is_valid or abs(area) < min_area:
            continue
        shapes.append((poly, area))

    if not shapes:
        return []

    exteriors: list[Polygon] = [poly for poly, area in shapes if area > 0]
    holes: list[Polygon] = [poly for poly, area in shapes if area < 0]
    if not exteriors:
        shapes = sorted(shapes, key=lambda item: abs(item[1]), reverse=True)
        exteriors = [shapes[0][0]]
        holes = [poly for poly, _ in shapes[1:]]

    polygons = []
    assigned_holes = set()
    for exterior in exteriors:
        hole_list = []
        for hole in holes:
            if exterior.contains(hole.representative_point()):
                hole_list.append(hole.exterior.coords)
                assigned_holes.add(hole)
        poly = Polygon(exterior.exterior.coords, hole_list)
        poly = poly.simplify(simplify_tolerance, preserve_topology=True)
        if poly.is_valid and poly.area >= min_area:
            polygons.append(poly)
        elif debug:
            print(f"      Skipped polygon: valid={poly.is_valid}, area={poly.area:.1f}")

    for hole in holes:
        if hole not in assigned_holes:
            poly = hole.simplify(simplify_tolerance, preserve_topology=True)
            if poly.is_valid and poly.area >= min_area:
                polygons.append(poly)
            elif debug:
                print(f"      Skipped unassigned hole: valid={poly.is_valid}, area={poly.area:.1f}")

    if debug:
        print(f"    Output: {len(polygons)} polygons (from {len(exteriors)} exteriors, {len(holes)} holes)")

    return polygons


def create_extruded_mesh(polygons, height_mm: float, scale: float, y_flip: bool = True):
    meshes = []
    for poly in polygons:
        exterior = [(x * scale, ((-y if y_flip else y) * scale)) for x, y in poly.exterior.coords]
        holes = [
            [(x * scale, ((-y if y_flip else y) * scale)) for x, y in ring.coords]
            for ring in poly.interiors
        ]
        polygon = Polygon(exterior, holes)
        if not polygon.is_valid or polygon.area == 0:
            continue
        try:
            mesh = trimesh.creation.extrude_polygon(polygon, height_mm)
            meshes.append(mesh)
        except Exception:
            continue

    if not meshes:
        return None

    # Process each mesh individually to ensure they're watertight
    watertight_meshes = []
    skipped_count = 0
    for mesh in meshes:
        # Try to make it watertight
        trimesh.repair.fill_holes(mesh)
        trimesh.repair.fix_normals(mesh)
        trimesh.repair.fix_winding(mesh)

        # If still not watertight, try to split and fix components
        if not mesh.is_watertight:
            # Split into connected components
            components = mesh.split(only_watertight=False)
            for comp in components:
                trimesh.repair.fill_holes(comp)
                if comp.is_watertight or comp.is_volume:
                    watertight_meshes.append(comp)
                else:
                    skipped_count += 1
        else:
            watertight_meshes.append(mesh)

    if not watertight_meshes:
        # Fallback: return combined mesh even if not perfect
        return trimesh.util.concatenate(meshes)

    return trimesh.util.concatenate(watertight_meshes)


def build_logo_meshes(svg_path: Path, width_mm: float, base_thickness: float, feature_height: float, png_width: int,
                      plate_width_mm: float = None, plate_height_mm: float = None, margin_mm: float = 0):
    """
    Generalized function to convert any SVG to 3D meshes.
    Automatically detects colors and creates one part per color.

    Args:
        svg_path: Path to SVG file
        width_mm: Target width in mm (if plate dimensions not specified)
        base_thickness: Thickness of base plate in mm
        feature_height: Height of raised features in mm
        png_width: Resolution for SVG rasterization
        plate_width_mm: Optional plate width - logo will be scaled and centered
        plate_height_mm: Optional plate height - logo will be scaled and centered
        margin_mm: Margin around logo when using plate dimensions

    Returns (base_mesh, color_parts_dict) where color_parts_dict maps color names to meshes.
    """
    image = render_svg_to_image(svg_path, pixel_width=png_width)
    width_px, height_px = image.size

    # Calculate scaling based on whether plate dimensions are specified
    if plate_width_mm is not None and plate_height_mm is not None:
        # Calculate available space after margins
        available_width = plate_width_mm - (2 * margin_mm)
        available_height = plate_height_mm - (2 * margin_mm)

        # Scale to fit within available space (maintaining aspect ratio)
        scale_x = available_width / width_px
        scale_y = available_height / height_px
        scale = min(scale_x, scale_y)

        # Actual logo dimensions after scaling
        logo_width_mm = width_px * scale
        logo_height_mm = height_px * scale

        # Calculate centering offsets
        x_offset = (plate_width_mm - logo_width_mm) / 2.0
        y_offset = (plate_height_mm - logo_height_mm) / 2.0

        # Use plate dimensions for base
        width_mm = plate_width_mm
        height_mm = plate_height_mm
    else:
        # Original behavior: use specified width, calculate height
        scale = width_mm / width_px
        logo_width_mm = width_mm
        logo_height_mm = height_px * scale
        height_mm = logo_height_mm
        x_offset = 0
        y_offset = 0

    # Automatically detect all color regions
    color_clusters = detect_color_clusters(image, tolerance=80, min_pixels=100)

    print(f"Detected {len(color_clusters)} color regions:")
    for name, rgb, mask in color_clusters:
        pixel_count = np.sum(mask)
        print(f"  {name}: RGB{rgb} - {pixel_count:,} pixels")

    # Create base plate
    base = trimesh.creation.box(extents=(width_mm, height_mm, base_thickness))
    base.apply_translation((width_mm / 2.0, height_mm / 2.0, base_thickness / 2.0))

    # Create mesh for each color
    color_parts = {}
    for name, rgb, mask in color_clusters:
        cleaned_mask = clean_mask(mask)
        polygons = mask_to_polygons(cleaned_mask)

        if polygons:
            mesh = create_extruded_mesh(polygons, feature_height, scale)
            if mesh is not None:
                # Position features to start exactly at the top of the base
                # Apply centering offset when plate dimensions are used
                mesh.apply_translation((x_offset, height_mm - y_offset, base_thickness))
                color_parts[name] = mesh

    if plate_width_mm is not None and plate_height_mm is not None:
        print(f"\nPlate dimensions: {plate_width_mm:.1f} x {plate_height_mm:.1f} mm")
        print(f"Logo dimensions: {logo_width_mm:.1f} x {logo_height_mm:.1f} mm")
        print(f"Margin: {margin_mm:.1f} mm")
        print(f"Centering offset: ({x_offset:.1f}, {y_offset:.1f}) mm")

    return base, color_parts


def _mesh_to_3mf_object(mesh: trimesh.Trimesh, object_id: int, name: str, material_id: int) -> ET.Element:
    obj = ET.Element('object', {
        'id': str(object_id),
        'name': name,
        'type': 'model'
    })
    mesh_el = ET.SubElement(obj, 'mesh')
    vertices_el = ET.SubElement(mesh_el, 'vertices')
    for vertex in mesh.vertices:
        ET.SubElement(vertices_el, 'vertex', {
            'x': str(float(vertex[0])),
            'y': str(float(vertex[1])),
            'z': str(float(vertex[2]))
        })
    triangles_el = ET.SubElement(mesh_el, 'triangles')
    for face in mesh.faces:
        ET.SubElement(triangles_el, 'triangle', {
            'v1': str(int(face[0])),
            'v2': str(int(face[1])),
            'v3': str(int(face[2])),
            'materialid': str(material_id)
        })
    return obj


def _add_basematerials(resources: ET.Element, color_names: list[str]) -> dict[str, int]:
    """
    Add basematerials for all colors.
    Returns dict mapping color name to material ID.
    """
    basematerials = ET.SubElement(resources, 'basematerials', {'id': '1'})

    # Color palette for display
    color_palette = {
        'base': '#000000',
        'black': '#000000',
        'white': '#FFFFFF',
        'red': '#FF0000',
        'green': '#00FF00',
        'blue': '#0000FF',
        'yellow': '#FFFF00',
        'magenta': '#FF00FF',
        'cyan': '#00FFFF',
    }

    material_map = {}
    for i, name in enumerate(['base'] + color_names):
        display_color = color_palette.get(name, '#808080')  # Default to gray
        ET.SubElement(basematerials, 'base', {
            'name': name.capitalize(),
            'displaycolor': display_color
        })
        material_map[name] = i

    return material_map


def save_3mf(meshes: dict[str, trimesh.Trimesh], path: Path):
    model = ET.Element('model', {
        'xmlns': 'http://schemas.microsoft.com/3dmanufacturing/core/2015/02',
        'xmlns:m': 'http://schemas.microsoft.com/3dmanufacturing/material/2015/02',
        'unit': 'millimeter'
    })
    resources = ET.SubElement(model, 'resources')

    # Get list of color names (excluding 'base')
    color_names = [name for name in meshes.keys() if name != 'base' and meshes[name] is not None]
    material_map = _add_basematerials(resources, color_names)

    object_ids = []
    current_id = 1
    for name, mesh in meshes.items():
        if mesh is None:
            continue
        material_id = material_map.get(name, 0)
        resource = _mesh_to_3mf_object(mesh, current_id, name, material_id)
        resources.append(resource)
        object_ids.append(current_id)
        current_id += 1

    build = ET.SubElement(model, 'build')

    # If we have only one object, add it directly to build (no assembly)
    # If we have multiple objects, create an assembly
    if len(object_ids) == 1:
        ET.SubElement(build, 'item', {'objectid': str(object_ids[0])})
    else:
        assembly_id = current_id
        assembly = ET.SubElement(resources, 'object', {
            'id': str(assembly_id),
            'name': 'assembly',
            'type': 'model'
        })
        components_el = ET.SubElement(assembly, 'components')
        for object_id in object_ids:
            ET.SubElement(components_el, 'component', {
                'objectid': str(object_id)
            })
        ET.SubElement(build, 'item', {'objectid': str(assembly_id)})

    xml_data = ET.tostring(model, encoding='utf-8', xml_declaration=True)
    with zipfile.ZipFile(path, 'w', compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr('[Content_Types].xml', '''<?xml version="1.0" encoding="utf-8"?>\n<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="model" ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/></Types>''')
        zf.writestr('_rels/.rels', '''<?xml version="1.0" encoding="utf-8"?>\n<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel" Target="/3D/3dmodel.model" Id="rel0"/></Relationships>''')
        zf.writestr('3D/3dmodel.model', xml_data)
    return path


def save_3mf_parts(meshes: dict[str, trimesh.Trimesh], path: Path):
    """
    Save meshes as separate objects in a 3MF assembly, similar to nameplate.3mf.
    Each mesh becomes a separate selectable component in BambuStudio.
    """
    model = ET.Element('model', {
        'xmlns': 'http://schemas.microsoft.com/3dmanufacturing/core/2015/02',
        'xmlns:m': 'http://schemas.microsoft.com/3dmanufacturing/material/2015/02',
        'unit': 'millimeter'
    })
    resources = ET.SubElement(model, 'resources')

    # Get list of color names (excluding 'base')
    color_names = [name for name in meshes.keys() if name != 'base' and meshes[name] is not None]
    material_map = _add_basematerials(resources, color_names)

    # Create separate object entries for each mesh
    object_ids = []
    current_id = 1
    for name, mesh in meshes.items():
        if mesh is None:
            continue
        material_id = material_map.get(name, 0)
        resource = _mesh_to_3mf_object(mesh, current_id, name, material_id)
        resources.append(resource)
        object_ids.append(current_id)
        current_id += 1

    # Create an assembly object that references all the individual objects
    # This is key: the assembly is what gets added to the build, not the individual objects
    assembly_id = current_id
    assembly = ET.SubElement(resources, 'object', {
        'id': str(assembly_id),
        'name': 'logo_assembly',
        'type': 'model'
    })
    components_el = ET.SubElement(assembly, 'components')

    # Add each object as a component in the assembly
    # Identity transform (no translation/rotation)
    for object_id in object_ids:
        ET.SubElement(components_el, 'component', {
            'objectid': str(object_id),
            'transform': '1 0 0 0 1 0 0 0 1 0 0 0'  # Identity matrix
        })

    # Add only the assembly to the build (not the individual objects)
    build = ET.SubElement(model, 'build')
    ET.SubElement(build, 'item', {'objectid': str(assembly_id)})

    xml_data = ET.tostring(model, encoding='utf-8', xml_declaration=True)
    with zipfile.ZipFile(path, 'w', compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr('[Content_Types].xml', '''<?xml version="1.0" encoding="utf-8"?>\n<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="model" ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/></Types>''')
        zf.writestr('_rels/.rels', '''<?xml version="1.0" encoding="utf-8"?>\n<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel" Target="/3D/3dmodel.model" Id="rel0"/></Relationships>''')
        zf.writestr('3D/3dmodel.model', xml_data)
    return path


def save_mesh(mesh, path: Path):
    if mesh is None:
        return None
    mesh.export(path)
    return path


def main():
    parser = argparse.ArgumentParser(description="Convert a color SVG logo into a layered STL for multi-color printing.")
    parser.add_argument("svg", help="Input SVG file path")
    parser.add_argument("output", help="Base output file path (without extension is fine)")
    parser.add_argument("--width-mm", type=float, default=100.0, help="Final model width in millimeters (if plate dimensions not specified)")
    parser.add_argument("--base-thickness", type=float, default=0.8, help="Thickness of the black base in mm")
    parser.add_argument("--feature-height", type=float, default=1.8, help="Height of the raised logo features above the base in mm")
    parser.add_argument("--png-width", type=int, default=2048, help="Rasterization width for SVG rendering")
    parser.add_argument("--plate-width", type=float, help="Plate width in mm - logo will be scaled and centered")
    parser.add_argument("--plate-height", type=float, help="Plate height in mm - logo will be scaled and centered")
    parser.add_argument("--margin", type=float, default=0, help="Margin around logo when using plate dimensions (default: 0)")
    args = parser.parse_args()

    svg_path = Path(args.svg)
    output_prefix = Path(args.output)
    output_prefix.parent.mkdir(parents=True, exist_ok=True)

    base_mesh, color_parts = build_logo_meshes(
        svg_path,
        args.width_mm,
        args.base_thickness,
        args.feature_height,
        args.png_width,
        plate_width_mm=args.plate_width,
        plate_height_mm=args.plate_height,
        margin_mm=args.margin,
    )

    # Save individual STL files for each color part
    base_path = output_prefix.with_name(output_prefix.stem + "_base.stl")
    assembled_stl_path = output_prefix.with_name(output_prefix.stem + "_assembled.stl")
    assembled_3mf_path = output_prefix.with_name(output_prefix.stem + "_assembled.3mf")
    parts_3mf_path = output_prefix.with_name(output_prefix.stem + "_parts.3mf")

    saved_base = save_mesh(base_mesh, base_path)
    saved_parts = {}
    for color_name, mesh in color_parts.items():
        part_path = output_prefix.with_name(output_prefix.stem + f"_{color_name}.stl")
        if save_mesh(mesh, part_path):
            saved_parts[color_name] = part_path

    # Concatenate all meshes into a single object
    all_meshes = [base_mesh] + list(color_parts.values())
    assembled_mesh = trimesh.util.concatenate([m for m in all_meshes if m is not None])
    assembled_mesh.export(assembled_stl_path)

    # For 3MF assembled version, export as a SINGLE mesh object
    # BambuStudio can then use the paint tool to assign colors
    save_3mf({
        'logo': assembled_mesh,
    }, assembled_3mf_path)

    # Export as separate parts in assembly for multi-color selection
    parts_dict = {'base': base_mesh}
    parts_dict.update(color_parts)
    save_3mf_parts(parts_dict, parts_3mf_path)

    print("\nCreated files:")
    if saved_base:
        print(f"  Base STL:               {saved_base}")
    for color_name, part_path in saved_parts.items():
        print(f"  {color_name.capitalize()} STL:              {part_path}")
    print(f"  Assembled STL:          {assembled_stl_path}")
    print(f"  Assembled 3MF:          {assembled_3mf_path}")
    print(f"  Parts 3MF:              {parts_3mf_path}")
    print(f"\nImport {parts_3mf_path.name} into BambuStudio for multi-color printing!")


if __name__ == "__main__":
    main()