#!/usr/bin/env python3
"""Generate with exact Z stretch: keep 0-10mm, keep 10.2-17mm, stretch 10-10.2mm."""
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 info(p):
    with open(p, "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

def gen_base():
    """Stretch X=80, Y=70, Z only between 10 and 10.2mm."""
    v, n = load_stl(os.path.join(ORIG, "Box.stl"))
    
    x0,x1 = v[:,0].min(), v[:,0].max()
    y0,y1 = v[:,1].min(), v[:,1].max()
    z0,z1 = v[:,2].min(), v[:,2].max()
    ow, oh, od = 60, 17, 60
    nw, nh, nd = 80, 32, 70
    cx = cy = 20
    xs = (nw - 2*cx) / (ow - 2*cx)
    ys = (nd - 2*cy) / (od - 2*cy)
    
    floor_h = 8.5      # keep 0-10mm
    mid_hi = 9.0        # stretch 10-10.2mm
    orig_mid = mid_hi - floor_h  # 0.2mm
    new_mid = nh - floor_h - (oh - mid_hi)  # 32-10-(17-10.2) = 32-10-6.8 = 15.2mm
    zs = new_mid / orig_mid if orig_mid > 0 else 1
    
    print("Z: keep 0-%.1fmm, stretch %.1f-%.1fmm (%.1f->%.1fmm, scale=%.1f), keep %.1f-%.1fmm" % (
        floor_h, floor_h, mid_hi, orig_mid, new_mid, zs, mid_hi, oh))
    
    r = v.copy()
    for idx in range(n*3):
        x, y, z = v[idx]
        xn = x if x < x0+cx else (x+nw-ow if x > x1-cx else x0+cx + (x-x0-cx)*xs)
        yn = y if y < y0+cy else (y+nd-od if y > y1-cy else y0+cy + (y-y0-cy)*ys)
        if z <= z0 + floor_h:
            zn = z
        elif z >= z0 + mid_hi:
            zn = z + nh - oh
        else:
            zn = z0 + floor_h + (z - z0 - floor_h) * zs
        r[idx] = [xn, yn, zn]
    
    save_stl(r, os.path.join(OUT, "base.stl"))
    _, xd, yd, zd = info(os.path.join(OUT, "base.stl"))
    print("  base: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, len(r)//3))
    return r

def gen_lid():
    """Stretch X=80, Y=70, keep Z=10mm original."""
    v, n = load_stl(os.path.join(ORIG, "Lid.stl"))
    x0,x1 = v[:,0].min(), v[:,0].max()
    y0,y1 = v[:,1].min(), v[:,1].max()
    xs = (80-40)/(60-40)
    ys = (70-40)/(60-40)
    r = v.copy()
    for idx in range(n*3):
        x, y, z = v[idx]
        xn = x if x < x0+20 else (x+20 if x > x1-20 else x0+20 + (x-x0-20)*xs)
        yn = y if y < y0+20 else (y+10 if y > y1-20 else y0+20 + (y-y0-20)*ys)
        r[idx] = [xn, yn, z]
    save_stl(r, os.path.join(OUT, "lid.stl"))
    _, xd, yd, zd = info(os.path.join(OUT, "lid.stl"))
    print("  lid: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, len(r)//3))

def gen_windows():
    """Create custom window panels with Build123d."""
    from build123d import Box, Pos, export_stl
    def mkw(w, h, t, cut=None):
        p = Box(w, t, h)
        for s in [-1, 1]:
            p = p + Pos(0, s*(t/2-0.5), 0) * Box(w-2, 1.0, h+0.5)
        if cut:
            for cx, cy, cw, ch in cut:
                p = p - Pos(cx, 0, cy) * Box(cw, t+0.4, ch)
        return p
    export_stl(mkw(36, 23, 4.8), os.path.join(OUT, "window_smooth.stl"))
    export_stl(mkw(36, 23, 4.8, [(-9, 0, 17, 5.3), (9, 0, 9, 9)]), os.path.join(OUT, "window_keystone_dc.stl"))
    export_stl(mkw(26, 23, 4.8), os.path.join(OUT, "window_smooth_y.stl"))
    export_stl(mkw(26, 23, 4.8, [(0, 0, 9, 5)]), os.path.join(OUT, "window_usb.stl"))
    print("  windows OK")

print("BASE...")
gen_base()
print("LID...")
gen_lid()
print("WINDOWS...")
gen_windows()
print("SCREW (original)...")
shutil.copy2(os.path.join(ORIG, "Screw.stl"), os.path.join(OUT, "screw.stl"))
print("\nDONE")