#!/usr/bin/env python3
"""Stretch original window STLs: keep edges, stretch 0.1mm at center."""
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_mid(v, ow, od, nw, nd, trans):
    """Keep edges, stretch a narrow trans zone at center of X and Y."""
    x0, x1 = v[:,0].min(), v[:,0].max()
    y0, y1 = v[:,1].min(), v[:,1].max()
    n = len(v)//3
    x_mid = (x0+x1)/2
    y_mid = (y0+y1)/2
    x_lo = x_mid - trans/2
    x_hi = x_mid + trans/2
    y_lo = y_mid - trans/2
    y_hi = y_mid + trans/2
    x_scale = (nw - ow + trans) / trans if trans > 0 else 1
    y_scale = (nd - od + trans) / trans if trans > 0 else 1
    print("  X: orig=%.1f->%.1f (scale=%.1f) Y: orig=%.1f->%.1f (scale=%.1f)" % (ow, nw, x_scale, od, nd, y_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
        if y > y_lo and y < y_hi:
            yn = y_lo + (y - y_lo) * y_scale
        elif y >= y_hi:
            yn = y + nd - od
        else:
            yn = y
        r[idx] = [xn, yn, z]
    return r

# Window definitions: (source, dest, target_width, target_height)
windows = [
    ("Window 20x60mm - liso.stl", "window_smooth.stl", 36, 23),
    ("Window 20x60mm.stl", "window_keystone_dc.stl", 36, 23),
    ("Window - Flex Gland.stl", "window_usb.stl", 26, 23),
]

for src, dst, nw, nd in windows:
    v, n = load_stl(os.path.join(ORIG, src))
    ow = v[:,0].max() - v[:,0].min()
    od = v[:,1].max() - v[:,1].min()
    print("%s -> %s: %dx%d -> %dx%d" % (src, dst, ow, od, nw, nd))
    r = stretch_mid(v, ow, od, nw, nd, 0.1)
    save_stl(r, os.path.join(OUT, dst))

# Y-axis smooth window (26mm wide)
v, n = load_stl(os.path.join(ORIG, "Window 20x60mm - liso.stl"))
ow = v[:,0].max() - v[:,0].min()
od = v[:,1].max() - v[:,1].min()
print("window_smooth_y.stl: %dx%d -> 26x%d" % (ow, od, 23))
r = stretch_mid(v, ow, od, 26, 23, 0.1)
save_stl(r, os.path.join(OUT, "window_smooth_y.stl"))

print("\nVerification:")
for f in sorted(os.listdir(OUT)):
    if f.endswith(".stl") and f not in ("screw.stl",):
        n, x, y, z = info(os.path.join(OUT, f))
        print("  %s: %.1fx%.1fx%.1f (%d tri)" % (f, x, y, z, n))