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 = 32) -> np.ndarray:
    if min_size > 0:
        cleaned = morphology.remove_small_objects(mask, max_size=min_size - 1)
    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 = 10.0, simplify_tolerance: float = 1.0):
    contours = measure.find_contours(mask.astype(np.uint8), 0.5)
    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)
        if poly.is_valid and poly.area >= min_area:
            polygons.append(poly)

    for hole in holes:
        if hole not in assigned_holes:
            poly = hole.simplify(simplify_tolerance)
            if poly.is_valid and poly.area >= min_area:
                polygons.append(poly)

    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
    return trimesh.util.concatenate(meshes) if meshes else None


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

    green_polys = mask_to_polygons(green_mask)
    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))

    green_mesh = create_extruded_mesh(green_polys, feature_height, scale)
    white_mesh = create_extruded_mesh(white_polys, feature_height, scale)
    for mesh in (green_mesh, white_mesh):
        if mesh is not None:
            mesh.apply_translation((0.0, height_mm, base_thickness))

    return base, green_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': 1, 'green': 2, 'white': 3}
    for name, mesh in meshes.items():
        if mesh is None:
            continue
        material_id = material_map.get(name, 1)
        resource = _mesh_to_3mf_object(mesh, current_id, name, material_id)
        resources.append(resource)
        object_ids.append(current_id)
        current_id += 1

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

    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_3mf_parts(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': 1, 'green': 2, 'white': 3}
    for name, mesh in meshes.items():
        if mesh is None:
            continue
        material_id = material_map.get(name, 1)
        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')
    for object_id in object_ids:
        ET.SubElement(build, 'item', {'objectid': str(object_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_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_path = output_prefix.with_name(output_prefix.stem + "_green.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 = save_mesh(green_mesh, green_path)
    saved_white = save_mesh(white_mesh, white_path)

    assembled_mesh = trimesh.util.concatenate([m for m in (base_mesh, green_mesh, white_mesh) if m is not None])
    assembled_mesh.export(assembled_stl_path)

    save_3mf({
        'base': base_mesh,
        'green': green_mesh,
        'white': white_mesh,
    }, assembled_3mf_path)
    save_3mf_parts({
        'base': base_mesh,
        'green': green_mesh,
        'white': white_mesh,
    }, parts_3mf_path)

    print("Created files:")
    if saved_base:
        print(f"  Base STL:      {saved_base}")
    if saved_green:
        print(f"  Green STL:     {saved_green}")
    if saved_white:
        print(f"  White 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()