#!/usr/bin/env python3
"""Generate all files: base, lid, windows, screw with corner=20mm."""
import struct, numpy as np, os, shutil

ORIG = "/app/projects/esp32-box/original"
OUT = "/app/projects/esp32-box/exports"
os.makedirs(OUT, exist_ok=True)

def load_stl(path):
    with open(path, "rb") as f:
        f.read(80)
        n = struct.unpack("<I", f.read(4))[0]
        verts = np.zeros((n*3, 3), dtype=np.float64)
        for i in range(n):
            f.read(12)
            for j in range(3):
                verts[i*3+j] = np.frombuffer(f.read(12), dtype=np.float32)
            f.read(2)
    return verts, n

def save_stl(verts, path):
    n = len(verts)//3
    with open(path, "wb") as f:
        f.write(b"\x00"*80)
        f.write(struct.pack("<I", n))
        for i in range(n):
            f.write(struct.pack("<fff", 0, 0, 0))
            for j in range(3):
                v = verts[i*3+j]
                f.write(struct.pack("<fff", float(v[0]), float(v[1]), float(v[2])))
            f.write(b"\x00\x00")

def stretch(v0, cx, cy, floor_h, lip_h, ow, od, oh, nw, nd, nh):
    x0, x1 = v0[:,0].min(), v0[:,0].max()
    y0, y1 = v0[:,1].min(), v0[:,1].max()
    z0, z1 = v0[:,2].min(), v0[:,2].max()
    n = len(v0)//3
    mo = oh - floor_h - lip_h
    mn = nh - floor_h - lip_h
    zs = mn/mo if mo > 0 else 1
    xs = (nw-2*cx)/(ow-2*cx) if (ow-2*cx) > 0 else 1
    ys = (nd-2*cy)/(od-2*cy) if (od-2*cy) > 0 else 1
    r = v0.copy()
    for idx in range(n*3):
        x, y, z = v0[idx]
        if z < z0 + floor_h:
            zn = z
        elif z > z1 - lip_h:
            zn = z + nh - oh
        else:
            zn = z0 + floor_h + (z - z0 - floor_h) * zs
        if x < x0 + cx:
            xn = x
        elif x > x1 - cx:
            xn = x + nw - ow
        else:
            xn = x0 + cx + (x - x0 - cx) * xs
        if y < y0 + cy:
            yn = y
        elif y > y1 - cy:
            yn = y + nd - od
        else:
            yn = y0 + cy + (y - y0 - cy) * ys
        r[idx] = [xn, yn, zn]
    return r

def info(path):
    with open(path, "rb") as f:
        f.read(80)
        n = struct.unpack("<I", f.read(4))[0]
        x1=y1=z1=9e9; x2=y2=z2=-9e9
        for i in range(n):
            f.read(12)
            for j in range(3):
                x,y,z = struct.unpack("<fff", f.read(12))
                x1,x2 = min(x1,x), max(x2,x)
                y1,y2 = min(y1,y), max(y2,y)
                z1,z2 = min(z1,z), max(z2,z)
            f.read(2)
    return n, x2-x1, y2-y1, z2-z1

# Clean old files
for f in os.listdir(OUT):
    if f.endswith(".stl"):
        os.remove(os.path.join(OUT, f))

print("BASE corner=20mm...")
v, n = load_stl(os.path.join(ORIG, "Box.stl"))
r = stretch(v, 20, 20, 4, 3, 60, 60, 17, 80, 70, 32)
save_stl(r, os.path.join(OUT, "base.stl"))
nw, xd, yd, zd = info(os.path.join(OUT, "base.stl"))
print("  base.stl: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, nw))

print("LID...")
v, n = load_stl(os.path.join(ORIG, "Lid.stl"))
r = stretch(v, 20, 20, 3, 3, 60, 60, 10, 80, 70, 18)
save_stl(r, os.path.join(OUT, "lid.stl"))
nw, xd, yd, zd = info(os.path.join(OUT, "lid.stl"))
print("  lid.stl: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, nw))

print("WINDOWS...")
for src, dst in [
    ("Window 20x60mm - liso.stl", "window_smooth.stl"),
    ("Window 20x60mm.stl", "window_keystone_dc.stl"),
    ("Window - Flex Gland.stl", "window_usb.stl"),
]:
    v, n = load_stl(os.path.join(ORIG, src))
    zh = v[:,2].max() - v[:,2].min()
    r = stretch(v, 14, 8, 0, 0, 60, 60, zh, 80, 70, zh)
    save_stl(r, os.path.join(OUT, dst))
    nw, xd, yd, zd = info(os.path.join(OUT, dst))
    print("  %s: %.1fx%.1fx%.1f (%d tri)" % (dst, xd, yd, zd, nw))

print("SCREW (original)...")
shutil.copy2(os.path.join(ORIG, "Screw.stl"), os.path.join(OUT, "screw.stl"))
nw, xd, yd, zd = info(os.path.join(OUT, "screw.stl"))
print("  screw.stl: %.1fx%.1fx%.1f (%d tri, ORIGINAL)" % (xd, yd, zd, nw))

print("\nDONE - 6 files in %s" % OUT)