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 white_color_mask(image: np.ndarray, min_brightness: int = 220) -> np.ndarray:
    alpha = image[..., 3] > 32
    rgb = image[..., :3]
    brightness = np.min(rgb, axis=-1)
    return alpha & (brightness >= min_brightness)


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):
    image = render_svg_to_image(svg_path, pixel_width=png_width)
    width_px, height_px = image.size
    scale = width_mm / width_px
    height_mm = height_px * scale

    img_array = np.array(image)
    green_mask = clean_mask(color_mask(img_array, (57, 233, 145), tolerance=80))
    white_mask = clean_mask(white_color_mask(img_array, min_brightness=220))

    # Split green mask into top (icon) and bottom (text) regions
    # Find where white text is located to use as separator
    white_rows = np.any(white_mask, axis=1)
    white_y_coords = np.where(white_rows)[0]

    if len(white_y_coords) > 0:
        white_y_mid = (white_y_coords[0] + white_y_coords[-1]) // 2

        # Green icon: everything above white text
        green_icon_mask = green_mask.copy()
        green_icon_mask[white_y_coords[0]:, :] = False  # Zero out everything from white text onwards

        # Green bottom text: everything below white text
        green_bottom_mask = green_mask.copy()
        green_bottom_mask[:white_y_coords[-1], :] = False  # Zero out everything before end of white text

        green_icon_polys = mask_to_polygons(green_icon_mask)
        green_bottom_polys = mask_to_polygons(green_bottom_mask)
    else:
        # No white text found, treat all green as one
        green_icon_polys = mask_to_polygons(green_mask)
        green_bottom_polys = []

    white_polys = mask_to_polygons(white_mask)

    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))

    # Features sit exactly on top of the base (no overlap to avoid slicing issues)
    green_icon_mesh = create_extruded_mesh(green_icon_polys, feature_height, scale)
    green_bottom_mesh = create_extruded_mesh(green_bottom_polys, feature_height, scale)
    white_mesh = create_extruded_mesh(white_polys, feature_height, scale)

    for mesh in (green_icon_mesh, green_bottom_mesh, white_mesh):
        if mesh is not None:
            # Position features to start exactly at the top of the base
            mesh.apply_translation((0.0, height_mm, base_thickness))

    return base, green_icon_mesh, green_bottom_mesh, white_mesh


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) -> int:
    basematerials = ET.SubElement(resources, 'basematerials', {'id': '1'})
    ET.SubElement(basematerials, 'base', {
        'name': 'Base',
        'displaycolor': '#000000'
    })
    ET.SubElement(basematerials, 'base', {
        'name': 'Green',
        'displaycolor': '#39e991'
    })
    ET.SubElement(basematerials, 'base', {
        'name': 'White',
        'displaycolor': '#ffffff'
    })
    return 1


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')
    _add_basematerials(resources)

    object_ids = []
    current_id = 1
    material_map = {'base': 0, 'green_icon': 1, 'green_text': 1, 'white_text': 2, 'logo': 0}
    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')
    basematerials_id = _add_basematerials(resources)

    # Create separate object entries for each mesh
    object_ids = []
    current_id = 1
    material_map = {'base': 0, 'green_icon': 1, 'green_text': 1, 'white_text': 2}
    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")
    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")
    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, green_icon_mesh, green_bottom_mesh, white_mesh = build_logo_meshes(
        svg_path,
        args.width_mm,
        args.base_thickness,
        args.feature_height,
        args.png_width,
    )

    base_path = output_prefix.with_name(output_prefix.stem + "_base.stl")
    green_icon_path = output_prefix.with_name(output_prefix.stem + "_green_icon.stl")
    green_bottom_path = output_prefix.with_name(output_prefix.stem + "_green_bottom.stl")
    white_path = output_prefix.with_name(output_prefix.stem + "_white.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_green_icon = save_mesh(green_icon_mesh, green_icon_path)
    saved_green_bottom = save_mesh(green_bottom_mesh, green_bottom_path)
    saved_white = save_mesh(white_mesh, white_path)

    # Concatenate all meshes into a single object
    assembled_mesh = trimesh.util.concatenate([m for m in (base_mesh, green_icon_mesh, green_bottom_mesh, white_mesh) 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
    save_3mf_parts({
        'base': base_mesh,
        'green_icon': green_icon_mesh,
        'white_text': white_mesh,
        'green_text': green_bottom_mesh,
    }, parts_3mf_path)

    print("Created files:")
    if saved_base:
        print(f"  Base STL:               {saved_base}")
    if saved_green_icon:
        print(f"  Green icon STL:         {saved_green_icon}")
    if saved_green_bottom:
        print(f"  Green bottom text STL:  {saved_green_bottom}")
    if saved_white:
        print(f"  White text STL:         {saved_white}")
    print(f"  Assembled STL:          {assembled_stl_path}")
    print(f"  Assembled 3MF:          {assembled_3mf_path}")
    print(f"  Parts 3MF:              {parts_3mf_path}")


if __name__ == "__main__":
    main()