#!/usr/bin/env python3
"""Stretch windows: only X (length) with 0.1mm transition, Z (height) PRESERVED."""
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 = 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 stretch_x_only(v, ow, nw, trans):
    """Stretch ONLY X (width), keep Y and Z untouched. 0.1mm transition at X center."""
    x0, x1 = v[:,0].min(), v[:,0].max()
    n = len(v)//3
    x_mid = (x0 + x1) / 2
    x_lo = x_mid - trans/2
    x_hi = x_mid + trans/2
    x_scale = (nw - ow + trans) / trans if trans > 0 else 1
    print("  X: %.1f -> %.1f (scale=%.1f), Y/Z preserved" % (ow, nw, x_scale))
    r = v.copy()
    for idx in range(n*3):
        x, y, z = v[idx]
        if x > x_lo and x < x_hi:
            xn = x_lo + (x - x_lo) * x_scale
        elif x >= x_hi:
            xn = x + nw - ow
        else:
            xn = x
        r[idx] = [xn, y, z]
    return r

# === Window_keystone_dc: stretch X only, preserve Y and Z ===
# Source: "Window 20x60mm.stl" (has keystone + DC cutout)
# Target X: 36mm, Y/Z: original
print("window_keystone_dc.stl (X stretch only, Y/Z preserved)...")
v, n = load_stl(os.path.join(ORIG, "Window 20x60mm.stl"))
ow = v[:,0].max() - v[:,0].min()
r = stretch_x_only(v, ow, 36, 0.1)
save_stl(r, os.path.join(OUT, "window_keystone_dc.stl"))
nt, xd, yd, zd = info(os.path.join(OUT, "window_keystone_dc.stl"))
print("  result: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, nt))

# === Window_smooth_y: stretch X only, preserve Y and Z ===
# Source: "Window 20x60mm - liso.stl" (smooth)
# Target X: 26mm, Y/Z: original
print("window_smooth_y.stl (X stretch only, Y/Z preserved)...")
v, n = load_stl(os.path.join(ORIG, "Window 20x60mm - liso.stl"))
ow = v[:,0].max() - v[:,0].min()
r = stretch_x_only(v, ow, 26, 0.1)
save_stl(r, os.path.join(OUT, "window_smooth_y.stl"))
nt, xd, yd, zd = info(os.path.join(OUT, "window_smooth_y.stl"))
print("  result: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, nt))

# Also regenerate window_smooth and window_usb with same approach (X only, Z preserved)
print("window_smooth.stl (X stretch only, Y/Z preserved)...")
v, n = load_stl(os.path.join(ORIG, "Window 20x60mm - liso.stl"))
ow = v[:,0].max() - v[:,0].min()
r = stretch_x_only(v, ow, 36, 0.1)
save_stl(r, os.path.join(OUT, "window_smooth.stl"))
nt, xd, yd, zd = info(os.path.join(OUT, "window_smooth.stl"))
print("  result: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, nt))

print("window_usb.stl (X stretch only, Y/Z preserved)...")
v, n = load_stl(os.path.join(ORIG, "Window - Flex Gland.stl"))
ow = v[:,0].max() - v[:,0].min()
r = stretch_x_only(v, ow, 26, 0.1)
save_stl(r, os.path.join(OUT, "window_usb.stl"))
nt, xd, yd, zd = info(os.path.join(OUT, "window_usb.stl"))
print("  result: %.1fx%.1fx%.1f (%d tri)" % (xd, yd, zd, nt))

print("\nDONE")