#!/usr/bin/env python3
"""Generate a 4-level terraced 3D badge from the Pyrrhic Victory tech icon.

The pixel art has exactly 3 palette colors (the rest is dither/A-A):
  black   #000000
  brown   ~#643A2A
  yellow  #FFFF00   (split into inner detail vs outer border ring)

Stacked bottom -> top (each 2.0mm, total 8mm, flat bottom, supported):
  L1  black          (lowest)
  L2  brown          (figure + steps)
  L3  yellow (inner) (face/trophy detail)
  L4  border-yellow  (outer rim, highest)
Outputs: pyrrhic_victory.3mf (per-level objects w/ materials),
         pyrrhic_victory.stl (combined, monochrome), legend.txt
"""
import shutil
import zipfile
from pathlib import Path

import numpy as np
from PIL import Image
from scipy import ndimage
from skimage.measure import find_contours
import manifold3d as M
import trimesh

SRC = Path("../mods/pycoalprocessinggraphics/graphics/technology/pyrrhic.png")
OUT = Path(__file__).resolve().parent
BASE = Path("/tmp/opencode/pyrrhic-build")

FOOTPRINT_MM = 51.2        # 128 px -> 4x supersample -> 512 px
PX = FOOTPRINT_MM / 512    # 0.1 mm per supersampled pixel
NLEVELS = 4
LEVEL_MM = 8.0 / NLEVELS   # 2.0 mm per level

# ---------------------------------------------------------------- load image
SUPER = 4
N = 128 * SUPER
im = Image.open(SRC).convert("RGBA").resize((N, N), Image.NEAREST)
a = np.array(im)
alpha = a[..., 3]
opaque = alpha > 128
r, g, b = a[..., 0], a[..., 1], a[..., 2]

def mean_hex(mask):
    rgb = a[mask][..., :3].mean(axis=0).astype(int)
    return f"#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}".upper()

# nearest-true-palette classification (the art is exactly 3 colors; the rest is dither/A-A)
PAL = np.array([[0, 0, 0], [100, 58, 42], [255, 255, 0]], dtype=np.float32)  # black, brown, yellow
rgb = a[..., :3].astype(np.float32)
d = np.linalg.norm(rgb[None, ...] - PAL[:, None, None, :], axis=3)  # (3,N,N)
cls = d.argmin(axis=0)  # 0=black 1=brown 2=yellow
black  = opaque & (cls == 0)
brown  = opaque & (cls == 1)
yellow = opaque & (cls == 2)

# split yellow into the outer border ring vs the inner detail (face/trophy)
lab, ncomp = ndimage.label(yellow)
big = max(range(1, ncomp + 1), key=lambda i: (lab == i).sum())
border   = lab == big
inner_y  = yellow & ~border
assert border.any() and inner_y.any(), "expected both border and inner yellow"
for name, m in (("black", black), ("brown", brown), ("inner-yellow", inner_y), ("border-yellow", border)):
    print(f"{name:<14} px(512^2)={m.sum():>6}")
assert ((black | brown | yellow) & ~opaque).sum() == 0, "palette color where transparent"
assert opaque.sum() == (black | brown | yellow).sum(), "opaque pixel not assigned to a palette color"

# canonical (true) colors: black / brown / yellow
TRUE_HEX = {"black": "#000000", "brown": mean_hex(brown), "yellow": "#FFFF00"}
print("canonical colors:", TRUE_HEX)

# cumulative stacking (each layer sits on the one below -> supported, no cantilevers).
# Height order bottom->top: black lowest, brown mid, inner-yellow, border-yellow highest.
full = black | brown | yellow
regions = [
    ("black",            full),          # base slab: whole badge
    ("brown",            brown | yellow),# mid slab: brown + all yellow
    ("yellow",           inner_y | border),  # level 3: inner detail + border
    ("border-yellow",    border),        # level 4 (top): border ring only
]

# ---------------------------------------------------------------- polygonize + extrude
def to_crosssection(mask):
    """Boolean mask -> manifold CrossSection in mm (y flipped: top of art = +Y).

    Each row->Y, col->X cell; contour coords from find_contours are (row, col).
    """
    contours = [[(x * PX, (N - y) * PX) for (y, x) in c]
                for c in find_contours(mask, fully_connected="low") if len(c) > 4]
    if not contours:
        return None
    return M.CrossSection(contours, M.FillRule.EvenOdd)

def manifold_to_trimesh(m):
    mm = m.to_mesh()
    return trimesh.Trimesh(np.asarray(mm.vert_properties),
                           np.asarray(mm.tri_verts), process=False)

def mask_to_scad_polymask(mask):
    """mask -> list of (outer_pts, [hole_pts...]) using shapely pixel-box union.

    unary_union does boolean topology, so holes + interlocking bands come out clean.
    """
    from shapely import box, unary_union
    idx = np.argwhere(mask)
    if len(idx) == 0:
        return []
    xs, ys = idx[:, 1], idx[:, 0]
    us = unary_union([box(x * PX, y * PX, (x + 1) * PX, (y + 1) * PX) for x, y in zip(xs, ys)])
    polys = list(us.geoms) if us.geom_type == "MultiPolygon" else [us]
    out = []
    for p in polys:
        if not p.is_valid:
            p = p.buffer(0)
        outer = [(x, y) for x, y in p.exterior.coords][:-1]
        holes = [[(x, y) for x, y in i.coords][:-1] for i in p.interiors]
        out.append((outer, holes))
    return out

shutil.rmtree(BASE / "meshes", ignore_errors=True)
(BASE / "meshes").mkdir(parents=True)

solids = []
for i, (name, m) in enumerate(regions):
    if not m.any():
        continue
    cross = to_crosssection(m)
    if cross is None or cross.is_empty():
        continue
    z0, z1 = i * LEVEL_MM, (i + 1) * LEVEL_MM
    mmesh = M.Manifold.extrude(cross, z1 - z0)
    mmesh = mmesh.translate([0.0, 0.0, z0])   # immutable API: reassign required
    mesh = manifold_to_trimesh(mmesh)
    # manifold3d guarantees a watertight, consistently-oriented manifold by construction
    assert mesh.is_watertight and mesh.is_volume, f"{name} not a watertight manifold"
    hexc = TRUE_HEX.get(name, TRUE_HEX["yellow"])
    solids.append(dict(name=name, hex=hexc, z0=z0, z1=z1, mesh=mesh))
    mesh.export(BASE / "meshes" / f"{i}_{name}.stl")
    print(f"level {i}: {name:<14} z={z0:.1f}-{z1:.1f}mm  "
          f"area={cross.area():.0f}mm^2  vol={mesh.volume:.0f}mm^3  {hexc}")

# ---------------------------------------------------------------- combined STL
comb = trimesh.util.concatenate([s["mesh"] for s in solids])
comb.export(OUT / "pyrrhic_victory.stl")

# ---------------------------------------------------------------- per-level STLs
# Loaded into Bambu Studio (or any slicer) at their absolute positions and
# assigned a filament each. Filenames are sort-ordered by level.
SUFFIX = {"black": "black-base", "brown": "brown-mid", "yellow": "yellow-inner",
          "border-yellow": "yellow-border-top"}
for i, s in enumerate(solids, start=1):
    s["stl_path"] = OUT / f"{i:02d}_{SUFFIX[s['name']]}.stl"
    s["mesh"].export(s["stl_path"])
    print(f"  stl  {s['stl_path'].name}")

# ---------------------------------------------------------------- OpenSCAD
# Vector, color-correct in any CAD viewer (OpenSCAD / Prusa / Cura / FreeCAD).
# Each level is a polygon() with holes (even_odd) extruded to its z-band.
LEVELS = [
    ("black",      full,                  "#000000"),  # full badge footprint (solid base, matches STL/3MF)
    ("brown",      brown | yellow,        TRUE_HEX["brown"]),
    ("yellow",     inner_y | border,      "#FFFF00"),
    ("border",     border,                "#FFFF00"),
]


def emit_scad():
    """One polygon() per band, with every outer + hole ring as a path.

    OpenSCAD polygon(points, paths) uses even-odd winding, so a ring set that
    alternates outer/hole (and any number of separate outer components) renders
    correctly (holes cut out, separate islands kept). This makes one band per
    level a single extrusion, which also fixes the 'only first polygon extruded'
    bug from multi-component bands.
    """
    out = [
        "// Pyrrhic Victory (Factorio tech icon) - 4-color terraced badge",
        "// Auto-generated by generate.py",
        f"// Footprint {FOOTPRINT_MM} x {FOOTPRINT_MM} mm, 8 mm tall ({NLEVELS} bands x {LEVEL_MM} mm), flat bottom at z=0.",
        "// Each band is one solid in its own z-range. In a multicolor slicer assign one",
        "// filament per solid (black / brown / yellow / yellow).",
        "",
    ]
    for i, (name, mask, hexc) in enumerate(LEVELS):
        if not mask.any():
            continue
        z0, h = i * LEVEL_MM, LEVEL_MM
        hx = hexc.lstrip("#").upper()
        cr, cg, cb = (int(hx[k:k + 2], 16) / 255 for k in (0, 2, 4))
        out.append(f"// level {i + 1}: {name}    z = {z0:.2f} .. {z0 + h:.2f}    color {hexc}")
        out.append(f"translate([0, 0, {z0:.4f}])")
        out.append(f"color([ {cr:.4f}, {cg:.4f}, {cb:.4f} ])")
        out.append(f"    linear_extrude(height = {h:.4f}, center = false)")
        # collect points and ring spans (a band may be a MultiPolygon with holes)
        all_pts, spans = [], []
        for (outer, holes) in mask_to_scad_polymask(mask):
            for ring in [outer] + holes:
                pts = [(float(x), float(y)) for (x, y) in ring]
                if len(pts) > 1 and pts[0] == pts[-1]:
                    pts = pts[:-1]              # drop the duplicated closing point
                start = len(all_pts)
                all_pts.extend(pts)
                spans.append((start, start + len(pts) - 1))   # inclusive [last, first-of-this-ring]
        pts_str = ", ".join(f"[{x:.4f}, {y:.4f}]" for (x, y) in all_pts)
        # explicit index list per ring: OpenSCAD paths syntax
        def ring(a, b):
            return "[" + ",".join(str(j) for j in range(a, b + 1)) + "]"
        paths_str = ", ".join(ring(a, b) for a, b in spans)
        out.append(f"        polygon(points = [\n            {pts_str}\n        ],")
        out.append(f"                 paths  = [\n            {paths_str}\n        ]);")
        out.append("")
    (OUT / "pyrrhic_victory.scad").write_text("\n".join(out) + "\n")
    print("  scad pyrrhic_victory.scad")

emit_scad()

# ---------------------------------------------------------------- 3MF writer
# Mirrors the exact structure of a Bambu-Studio-authored multicolor 3MF
# (/vault/3d-model/... LONGBU): each object in its own 3D/Objects/object_N.model,
# a composite 3dmodel.model holding <components> refs, a
# 3D/_rels/3dmodel.model.rels relationship file, and <item> elements in <build>.
import uuid as _uuid

TR = "1 0 0 0 1 0 0 0 1 0 0 0"
HDR = ('unit="millimeter" xml:lang="en-US" '
       'xmlns="http://schemas.microsoft.com/3dmanufacturing/core/2015/02" '
       'xmlns:BambuStudio="http://schemas.bambulab.com/package/2021" '
       'xmlns:p="http://schemas.microsoft.com/3dmanufacturing/production/2015/06" '
       'requiredextensions="p"')

for i, s in enumerate(solids):
    s["id"] = 2 + 2 * i          # composite ids in main model (even, like Bambu)
    s["uuid"] = str(_uuid.uuid4())
    s["cuuid"] = str(_uuid.uuid4())

def object_file_xml(s):
    verts = "".join(f'<vertex x="{x:.6f}" y="{y:.6f}" z="{z:.6f}" m:materialindex="0"/>'
                    for (x, y, z) in s["mesh"].vertices)
    tris = "".join(f'<triangle v1="{a}" v2="{b}" v3="{c}" m:materialindex="0"/>'
                   for (a, b, c) in s["mesh"].faces)
    return (f'<?xml version="1.0" encoding="UTF-8"?>\n'
            f'<model unit="millimeter" xml:lang="en-US" '
            f'xmlns="http://schemas.microsoft.com/3dmanufacturing/core/2015/02" '
            f'xmlns:m="http://schemas.microsoft.com/3dmanufacturing/material/2015/02" '
            f'xmlns:BambuStudio="http://schemas.bambulab.com/package/2021" '
            f'xmlns:p="http://schemas.microsoft.com/3dmanufacturing/production/2015/06" '
            f'requiredextensions="p">\n'
            ' <metadata name="BambuStudio:3mfVersion">1</metadata>\n'
            ' <resources>\n'
            f'  <object id="1" name="{s["name"]}" p:UUID="{s["uuid"]}" type="model">\n'
            f'   <m:basematerials><m:base id="0" name="{s["name"]}" color="{s["hex"]}FF"/></m:basematerials>\n'
            '   <mesh>\n'
            f'    <vertices>{verts}</vertices>\n'
            f'    <triangles>{tris}</triangles>\n'
            '   </mesh>\n'
            '  </object>\n'
            ' </resources>\n'
            ' <build/>\n'
            '</model>\n')

res_objs = []
for s in solids:
    res_objs.append(
        f'  <object id="{s["id"]}" p:UUID="{s["cuuid"]}" type="model">\n'
        '   <components>\n'
        f'    <component p:path="/3D/Objects/object_{s["id"]}.model" objectid="1" '
        f'p:UUID="{s["uuid"]}" transform="{TR}"/>\n'
        '   </components>\n'
        '  </object>')
items = "".join(
    f'<item objectid="{s["id"]}" p:UUID="{s["cuuid"]}" transform="{TR}" printable="1"/>'
    for s in solids)
main_xml = (f'<?xml version="1.0" encoding="UTF-8"?>\n'
            f'<model {HDR}>\n'
            ' <metadata name="Application">opencode-3dprint</metadata>\n'
            ' <metadata name="BambuStudio:3mfVersion">1</metadata>\n'
            ' <metadata name="Title">Pyrrhic Victory</metadata>\n'
            ' <metadata name="Description">Factorio Pyrrhic Victory tech icon, 4-level color badge</metadata>\n'
            ' <metadata name="Thumbnail_Middle">/Metadata/cover.png</metadata>\n'
            ' <metadata name="Thumbnail_Small">/Metadata/cover.png</metadata>\n'
            ' <resources>\n' + "\n".join(res_objs) + '\n  </resources>\n'
            ' <build>\n   ' + items + '\n  </build>\n'
            '</model>\n')

model_rels = ('<?xml version="1.0" encoding="UTF-8"?>\n'
              '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">\n'
              + "".join(
                  f' <Relationship Target="/3D/Objects/object_{s["id"]}.model" Id="rel-{i+1}" '
                  f'Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"/>\n'
                  for i, s in enumerate(solids))
              + '</Relationships>\n')

main_rels = ('<?xml version="1.0" encoding="UTF-8"?>\n'
             '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">\n'
             ' <Relationship Target="/3D/3dmodel.model" Id="rel-1" '
             'Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"/>\n'
             ' <Relationship Target="/Metadata/cover.png" Id="rel-2" '
             'Type="http://schemas.openxmlformats.org/package/2006/relationships/metadata/thumbnail"/>\n'
             ' <Relationship Target="/Metadata/cover.png" Id="rel-3" '
             'Type="http://schemas.bambulab.com/package/2021/cover-thumbnail-middle"/>\n'
             ' <Relationship Target="/Metadata/cover.png" Id="rel-4" '
             'Type="http://schemas.bambulab.com/package/2021/cover-thumbnail-small"/>\n'
             '</Relationships>\n')

content_types = ('<?xml version="1.0" encoding="UTF-8"?>\n'
                 '<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">\n'
                 ' <Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>\n'
                 ' <Default Extension="model" ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/>\n'
                 ' <Default Extension="png" ContentType="image/png"/>\n'
                 '</Types>\n')

cover = (OUT / "pyrrhic_src.png").read_bytes()

tmp3mf = OUT / "pyrrhic_victory.3mf.tmp"
with zipfile.ZipFile(tmp3mf, "w", zipfile.ZIP_DEFLATED) as z:
    z.writestr("[Content_Types].xml", content_types)
    z.writestr("_rels/.rels", main_rels)
    z.writestr("3D/3dmodel.model", main_xml)
    z.writestr("3D/_rels/3dmodel.model.rels", model_rels)
    z.writestr("Metadata/cover.png", cover)
    for s in solids:
        z.writestr(f"3D/Objects/object_{s['id']}.model", object_file_xml(s))
tmp3mf.replace(OUT / "pyrrhic_victory.3mf")

# ---------------------------------------------------------------- legend
legend = ["Pyrrhic Victory 3D print - filament legend", "=" * 52,
          f"footprint: {FOOTPRINT_MM} x {FOOTPRINT_MM} mm, total height: 8.0 mm",
          f"{len(solids)} stacked levels, ~{LEVEL_MM:.2f} mm each, flat bottom",
          "",
          f"  {'level':<6} {'layer region':<14} {'z(mm)':<11} {'color':<9} filament to assign",
          "-" * 70]
for i, s in enumerate(solids):
    legend.append(f"  {i+1:<6} {s['name']:<14} {s['z0']:.2f}-{s['z1']:.2f}{'':<6} {s['hex']:<9} closest to {s['hex']}")
legend += ["",
           "Layers are cumulative (each sits on the one below), so the print is",
           "supported. Terraces expose ~2mm of each color's band.",
           "Height order bottom->top: black < brown < inner-yellow < border.",
           "The outer RIM is the highest band; inner face/trophy detail is one",
           "step below it (both are yellow in the art).",
           "",
           "In Bambu Studio: enable multi-material (AMS/AMS Lite), then set a",
           "filament per object: black->black, brown->brown, yellow->yellow",
           "(inner-yellow and border-yellow both take your yellow spool)."]
(OUT / "legend.txt").write_text("\n".join(legend) + "\n")

print("\nwrote:")
for f in ("pyrrhic_victory.3mf", "pyrrhic_victory.stl", "legend.txt"):
    print("  ", OUT / f, (OUT / f).stat().st_size, "bytes")