#!/usr/bin/env python3
"""Per-vertex stretch with larger corner zone to protect holes."""
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_v(v0,new_w,new_d,new_h,orig_w,orig_d,orig_h,cx,cy,floor_h,lip_h):
    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
    
    mid_orig=orig_h-floor_h-lip_h
    mid_new=new_h-floor_h-lip_h
    z_scale=mid_new/mid_orig if mid_orig>0 else 1
    x_scale=(new_w-2*cx)/(orig_w-2*cx) if (orig_w-2*cx)>0 else 1
    y_scale=(new_d-2*cy)/(orig_d-2*cy) if (orig_d-2*cy)>0 else 1
    
    r=v0.copy()
    for idx in range(n*3):
        x,y,z=v0[idx]
        if z<z0+floor_h: z_new=z
        elif z>z1-lip_h: z_new=z+new_h-orig_h
        else: z_new=z0+floor_h+(z-z0-floor_h)*z_scale
        if x<x0+cx: x_new=x
        elif x>x1-cx: x_new=x+new_w-orig_w
        else: x_new=x0+cx+(x-x0-cx)*x_scale
        if y<y0+cy: y_new=y
        elif y>y1-cy: y_new=y+new_d-orig_d
        else: y_new=y0+cy+(y-y0-cy)*y_scale
        r[idx]=[x_new,y_new,z_new]
    
    print("  -> %.1fx%.1fx%.1f (x_scale=%.3f y_scale=%.3f z_scale=%.3f)" % (
        r[:,0].max()-r[:,0].min(), r[:,1].max()-r[:,1].min(), r[:,2].max()-r[:,2].min(),
        x_scale, y_scale, z_scale))
    return r

print("BASE (corner=17mm):")
v,n=load_stl(os.path.join(ORIG,"Box.stl"))
r=stretch_v(v,80,70,32,60,60,17,17,17,4,3)
save_stl(r,os.path.join(OUT,"base.stl"))

print("LID:")
v,n=load_stl(os.path.join(ORIG,"Lid.stl"))
r=stretch_v(v,80,70,18,60,60,10,17,17,3,3)
save_stl(r,os.path.join(OUT,"lid.stl"))

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(v,80,70,zh,60,60,zh,12,6,0,0)
    save_stl(r,os.path.join(OUT,dst))

print("Screw (original):")
shutil.copy2(os.path.join(ORIG,"Screw.stl"),os.path.join(OUT,"screw.stl"))
v,n=load_stl(os.path.join(OUT,"screw.stl"))
print("  %.1fx%.1fx%.1f (unchanged)" % (v[:,0].max()-v[:,0].min(),v[:,1].max()-v[:,1].min(),v[:,2].max()-v[:,2].min()))

print("\nDONE")