#!/usr/bin/env python3
"""Deform box STL only. Screw and windows get custom treatment."""
import struct, numpy as np, os
import 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_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
    mid_min = mn + corner
    mid_max = mx - corner
    new_mid_min = mn + corner
    new_mid_max = mx - corner + (new_len - orig_len)
    scale = (new_mid_max - new_mid_min) / (mid_max - mid_min)
    
    result = np.empty_like(vals)
    result[vals <= mid_min] = vals[vals <= mid_min]
    mid_mask = (vals > mid_min) & (vals < mid_max)
    result[mid_mask] = new_mid_min + (vals[mid_mask] - mid_min) * scale
    result[vals >= mid_max] = new_mid_max + (vals[vals >= mid_max] - mid_max)
    return result

# === BOX: 60x60x17 -> 80x70x32 ===
print("=== BASE ===")
v, n = load_stl(os.path.join(ORIG, "Box.stl"))
v[:, 0] = stretch_dim(v[:, 0], 60, 80, 12)
v[:, 1] = stretch_dim(v[:, 1], 60, 70, 12)
v[:, 2] = stretch_dim(v[:, 2], 17, 32, 3)
write_stl(v, n, os.path.join(OUT, "base.stl"))
print("  base: %.0fx%.0fx%.0fmm" % (v[:,0].max()-v[:,0].min(), v[:,1].max()-v[:,1].min(), v[:,2].max()-v[:,2].min()))

# === LID: 60x60x10 -> 80x70x18 ===
print("=== LID ===")
v, n = load_stl(os.path.join(ORIG, "Lid.stl"))
v[:, 0] = stretch_dim(v[:, 0], 60, 80, 12)
v[:, 1] = stretch_dim(v[:, 1], 60, 70, 12)
v[:, 2] = stretch_dim(v[:, 2], 10, 18, 3)
write_stl(v, n, os.path.join(OUT, "lid.stl"))
print("  lid: %.0fx%.0fx%.0fmm" % (v[:,0].max()-v[:,0].min(), v[:,1].max()-v[:,1].min(), v[:,2].max()-v[:,2].min()))

# === WINDOWS: scale width only ===
print("=== WINDOWS ===")
for src, dst, cx, cy in [
    ("Window 20x60mm - liso.stl", "window_smooth.stl", 8, 4),
    ("Window 20x60mm.stl", "window_keystone_dc.stl", 8, 4),
    ("Window - Flex Gland.stl", "window_usb.stl", 8, 4),
]:
    v, n = load_stl(os.path.join(ORIG, src))
    oz = v[:,2].max() - v[:,2].min()
    v[:, 0] = stretch_dim(v[:, 0], 60, 80, cx)
    v[:, 1] = stretch_dim(v[:, 1], 60, 70, cy)
    # Keep Z same
    write_stl(v, n, os.path.join(OUT, dst))
    print("  %s: orig Z=%.1f -> %.0fx%.0fx%.1fmm" % (dst, oz, v[:,0].max()-v[:,0].min(), v[:,1].max()-v[:,1].min(), v[:,2].max()-v[:,2].min()))

# === SCREW: 100% ORIGINAL, NO CHANGE ===
print("=== SCREW ===")
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("  screw: %.1fx%.1fx%.1fmm (original)" % (v[:,0].max()-v[:,0].min(), v[:,1].max()-v[:,1].min(), v[:,2].max()-v[:,2].min()))

print("\nDONE")