#!/usr/bin/env python3
"""Slice STL: keep floor/top, stretch only middle wall section in X/Y/Z."""
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.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")
    return n

def process_stl(src, dst, new_w, new_d, new_h, orig_w, orig_d, orig_h,
                floor_h, lip_h, corner_x, corner_y):
    v_orig, n_tri = load_stl(src)
    v = v_orig.copy()
    
    # Original bounds
    x_min, x_max = v[:,0].min(), v[:,0].max()
    y_min, y_max = v[:,1].min(), v[:,1].max()
    z_min, z_max = v[:,2].min(), v[:,2].max()
    
    orig_span = (x_max-x_min, y_max-y_min, z_max-z_min)
    
    # Per-triangle Z classification
    tri_min = v.reshape(-1, 3, 3)[:, :, 2].min(axis=1)
    tri_max = v.reshape(-1, 3, 3)[:, :, 2].max(axis=1)
    
    floor_end = z_min + floor_h
    top_start = z_max - lip_h
    
    is_floor = tri_max <= floor_end + 0.5
    is_top = tri_min >= top_start - 0.5
    is_mid = ~is_floor & ~is_top
    
    # Scales
    mid_orig = orig_span[2] - floor_h - lip_h
    mid_new = new_h - floor_h - lip_h
    z_scale = mid_new / mid_orig if mid_orig > 0 else 1.0
    
    x_mid_orig = orig_w - 2*corner_x
    y_mid_orig = orig_d - 2*corner_y
    x_scale = (new_w - 2*corner_x) / x_mid_orig if x_mid_orig > 0 else 1.0
    y_scale = (new_d - 2*corner_y) / y_mid_orig if y_mid_orig > 0 else 1.0
    
    result = v.copy()
    
    for i in range(n_tri):
        if is_top[i]:
            # Top rim: shift up to new total height
            for j in range(3):
                idx = i*3 + j
                result[idx, 2] += new_h - orig_h
        elif is_mid[i]:
            for j in range(3):
                idx = i*3 + j
                x, y, z = v[idx]
                
                # X stretch in middle only
                if x < x_min + corner_x:
                    pass  # left corner
                elif x > x_max - corner_x:
                    result[idx, 0] = x + (new_w - orig_w)  # shift right corner outward
                else:
                    result[idx, 0] = x_min + corner_x + (x - x_min - corner_x) * x_scale
                
                # Y stretch
                if y > y_min + corner_y and y < y_max - corner_y:
                    result[idx, 1] = y_min + corner_y + (y - y_min - corner_y) * y_scale
                elif y >= y_max - corner_y:
                    result[idx, 1] = y + (new_d - orig_d)
                
                # Z stretch in middle
                if z > floor_end + z_min:
                    result[idx, 2] = z_min + floor_h + (z - z_min - floor_h) * z_scale
    
    # Verify
    print("  %s: %dx%dx%d -> %dx%dx%d (%d tri, floor=%d, mid=%d, top=%d)" % (
        dst, orig_span[0], orig_span[1], orig_span[2],
        result[:,0].max()-result[:,0].min(), 
        result[:,1].max()-result[:,1].min(),
        result[:,2].max()-result[:,2].min(),
        n_tri, is_floor.sum(), is_mid.sum(), is_top.sum()))
    
    save_stl(result, os.path.join(OUT, dst))

# ==== BASE ====
print("=== BASE ===")
process_stl(
    src=os.path.join(ORIG, "Box.stl"),
    dst="base.stl",
    new_w=80, new_d=70, new_h=32,
    orig_w=60, orig_d=60, orig_h=17,
    floor_h=4.0, lip_h=3.0,
    corner_x=12.0, corner_y=12.0
)

# ==== LID ====
print("\n=== LID ===")
process_stl(
    src=os.path.join(ORIG, "Lid.stl"),
    dst="lid.stl",
    new_w=80, new_d=70, new_h=18,
    orig_w=60, orig_d=60, orig_h=10,
    floor_h=3.0, lip_h=3.0,
    corner_x=12.0, corner_y=12.0
)

# ==== WINDOWS: only X/Y stretch ====
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))
    z = v[:,2].max() - v[:,2].min()
    # Windows: stretch X/Y, keep Z
    process_stl(
        src=os.path.join(ORIG, src),
        dst=dst,
        new_w=80, new_d=70, new_h=z,
        orig_w=60, orig_d=60, orig_h=z,
        floor_h=0, lip_h=0,
        corner_x=8.0, corner_y=4.0
    )

# ==== SCREW ====
print("\n=== SCREW ===")
import shutil
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")