"""AI 手机仓 — printable parts for Bambu 256x256 bed.
Every part is built directly in its PRINT orientation (Z up = build direction), units mm.
Outputs: STL per part + one generic 3MF per plate with parts laid out on the bed.
"""
import os, struct, zipfile, math
import numpy as np
from manifold3d import Manifold, CrossSection

OUT = os.path.join(os.path.dirname(__file__), 'out')
os.makedirs(OUT, exist_ok=True)
SEG = 48

# ---------------- primitives ----------------
def box(sx, sy, sz, x=0, y=0, z=0):
    return Manifold.cube([sx, sy, sz]).translate([x, y, z])

def cylz(r, h, x=0, y=0, z=0, seg=SEG):
    return Manifold.cylinder(h, r, r, seg).translate([x, y, z])

def cylx(r, h, x0, y, z, seg=SEG):
    """cylinder along +X from x0, length h"""
    return Manifold.cylinder(h, r, r, seg).rotate([0, 90, 0]).translate([x0, y, z])

def cyly(r, h, x, y0, z, seg=SEG):
    """cylinder along +Y from y0, length h"""
    return Manifold.cylinder(h, r, r, seg).rotate([-90, 0, 0]).translate([x, y0, z])

def union(*ms):
    out = ms[0]
    for m in ms[1:]:
        out = out + m
    return out

def hole_z(r, depth, x, y, ztop):
    """blind hole from the top face at ztop going down"""
    return cylz(r, depth + 0.01, x, y, ztop - depth)

# insert hole radii (heat-set): M2 3.2, M2.5 3.6, M3 4.0 ; through holes: M2 2.4, M2.5 2.9, M3 3.4
R_M2I, R_M25I, R_M3I = 1.6, 1.8, 2.0
R_M2T, R_M3T = 1.2, 1.7

# ---- global envelope: housing 123 x 230 x 172 (bottom shell 154 + lid 18); phone pocket 83 x 186 (phones <= 79 wide + case)
W, D, HS, T, TB, LIDH = 123, 230, 154, 2.4, 3, 18
HW = W / 2                      # 61.5
TRAY_W, TRAY_L, WALL = 93, 204, 5   # tray ±46.5, pocket ±41.5 (83 wide), pocket Y 8..194 (186 long)
LIP_W = TRAY_W + 20

# ---------------- 1. 底壳 bottom shell ----------------
def bottom_shell():
    H = HS
    body = box(W, D - 3, H, -HW, 3, 0)
    body -= box(W - 2 * T, D - 3 - T + 1, H + 10, -HW + T, 2, TB)          # cavity, open top & front
    body += box(1.2, D - 3 - T, 2, HW - 3.4, 3, H) + box(1.2, D - 3 - T, 2, -(HW - 2.2), 3, H)   # 止口 lip sides
    body += box(W - 6.8, 1.2, 2, -(HW - 3.4), D - T - 1.2, H)                                      # lip rear
    body -= box(12, 12, 4, HW - 14, D - 14, H) + box(12, 12, 4, -(HW - 2), D - 14, H)             # lip cut at lid tabs
    for sx in (1, -1):                                                                             # rail bosses + M3 holes
        body += box(2.6, D - 3 - T, 25, HW - 5 if sx > 0 else -(HW - 2.4), 3, 5)
        for yy in (30, 105, 180):
            body -= cylx(R_M3I, 4.6, HW - 5.1, yy, 18.5) if sx > 0 else cylx(R_M3I, 4.6, -(HW - 0.5), yy, 18.5)
    for sx in (1, -1):                                                                             # thumb-screw bosses
        bx = HW - 12 if sx > 0 else -(HW - 4)
        body += box(8, 8, 17, bx, D - T - 8, H - 17)
        body -= hole_z(R_M3I, 6, bx + 4, D - T - 4, H)
    body += box(2.6, 14, 18, -(HW - 2.4), 6, 30)                                                   # damper boss (left)
    for yy in (10, 16):
        body -= cylx(R_M2I, 4.6, -(HW - 0.5), yy, 38)
    body += box(8, 8, 11.5, -39, 4, TB)                                                            # spring anchor post
    body -= cyly(R_M3I, 6.5, -35, 4, 10)
    body += box(14, 14, 5, 13, 6, TB)                                                              # pogo block
    body -= hole_z(R_M2I, 4, 16, 13, TB + 5) + hole_z(R_M2I, 4, 24, 13, TB + 5)
    for xx in (12, -12):                                                                           # lock seat
        body += cylz(4, 4, xx, D - 13, TB)
        body -= hole_z(R_M3I, 3.5, xx, D - 13, TB + 4)
    for xx, yy in ((22, D - 19), (40, D - 19), (22, D - 8), (40, D - 8)):                           # hub standoffs
        body += cylz(2.5, 4, xx, yy, TB)
        body -= hole_z(R_M2I, 3.5, xx, yy, TB + 4)
    for sx in (1, -1):                                                                             # front-panel ears
        ex = HW - 12 if sx > 0 else -(HW - 4)
        for zz in (36, 96):
            body += box(8, 6, 8, ex, 3, zz)
            body -= cyly(1.45, 8, ex + 4, 2, zz + 4)
    body -= box(28, 6, 28, -46, D - 5, 41)                                                        # fan opening + grille
    for xx in (-39, -32, -25):
        body += box(2, T, 28, xx - 1, D - T, 41)
    for xx, zz in ((-44, 43), (-20, 43), (-44, 67), (-20, 67)):
        body -= cyly(0.9, 6, xx, D - 5, zz)
    for sx in (1, -1):                                                                             # side vents
        for i in range(5):
            body -= box(4, 60, 1.5, HW - 3.5 if sx > 0 else -(HW + 0.5), 100, 80 + i * 5)
    for sx in (1, -1):                                                                             # light-strip ribs
        rx = -(HW - 2.4) if sx < 0 else HW - 3.9
        body += box(1.5, 100, 1.2, rx, 40, 36) + box(1.5, 100, 1.2, rx, 40, 43.2)
    for xx, yy in ((-(HW - 10), 15), (HW - 10, 15), (-(HW - 10), D - 15), (HW - 10, D - 15)):      # foot recesses
        body -= cylz(6, 1.01, xx, yy, -0.01)
    return body

# ---------------- 2. 电子舱盖 lid (printed top face DOWN) ----------------
def lid():
    Hl, Tl, TT = LIDH, 2.0, 2.4
    body = box(W, D, Hl, -HW, 0, 0)
    body -= box(W - 2 * Tl, D - 2 * Tl, Hl, -(HW - Tl), 2, TT)
    for xx, yy in ((25, 127), (-25, 127), (25, 211), (-25, 211), (13, 89), (-13, 89)):            # PCB + camera bosses (M2.5)
        body += cylz(3.5, 6, xx, yy, TT)
        body -= hole_z(R_M25I, 5, xx, yy, TT + 6)
    body -= box(10, 6, 4, -5, D - 4, 5)                                                            # USB-C opening
    for sx in (1, -1):                                                                             # thumb-screw tabs
        tx = HW - 12 if sx > 0 else -(HW - 2)
        body += box(10, 10, 3, tx, D - Tl - 10, Hl - 3)
        body -= cylz(R_M3T, 4, tx + 5, D - Tl - 5, Hl - 3.5)
    for i in range(6):                                                                             # vents over the ESP32
        body -= box(24, 1.5, 4, 3, 124 + i * 4.5, -1)
    return body

# ---------------- 3. 前面板 front panel (outer face DOWN; local y = world Z) ----------------
def front_panel():
    body = box(W, HS, 3, -HW, 0, 0)
    body -= box(LIP_W + 1.2, 20.2, 5, -(LIP_W + 1.2) / 2, 14.4, -1)                                # tray opening
    body -= cylz(3.7, 5, 45, 50, -1)                                                               # button
    body -= box(30, 4, 5, 30, 76, -1)                                                              # light pipe slot
    for xx, yy in ((HW - 8, 40), (-(HW - 8), 40), (HW - 8, 100), (-(HW - 8), 100)):               # bosses w/ M2.5 inserts
        body += cylz(3.5, 6, xx, yy, 3)
        body -= hole_z(R_M25I, 5, xx, yy, 9)
    return body

# ---------------- 4. 托盘 tray (floor DOWN; local z = world Z - 15) ----------------
def tray():
    L, hw, pw = TRAY_L, TRAY_W / 2, TRAY_W / 2 - WALL      # hw 48, pw 43
    body = box(TRAY_W, L - 2, 3, -hw, 2, 0) + box(LIP_W, 2, 19, -LIP_W / 2, 0, 0)                # floor + front lip
    body += box(WALL, L - 2, 12, hw - WALL, 2, 0) + box(WALL, L - 2, 12, -hw, 2, 0)                # side walls
    body += box(2 * pw, 4, 12, -pw, L - 10, 0)                                                     # pocket rear wall (pocket Y 8..L-10)
    body += cylz(5, 8, -35, L - 9, 0)                                                              # spring axle boss
    body -= cylz(R_M3I, 8, -35, L - 9, -1)
    body += box(10, 4, 9, -5, L, 0)                                                                # latch tab + window
    body -= box(6, 6, 4, -3, L - 1, 4)
    for yy in (30, 105, 180):                                                                      # rail inner-member M3 holes
        body -= cylx(R_M3I, 4.6, hw - 0.1, yy, 3.5) + cylx(R_M3I, 4.6, -(hw + 4.5), yy, 3.5)
    for yy in (20, 180):                                                                           # rack holes on LEFT wall top
        body -= hole_z(R_M2I, 5, -(hw - 2.5), yy, 12)
    body -= box(5, 5, 0.6, pw, 67.5, 11.5) + box(5, 5, 0.6, -hw, 105.5, 11.5)                      # fiducial recesses
    body -= box(15.4, 25.4, 1.6, 12.3, 7.8, -0.01)                                                 # pad-board recess
    body -= box(3, 6, 5, 18.5, 8, -1)                                                              # lead slot
    body += box(2, 6, 6, -7.5, 2, 3) + box(2, 6, 6, 5.5, 2, 3)                                     # plug cradle
    body -= box(13.5, 2, 2.1, 5.5, 5, 1)                                                           # lead groove
    return body

# ---------------- 5. 齿条 rack, module 1, printed teeth UP ----------------
def rack():
    L, Wd, body_t = 180, 6, 3
    m = 1.0; pitch = math.pi * m; ha, hf = m, 1.25 * m; tan20 = math.tan(math.radians(20))
    half_p = pitch / 4                                   # half tooth thickness at pitch line
    profile = []
    n = int(L // pitch)
    x0 = (L - n * pitch) / 2
    for i in range(n):
        c = x0 + i * pitch + pitch / 2
        profile.append([
            [c - half_p - hf * tan20, body_t - 0.01], [c + half_p + hf * tan20, body_t - 0.01],
            [c + half_p - ha * tan20, body_t + hf + ha], [c - half_p + ha * tan20, body_t + hf + ha]])
    teeth = Manifold.extrude(CrossSection(profile), Wd)
    base = box(L, body_t + 0.001, Wd, 0, 0, 0)
    r = (base + teeth).rotate([90, 0, 0]).translate([0, Wd, 0])        # width now along Y, height along Z
    for xx in (10, 170):
        r -= cylz(R_M2T, 10, xx, Wd / 2, -1)
    return r

# ---------------- 6. 摄像头 U 形支架 (bottom plate DOWN) ----------------
def cam_bracket():
    body = box(32, 20, 2, -16, -10, 0)
    body -= box(14, 14, 5, -7, -7, -1)                            # lens barrel hole (13 x 13 barrel + clearance)
    for xx in (9.25, -9.25):                                      # ear bosses with M2 inserts (ears pitch 18)
        body += cylz(2.6, 4, xx, 0, 0)
        body -= hole_z(R_M2I, 3.5, xx, 0, 4)
    body -= box(14, 14, 5, -7, -7, -1)
    body += box(2, 20, 23, 14, -10, 0) + box(2, 20, 23, -16, -10, 0)   # legs
    body += box(32, 20, 3, -16, -10, 20)                           # top plate (28 mm bridge)
    for xx in (13, -13):                                          # adjustment slots along Y (M2.5)
        body -= box(2.9, 3.1, 5, xx - 1.45, -1.55, 19) + cylz(1.45, 5, xx, -1.55, 19) + cylz(1.45, 5, xx, 1.55, 19)
    return body

# ---------------- 7. 按键帽 / 8. 导光条 ----------------
def button_cap():
    return cylz(4.5, 1.5, 0, 0, 0) + cylz(3.4, 6, 0, 0, 1.4)

def light_pipe():
    return box(29.6, 3.6, 3, 0, 0, 0)

# ---------------- 9. 手机卡托 (TPU) — per size: 2 side strips + 1 rear block; plus 1 common front bar ----------------
PW, PL, IH = 83 - 0.4, 186 - 0.4, 4            # pocket 83 x 186 with 0.2 clearance per side
def insert_side(inner_w):
    w = (PW - inner_w) / 2
    return box(w, PL - 6, IH, 0, 0, 0)
def insert_rear(inner_l):
    d = PL - 3 - inner_l
    return box(PW, d, IH, 0, 0, 0)
def front_bar():
    return box(PW, 3, IH, 0, 0, 0) - box(16, 5, IH, (PW - 16) / 2, -1, 1.5)
SIZES = {'XS_68x140': (68, 140), 'S_72x152': (72, 152), 'M_76x162': (76, 162), 'L_79x168': (79, 168)}

PARTS = {
    'P01_底壳_bottom_shell': bottom_shell,
    'P02_电子舱盖_lid': lid,
    'P03_前面板_front_panel': front_panel,
    'P04_托盘_tray': tray,
    'P05_齿条_rack_m1': rack,
    'P06_摄像头支架_cam_bracket': cam_bracket,
    'P07_按键帽_button_cap': button_cap,
    'P08_导光条_light_pipe': light_pipe,
    'P09_卡托前条_front_bar': front_bar,
}
for _n, (_w, _l) in SIZES.items():
    PARTS['P10_卡托侧条_%s' % _n] = (lambda w=_w: insert_side(w))
    PARTS['P11_卡托后块_%s' % _n] = (lambda l=_l: insert_rear(l))

# ---------------- export helpers ----------------
def mesh_arrays(m):
    mesh = m.to_mesh()
    v = np.asarray(mesh.vert_properties, dtype=np.float32)[:, :3]
    t = np.asarray(mesh.tri_verts, dtype=np.uint32)
    return v, t

def write_stl(path, v, t):
    tri = v[t]                                   # (n,3,3)
    n = np.cross(tri[:, 1] - tri[:, 0], tri[:, 2] - tri[:, 0])
    n /= np.maximum(np.linalg.norm(n, axis=1, keepdims=True), 1e-12)
    data = np.zeros(len(t), dtype=[('n', '<f4', 3), ('v', '<f4', (3, 3)), ('a', '<u2')])
    data['n'] = n; data['v'] = tri
    with open(path, 'wb') as f:
        f.write(b'AI phone cabin part'.ljust(80, b'\0')); f.write(struct.pack('<I', len(t))); f.write(data.tobytes())

def write_3mf(path, items):
    """items: list of (name, v, t, (tx,ty)) — generic 3MF, one build with translations."""
    objs, builds = [], []
    for i, (name, v, t, (tx, ty)) in enumerate(items, start=1):
        verts = ''.join(f'<vertex x="{x:.4f}" y="{y:.4f}" z="{z:.4f}"/>' for x, y, z in v)
        tris = ''.join(f'<triangle v1="{a}" v2="{b}" v3="{c}"/>' for a, b, c in t)
        objs.append(f'<object id="{i}" name="{name}" type="model"><mesh><vertices>{verts}</vertices><triangles>{tris}</triangles></mesh></object>')
        builds.append(f'<item objectid="{i}" transform="1 0 0 0 1 0 0 0 1 {tx:.3f} {ty:.3f} 0"/>')
    model = ('<?xml version="1.0" encoding="UTF-8"?>'
             '<model unit="millimeter" xml:lang="en-US" xmlns="http://schemas.microsoft.com/3dmanufacturing/core/2015/02">'
             '<metadata name="Application">AI phone cabin builder</metadata>'
             f'<resources>{"".join(objs)}</resources><build>{"".join(builds)}</build></model>')
    with zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED) as z:
        z.writestr('[Content_Types].xml', '<?xml version="1.0" encoding="UTF-8"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/><Default Extension="model" ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/></Types>')
        z.writestr('_rels/.rels', '<?xml version="1.0" encoding="UTF-8"?><Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Target="/3D/3dmodel.model" Id="rel0" Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"/></Relationships>')
        z.writestr('3D/3dmodel.model', model)

if __name__ == '__main__':
    built = {}
    rows = []
    for name, fn in PARTS.items():
        m = fn()
        assert m.status().name == 'NoError', (name, m.status())
        v, t = mesh_arrays(m)
        bb = m.bounding_box()
        # normalise: min corner to (0,0,0) so plates can place by size
        v = v - np.array([bb[0], bb[1], bb[2]], dtype=np.float32)
        size = (bb[3] - bb[0], bb[4] - bb[1], bb[5] - bb[2])
        built[name] = (v, t, size, m.volume() / 1000.0)
        write_stl(os.path.join(OUT, name + '.stl'), v, t)
        rows.append((name, size, m.volume() / 1000.0, len(t), m.genus()))
    for r in rows:
        print(f'{r[0]:34s} {r[1][0]:6.1f} x {r[1][1]:6.1f} x {r[1][2]:5.1f} mm  vol {r[2]:7.1f} cm3  tris {r[3]:6d} genus {r[4]}')

    # -------- plates (Bambu 256 x 256), positions = min-corner of each part --------
    def item(name, tx, ty):
        v, t, size, vol = built[name]; return (name, v, t, (tx, ty))
    plates = {
        'Plate1_底壳+托盘+齿条+小件_PETG': [
            item('P01_底壳_bottom_shell', 3, 12),
            item('P04_托盘_tray', 131, 12),
            item('P05_齿条_rack_m1', 247, 12),       # rotated below (lies along Y)
            item('P06_摄像头支架_cam_bracket', 131, 226),
            item('P07_按键帽_button_cap', 170, 226),
            item('P08_导光条_light_pipe', 186, 226),
        ],
        'Plate2_舱盖+前面板_PETG': [
            item('P02_电子舱盖_lid', 3, 12),
            item('P03_前面板_front_panel', 130, 12),
        ],
        'Plate3_卡托_TPU': [],
    }
    # TPU strips: side strips along Y in columns, rear blocks + front bar stacked at the right
    x = 6
    for n, (w, l) in SIZES.items():
        sw = (PW - w) / 2
        for k in range(2):
            plates['Plate3_卡托_TPU'].append(item('P10_卡托侧条_%s' % n, x, 6)); x += sw + 5
    y = 6
    for n, (w, l) in SIZES.items():
        plates['Plate3_卡托_TPU'].append(item('P11_卡托后块_%s' % n, 150, y)); y += (PL - 3 - l) + 5
    plates['Plate3_卡托_TPU'].append(item('P09_卡托前条_front_bar', 150, y))
    # rotate rack 90° (lie along Y) and L insert 90°
    def rot90(it):
        name, v, t, (tx, ty) = it
        v2 = v.copy(); v2[:, 0], v2[:, 1] = -v[:, 1], v[:, 0]
        v2[:, 0] -= v2[:, 0].min()
        return (name, v2, t, (tx, ty))
    plates['Plate1_底壳+托盘+齿条+小件_PETG'][2] = rot90(plates['Plate1_底壳+托盘+齿条+小件_PETG'][2])
    for pname, items in plates.items():
        # bed check
        for name, v, t, (tx, ty) in items:
            mx, my = v[:, 0].max() + tx, v[:, 1].max() + ty
            assert mx <= 256 and my <= 256, (pname, name, mx, my)
        write_3mf(os.path.join(OUT, pname + '.3mf'), items)
        print('plate', pname, [(n, round(float(v[:, 0].max() + tx), 1), round(float(v[:, 1].max() + ty), 1)) for n, v, t, (tx, ty) in items])