#!/usr/bin/env python3
"""Generate a closed cylinder STL for the drone body (92mm OD, 163mm height).
The cylinder is centered at the origin with axis along Z.
Units: metres (0.092m diameter, 0.163m height).
"""
import struct, math, sys

def write_binary_stl(filename, triangles):
    with open(filename, 'wb') as f:
        f.write(b'\0' * 80)  # header
        f.write(struct.pack('<I', len(triangles)))
        for normal, v1, v2, v3 in triangles:
            for val in normal + v1 + v2 + v3:
                f.write(struct.pack('<f', val))
            f.write(struct.pack('<H', 0))

def cylinder_mesh(radius, height, n_sides=60, n_z=10):
    triangles = []
    dtheta = 2 * math.pi / n_sides
    dz = height / n_z
    z_off = -height / 2  # center vertically

    # Side wall
    for i in range(n_sides):
        t0 = i * dtheta
        t1 = (i + 1) * dtheta
        c0, s0 = math.cos(t0), math.sin(t0)
        c1, s1 = math.cos(t1), math.sin(t1)
        nx = (c0 + c1) / 2
        ny = (s0 + s1) / 2
        ln = math.sqrt(nx*nx + ny*ny)
        nx, ny = nx/ln, ny/ln
        for j in range(n_z):
            z0 = z_off + j * dz
            z1 = z_off + (j + 1) * dz
            p00 = (radius*c0, radius*s0, z0)
            p10 = (radius*c1, radius*s1, z0)
            p01 = (radius*c0, radius*s0, z1)
            p11 = (radius*c1, radius*s1, z1)
            triangles.append(((nx, ny, 0.0), p00, p10, p11))
            triangles.append(((nx, ny, 0.0), p00, p11, p01))

    # Top cap (z = +height/2)
    z_top = height / 2
    cx, cy = 0.0, 0.0
    for i in range(n_sides):
        t0 = i * dtheta
        t1 = (i + 1) * dtheta
        p0 = (radius*math.cos(t0), radius*math.sin(t0), z_top)
        p1 = (radius*math.cos(t1), radius*math.sin(t1), z_top)
        triangles.append(((0, 0, 1.0), (cx, cy, z_top), p0, p1))

    # Bottom cap (z = -height/2)
    z_bot = -height / 2
    for i in range(n_sides):
        t0 = i * dtheta
        t1 = (i + 1) * dtheta
        p0 = (radius*math.cos(t0), radius*math.sin(t0), z_bot)
        p1 = (radius*math.cos(t1), radius*math.sin(t1), z_bot)
        triangles.append(((0, 0, -1.0), (cx, cy, z_bot), p1, p0))

    return triangles

if __name__ == '__main__':
    R = 0.046  # 92mm diameter / 2
    H = 0.163  # 163mm height
    tris = cylinder_mesh(R, H, n_sides=60, n_z=10)
    out = 'constant/triSurface/drone_body.stl'
    write_binary_stl(out, tris)
    print(f"Wrote {len(tris)} triangles to {out}")