import bpy
import bmesh
import math
from mathutils import Vector, Matrix


# ---------------------------------------------------------------------------
# Lookup tables
# ---------------------------------------------------------------------------

def _build_unit_circle(segments: int) -> list:
    """(cos, sin) sampled once per segment. Reused by all primitives."""
    return [
        (math.cos(2 * math.pi * i / segments),
         math.sin(2 * math.pi * i / segments))
        for i in range(segments)
    ]


def _build_latitude_stack(segments: int) -> list:
    """(sin_phi, cos_phi) for each latitude ring, poles included."""
    rings = max(segments // 2, 2)
    return [
        (math.sin(math.pi * r / rings),
         math.cos(math.pi * r / rings))
        for r in range(rings + 1)
    ]


class _Tables:
    """Precomputed trig tables for a given segment count."""
    def __init__(self, segments: int):
        self.segments  = segments
        self.rings     = max(segments // 2, 2)
        self.circle    = _build_unit_circle(segments)
        self.latitudes = _build_latitude_stack(segments)


# ---------------------------------------------------------------------------
# Canonical primitive builders  (Z-up, origin at centre)
# ---------------------------------------------------------------------------

def _sphere_positions(radius: float, tables: _Tables) -> list:
    out = []
    for sin_phi, cos_phi in tables.latitudes:
        for c, s in tables.circle:
            out.append(Vector((
                radius * sin_phi * c,
                radius * sin_phi * s,
                radius * cos_phi,
            )))
    return out


def _ring_positions(radius: float, tables: _Tables) -> list:
    """Flat ring in the XY plane — caller's matrix orients it."""
    return [Vector((c * radius, s * radius, 0.0)) for c, s in tables.circle]


def _apply(mat: Matrix, vecs: list) -> list:
    return [(mat @ v.to_4d()).to_3d() for v in vecs]


# ---------------------------------------------------------------------------
# BMesh helpers
# ---------------------------------------------------------------------------

def add_sphere(bm: bmesh.types.BMesh, mat: Matrix, radius: float, tables: _Tables):
    """Sphere built at origin/Z-up, then placed by `mat`."""
    segs  = tables.segments
    rings = tables.rings
    verts = [bm.verts.new(p) for p in _apply(mat, _sphere_positions(radius, tables))]
    for r in range(rings):
        for s in range(segs):
            sn = (s + 1) % segs
            try:
                bm.faces.new((
                    verts[ r      * segs + s ],
                    verts[ r      * segs + sn],
                    verts[(r + 1) * segs + sn],
                    verts[(r + 1) * segs + s ],
                ))
            except ValueError:
                pass


def add_capped_frustum(bm: bmesh.types.BMesh,
                        head_mat: Matrix, head_radius: float,
                        tail_mat: Matrix, tail_radius: float,
                        tables: _Tables):
    """
    Tapered cylinder. Rings are built in XY then placed by per-end matrices —
    no axis arithmetic needed.
    """
    head_pos = (head_mat @ Vector((0, 0, 0, 1))).to_3d()
    tail_pos = (tail_mat @ Vector((0, 0, 0, 1))).to_3d()
    if (tail_pos - head_pos).length < 1e-6:
        return

    segs = tables.segments
    hr = [bm.verts.new(p) for p in _apply(head_mat, _ring_positions(head_radius, tables))]
    tr = [bm.verts.new(p) for p in _apply(tail_mat, _ring_positions(tail_radius, tables))]

    # Side quads
    for i in range(segs):
        nxt = (i + 1) % segs
        try:
            bm.faces.new((hr[i], hr[nxt], tr[nxt], tr[i]))
        except ValueError:
            pass

    # Head cap
    hcv = bm.verts.new(head_pos)
    for i in range(segs):
        try:
            bm.faces.new((hcv, hr[(i + 1) % segs], hr[i]))
        except ValueError:
            pass

    # Tail cap
    tcv = bm.verts.new(tail_pos)
    for i in range(segs):
        try:
            bm.faces.new((tcv, tr[i], tr[(i + 1) % segs]))
        except ValueError:
            pass


# ---------------------------------------------------------------------------
# Bone → per-end matrices
# ---------------------------------------------------------------------------

def _end_matrices(world_mat: Matrix, bone_matrix: Matrix, bone_length: float):
    """
    Return (head_mat, tail_mat) — world-space matrices for each bone end.

    The bone matrix (armature-local) columns:
      col[0] = bone X axis
      col[1] = bone Y axis  ← along the bone
      col[2] = bone Z axis
      col[3] = head position

    We remap columns so our local Z = bone Y, meaning _ring_positions()
    (which lies in XY, normal = +Z) is automatically perpendicular to the
    bone — no cross products or acos required.
      our X = bone X
      our Y = bone Z
      our Z = bone Y  (ring normal → along bone)
    """
    bx = bone_matrix.col[0].to_3d()
    by = bone_matrix.col[1].to_3d()   # along bone
    bz = bone_matrix.col[2].to_3d()

    orient = Matrix((
        (bx.x, bz.x, by.x),
        (bx.y, bz.y, by.y),
        (bx.z, bz.z, by.z),
    ))

    head_local = bone_matrix.col[3].to_3d()
    tail_local  = head_local + by * bone_length

    def _make(pos_local):
        local_mat = orient.to_4x4()
        local_mat.col[3][:3] = pos_local
        return world_mat @ local_mat

    return _make(head_local), _make(tail_local)


# ---------------------------------------------------------------------------
# Operator
# ---------------------------------------------------------------------------

class ARMATURE_OT_build_envelope_mesh(bpy.types.Operator):
    bl_idname  = "armature.build_envelope_mesh"
    bl_label   = "Build Envelope Mesh"
    bl_description = (
        "Convert the selected armature into an envelope-style mesh. "
        "Object/Pose mode = posed bones | Edit mode = edit bones."
    )
    bl_options = {'REGISTER', 'UNDO'}

    segments: bpy.props.IntProperty(
        name="Segments",
        description="Sphere / cylinder resolution",
        default=16, min=4, max=64,
    )

    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return obj is not None and obj.type == 'ARMATURE'

    def execute(self, context):
        arm_obj   = context.active_object
        world_mat = arm_obj.matrix_world
        bm        = bmesh.new()
        tables    = _Tables(self.segments)

        def meshify(bone_matrix, bone_length, head_radius, tail_radius, draw_head=True):
            head_mat, tail_mat = _end_matrices(world_mat, bone_matrix, bone_length)
            if draw_head:
                add_sphere(bm, head_mat, max(head_radius, 0.001), tables)
            add_sphere(bm, tail_mat, max(tail_radius, 0.001), tables)
            add_capped_frustum(bm,
                               head_mat, max(head_radius, 0.001),
                               tail_mat, max(tail_radius, 0.001),
                               tables)

        if arm_obj.mode == 'EDIT':
            for bone in arm_obj.data.edit_bones:
                meshify(bone.matrix, bone.length,
                        bone.head_radius, bone.tail_radius,
                        bone.parent is None)
        else:
            for bone in arm_obj.pose.bones:
                meshify(bone.matrix, bone.bone.length,
                        bone.bone.head_radius, bone.bone.tail_radius,
                        bone.parent is None)

        mesh = bpy.data.meshes.new(arm_obj.name + "_envelope_mesh")
        bm.to_mesh(mesh)
        bm.free()
        mesh.update()

        result_obj = bpy.data.objects.new(arm_obj.name + "_envelope", mesh)
        context.collection.objects.link(result_obj)

        return {'FINISHED'}


# ---------------------------------------------------------------------------
# Panel
# ---------------------------------------------------------------------------

class VIEW3D_PT_armature_mesher(bpy.types.Panel):
    bl_label       = "Armature Mesher"
    bl_idname      = "VIEW3D_PT_armature_mesher"
    bl_space_type  = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category    = "Armature"

    def draw(self, context):
        layout = self.layout
        op = layout.operator(
            ARMATURE_OT_build_envelope_mesh.bl_idname,
            text="Build Envelope Mesh",
            icon='OUTLINER_OB_MESH',
        )
        op.segments = 16


# ---------------------------------------------------------------------------
# Registration
# ---------------------------------------------------------------------------

classes = (ARMATURE_OT_build_envelope_mesh, VIEW3D_PT_armature_mesher)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)

if __name__ == "__main__":
    register()