#!/usr/bin/env python3
"""Deform STL: normalize to origin, stretch middle, keep corners."""
import struct, numpy as np, os

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_tri = struct.unpack("<I", f.read(4))[0]
        verts = np.zeros((n_tri * 3, 3), dtype=np.float64)
        for i in range(n_tri):
            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_tri

def write_stl(verts, n_tri, path):
    with open(path, "wb") as f:
        f.write(b"\x00" * 80)
        f.write(struct.pack("<I", n_tri))
        for i in range(n_tri):
            f.write(struct.pack("<fff", 0.0, 0.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_dim(vals, orig_len, new_len, corner):
    mn, mx = vals.min(), vals.max()
    if mx - mn < 0.01:
        return vals
    
    orig_min = mn
    orig_max = mx
    mid_min = orig_min + corner
    mid_max = orig_max - corner
    new_mid_min = orig_min + corner
    new_mid_max = orig_max - corner + (new_len - orig_len)
    scale = (new_mid_max - new_mid_min) / (mid_max - mid_min)
    
    result = np.empty_like(vals)
    mask_l = vals <= mid_min
    mask_m = (vals > mid_min) & (vals < mid_max)
    mask_r = vals >= mid_max
    
    result[mask_l] = vals[mask_l]
    result[mask_m] = new_mid_min + (vals[mask_m] - mid_min) * scale
    result[mask_r] = new_mid_max + (vals[mask_r] - mid_max)
    
    return result

def stretch_stl(verts, orig_size, new_size, corner_xy, corner_z):
    ox, oy, oz = orig_size
    nx, ny, nz = new_size
    cx, cy = corner_xy
    cz = corner_z
    v = verts.copy()
    v[:, 0] = stretch_dim(v[:, 0], ox, nx, cx)
    v[:, 1] = stretch_dim(v[:, 1], oy, ny, cy)
    v[:, 2] = stretch_dim(v[:, 2], oz, nz, cz)
    return v

# BOX
print("=== BASE ===")
v, n = load_stl(os.path.join(ORIG, "Box.stl"))
ox = v[:,0].max() - v[:,0].min()
oy = v[:,1].max() - v[:,1].min()
oz = v[:,2].max() - v[:,2].min()
print("  Orig: %.1fx%.1fx%.1f" % (ox, oy, oz))
v2 = stretch_stl(v, (60, 60, 17), (80, 70, 32), (12, 12), 3.0)
write_stl(v2, n, os.path.join(OUT, "base.stl"))
print("  New: %.1fx%.1fx%.1f" % (v2[:,0].max()-v2[:,0].min(), v2[:,1].max()-v2[:,1].min(), v2[:,2].max()-v2[:,2].min()))

# LID
print("\n=== LID ===")
v, n = load_stl(os.path.join(ORIG, "Lid.stl"))
ox = v[:,0].max() - v[:,0].min()
oy = v[:,1].max() - v[:,1].min()
oz = v[:,2].max() - v[:,2].min()
print("  Orig: %.1fx%.1fx%.1f" % (ox, oy, oz))
v2 = stretch_stl(v, (60, 60, 10), (80, 70, 18), (12, 12), 3.0)
write_stl(v2, n, os.path.join(OUT, "lid.stl"))
print("  New: %.1fx%.1fx%.1f" % (v2[:,0].max()-v2[:,0].min(), v2[:,1].max()-v2[:,1].min(), v2[:,2].max()-v2[:,2].min()))

# WINDOWS
print("\n=== 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))
    ox = v[:,0].max() - v[:,0].min()
    oy = v[:,1].max() - v[:,1].min()
    oz = v[:,2].max() - v[:,2].min()
    v2 = stretch_stl(v, (60, 60, oz), (80, 70, oz), (12, 12), 0)
    write_stl(v2, n, os.path.join(OUT, dst))
    print("  %s: %.1fx%.1fx%.1f -> %.1fx%.1fx%.1f" % (dst, ox, oy, oz, v2[:,0].max()-v2[:,0].min(), v2[:,1].max()-v2[:,1].min(), v2[:,2].max()-v2[:,2].min()))

# SCREW
print("\n=== SCREW ===")
v, n = load_stl(os.path.join(ORIG, "Screw.stl"))
ox = v[:,0].max() - v[:,0].min()
oy = v[:,1].max() - v[:,1].min()
oz = v[:,2].max() - v[:,2].min()
print("  Orig: %.1fx%.1fx%.1f" % (ox, oy, oz))
v2 = stretch_stl(v, (60, 60, oz), (80, 70, oz), (10, 10), 0)
write_stl(v2, n, os.path.join(OUT, "screw.stl"))
print("  New: %.1fx%.1fx%.1f" % (v2[:,0].max()-v2[:,0].min(), v2[:,1].max()-v2[:,1].min(), v2[:,2].max()-v2[:,2].min()))

print("\nDONE")