import bpy
import bmesh
import math
from mathutils import Vector, Euler, Quaternion

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
try:
    if bpy.context.object is not None and bpy.context.object.mode != 'OBJECT':
        bpy.ops.object.mode_set(mode='OBJECT')
except Exception:
    pass

for _o in list(bpy.data.objects):
    if _o.type in {'MESH', 'CURVE', 'EMPTY'}:
        bpy.data.objects.remove(_o, do_unlink=True)

PARTS = []


def active():
    o = bpy.context.view_layer.objects.active
    if o is None:
        sel = bpy.context.selected_objects
        o = sel[-1] if sel else None
    if o is None:
        raise RuntimeError("No active object after primitive creation")
    return o


def make_mat(name, rgb, rough=0.5, metal=0.0, coat=0.0):
    m = bpy.data.materials.new(name)
    try:
        m.diffuse_color = (rgb[0], rgb[1], rgb[2], 1.0)
    except Exception:
        pass
    try:
        m.use_nodes = True
        bsdf = m.node_tree.nodes.get("Principled BSDF")
        if bsdf is not None:
            try:
                bsdf.inputs["Base Color"].default_value = (rgb[0], rgb[1], rgb[2], 1.0)
            except Exception:
                pass
            try:
                bsdf.inputs["Roughness"].default_value = float(rough)
            except Exception:
                pass
            try:
                bsdf.inputs["Metallic"].default_value = float(metal)
            except Exception:
                pass
            if coat > 0:
                try:
                    bsdf.inputs["Coat Weight"].default_value = float(coat)
                except Exception:
                    pass
    except Exception:
        pass
    return m


M = {
    'grey': make_mat("Flank_Vermiculated_Grey", (0.55, 0.55, 0.52), 0.8),
    'chestnut': make_mat("Breast_Chestnut", (0.30, 0.11, 0.05), 0.7),
    'black': make_mat("Rump_Black", (0.02, 0.02, 0.025), 0.5),
    'belly': make_mat("Belly_PaleGrey", (0.76, 0.75, 0.71), 0.8),
    'back': make_mat("Mantle_GreyBrown", (0.36, 0.31, 0.26), 0.8),
    'green': make_mat("Head_Iridescent_Green", (0.02, 0.22, 0.08), 0.3, 0.4, 0.5),
    'white': make_mat("Collar_White", (0.92, 0.92, 0.90), 0.7),
    'bill': make_mat("Bill_Yellow", (0.78, 0.68, 0.16), 0.45),
    'nail': make_mat("Bill_Nail_Dark", (0.06, 0.06, 0.04), 0.4),
    'eye': make_mat("Eye_DarkBrown", (0.08, 0.04, 0.02), 0.05),
    'leg': make_mat("Leg_Orange", (0.95, 0.42, 0.06), 0.55),
    'wing': make_mat("Wing_Coverts_GreyBrown", (0.42, 0.38, 0.33), 0.8),
    'primary': make_mat("Primaries_DarkBrown", (0.27, 0.24, 0.21), 0.75),
    'spec': make_mat("Speculum_Blue", (0.14, 0.10, 0.58), 0.25, 0.5, 0.4),
    'tail': make_mat("Rectrices_WhiteGrey", (0.82, 0.82, 0.80), 0.8),
    'claw': make_mat("Claw_Dark", (0.10, 0.08, 0.06), 0.4),
}


def finish(o, name, mats, smooth=True, assign=None):
    o.name = name
    for m in mats:
        o.data.materials.append(m)
    if assign is not None:
        for p in o.data.polygons:
            p.material_index = assign(p.center)
    if smooth:
        for p in o.data.polygons:
            p.use_smooth = True
    PARTS.append(o)
    return o


def set_rot(o, rot):
    if isinstance(rot, Quaternion):
        o.rotation_mode = 'QUATERNION'
        o.rotation_quaternion = rot
    else:
        o.rotation_mode = 'XYZ'
        o.rotation_euler = Euler(rot, 'XYZ')


def ellipsoid(name, loc, scale, mats, rot=(0, 0, 0), deform=None, assign=None, seg=32, rings=16):
    bpy.ops.mesh.primitive_uv_sphere_add(segments=seg, ring_count=rings, radius=1.0, location=(0, 0, 0))
    o = active()
    for v in o.data.vertices:
        x, y, z = v.co
        if deform is not None:
            x, y, z = deform(x, y, z)
        v.co = Vector((x * scale[0], y * scale[1], z * scale[2]))
    o.data.update()
    set_rot(o, rot)
    o.location = Vector(loc)
    if not isinstance(mats, (list, tuple)):
        mats = [mats]
    return finish(o, name, mats, assign=assign)


def sphere(name, loc, r, mat):
    return ellipsoid(name, loc, (r, r, r), mat, seg=16, rings=8)


def limb(name, p1, p2, r1, r2, mat, verts=16):
    p1, p2 = Vector(p1), Vector(p2)
    d = p2 - p1
    L = max(d.length, 1e-5)
    bpy.ops.mesh.primitive_cone_add(vertices=verts, radius1=r1, radius2=r2, depth=L,
                                    end_fill_type='NGON', location=(p1 + p2) / 2)
    o = active()
    o.rotation_mode = 'QUATERNION'
    o.rotation_quaternion = Vector((0, 0, 1)).rotation_difference(d.normalized())
    return finish(o, name, [mat])


def axis_of(rot):
    return Euler(rot, 'XYZ').to_matrix() @ Vector((1, 0, 0))


# ---------------------------------------------------------------------------
# Body (torso) — drake plumage zones by face position
# ---------------------------------------------------------------------------
def body_def(x, y, z):
    if x < 0:
        f = 1.0 - 0.45 * ((-x) ** 1.6)
        y *= f
        z *= f
        z += 0.25 * (x * x)
    elif z < 0:
        z *= 1.0 + 0.1 * x
    if z > 0:
        z *= 0.85
    else:
        z *= 0.92
    return x, y, z


def body_assign(c):
    if c.x < -0.14:
        return 2
    if c.x > 0.10 and c.z > -0.07:
        return 1
    if c.x > 0.13:
        return 1
    if c.z < -0.055:
        return 3
    if c.z > 0.05 and c.x > -0.06:
        return 4
    return 0


body = ellipsoid("Body_Torso", (0, 0, 0.17), (0.20, 0.11, 0.095),
                 [M['grey'], M['chestnut'], M['black'], M['belly'], M['back']],
                 deform=body_def, assign=body_assign, seg=48, rings=24)

# Under-tail and upper-tail coverts (black)
ellipsoid("UpperTailCoverts", (-0.195, 0, 0.228), (0.055, 0.045, 0.028), M['black'], rot=(0, math.radians(12), 0))
ellipsoid("UnderTailCoverts", (-0.185, 0, 0.158), (0.05, 0.05, 0.035), M['black'])

# ---------------------------------------------------------------------------
# Neck (chestnut base, white collar, green upper)
# ---------------------------------------------------------------------------
neck_p1 = Vector((0.13, 0, 0.21))
neck_p2 = Vector((0.235, 0, 0.36))
neck_d = neck_p2 - neck_p1
neck_half = neck_d.length / 2.0
neck_q = Vector((0, 0, 1)).rotation_difference(neck_d.normalized())


def neck_def(x, y, z):
    f = 1.1 - 0.28 * (z + 1) / 2
    return x * f, y * f, z


def neck_assign(c):
    t = (c.z + neck_half) / (2 * neck_half)
    if t < 0.28:
        return 0
    if t < 0.36:
        return 1
    return 2


ellipsoid("Neck", (neck_p1 + neck_p2) / 2, (0.036, 0.032, neck_half),
          [M['chestnut'], M['white'], M['green']], rot=neck_q,
          deform=neck_def, assign=neck_assign, seg=32, rings=24)

# ---------------------------------------------------------------------------
# Head
# ---------------------------------------------------------------------------
def head_def(x, y, z):
    if x > 0:
        y *= 1.0 - 0.15 * x
    if z > 0 and x < 0.3:
        z *= 1.05
    return x, y, z


ellipsoid("Head", (0.255, 0, 0.38), (0.058, 0.036, 0.042), M['green'],
          rot=(0, math.radians(-4), 0), deform=head_def, seg=40, rings=20)

# Eyes (dark brown iris)
for s in (1, -1):
    sphere("Eye_" + ("L" if s > 0 else "R"), (0.272, s * 0.031, 0.391), 0.0065, M['eye'])

# ---------------------------------------------------------------------------
# Bill: upper mandible, lower mandible, nail, nostrils
# ---------------------------------------------------------------------------
pitch = math.radians(15)
bd = Vector((math.cos(pitch), 0, -math.sin(pitch)))
bup = Vector((math.sin(pitch), 0, math.cos(pitch)))
B = Vector((0.302, 0, 0.372))


def bill_def(x, y, z):
    u = (x + 1) / 2
    y *= 0.78 + 0.28 * u
    z *= 1.0 - 0.45 * u
    if z > 0:
        z *= 0.75
    return x, y, z


ellipsoid("Bill_UpperMandible", B + bd * 0.026, (0.028, 0.020, 0.011), M['bill'],
          rot=(0, pitch, 0), deform=bill_def)
ellipsoid("Bill_LowerMandible", B + bd * 0.024 - bup * 0.005, (0.025, 0.017, 0.005), M['bill'],
          rot=(0, pitch, 0), deform=bill_def)
ellipsoid("Bill_Nail", B + bd * 0.051 + bup * 0.0025, (0.006, 0.006, 0.0028), M['nail'],
          rot=(0, pitch, 0), seg=16, rings=8)
for s in (1, -1):
    ellipsoid("Nostril_" + ("L" if s > 0 else "R"),
              B + bd * 0.018 + bup * 0.0068 + Vector((0, s * 0.0065, 0)),
              (0.0042, 0.0015, 0.0015), M['nail'], rot=(0, pitch, 0), seg=12, rings=6)

# ---------------------------------------------------------------------------
# Wings (folded): coverts, speculum with white bars, tertials, primaries
# ---------------------------------------------------------------------------
for s in (1, -1):
    side = "L" if s > 0 else "R"
    ellipsoid("Wing_Coverts_" + side, (-0.02, s * 0.095, 0.215), (0.15, 0.03, 0.055), M['wing'],
              rot=(s * math.radians(-6), math.radians(5), 0))
    ellipsoid("Wing_Speculum_" + side, (-0.05, s * 0.122, 0.205), (0.055, 0.008, 0.016), M['spec'],
              rot=(0, math.radians(5), 0))
    for zb in (0.2235, 0.1865):
        ellipsoid("Wing_SpeculumBar_" + side, (-0.05, s * 0.1205, zb), (0.057, 0.007, 0.0042), M['white'],
                  rot=(0, math.radians(5), 0), seg=16, rings=8)
    for i in range(3):
        ellipsoid("Wing_Tertial_%s_%d" % (side, i),
                  (-0.07 - i * 0.016, s * (0.082 - i * 0.004), 0.242 - i * 0.002),
                  (0.052, 0.006, 0.02), M['back'],
                  rot=(s * math.radians(25), math.radians(4), s * math.radians(4)), seg=20, rings=10)
    for i in range(9):
        ellipsoid("Wing_Primary_%s_%d" % (side, i),
                  (-0.13 - i * 0.006, s * (0.072 - i * 0.004), 0.236 + i * 0.002),
                  (0.075, 0.004, 0.013), M['primary'],
                  rot=(s * math.radians(20), math.radians(4), s * math.radians(6 - i * 0.6)), seg=20, rings=10)

# ---------------------------------------------------------------------------
# Tail: rectrices fan + drake curled central feathers
# ---------------------------------------------------------------------------
tail_base = Vector((-0.205, 0, 0.208))
n_rect = 14
for i in range(n_rect):
    a = math.radians(-35 + 70 * i / (n_rect - 1))
    rot = (0, math.radians(-12), math.pi + a)
    ax = axis_of(rot)
    ellipsoid("Tail_Rectrix_%02d" % i, tail_base + ax * 0.04 + Vector((0, 0, 0.0015 * abs(i - 6.5) * -0.3)),
              (0.045, 0.012, 0.0025), M['tail'], rot=rot, seg=20, rings=8)

for k, (dx, dy) in enumerate(((0.0, 0.003), (-0.006, -0.003))):
    try:
        bpy.ops.mesh.primitive_torus_add(major_radius=0.011, minor_radius=0.0024,
                                         major_segments=24, minor_segments=8,
                                         location=(-0.212 + dx, dy, 0.258),
                                         rotation=(math.radians(90), 0, 0))
        finish(active(), "Tail_DrakeCurl_%d" % k, [M['black']])
    except Exception:
        pass

# ---------------------------------------------------------------------------
# Legs and webbed feet
# ---------------------------------------------------------------------------
def toe_point(F, yaw, L, t, z_tip=0.004):
    return Vector((F.x + math.cos(yaw) * L * t,
                   F.y + math.sin(yaw) * L * t,
                   F.z + (z_tip - F.z) * t + 0.0025 * math.sin(math.pi * t)))


def make_web(name, F, toeA, toeB):
    ts = [0.05, 0.3, 0.55, 0.8, 0.95]
    bm = bmesh.new()
    grid = []
    for i, t in enumerate(ts):
        a = toe_point(F, toeA[0], toeA[1], t) - Vector((0, 0, 0.0008))
        b = toe_point(F, toeB[0], toeB[1], t) - Vector((0, 0, 0.0008))
        mid = (a + b) / 2
        if i == len(ts) - 1:
            mid = F + (mid - F) * 0.78
        elif i == len(ts) - 2:
            mid = F + (mid - F) * 0.96
        mid.z -= 0.0005
        grid.append([bm.verts.new(a), bm.verts.new(mid), bm.verts.new(b)])
    for i in range(len(ts) - 1):
        for j in range(2):
            bm.faces.new((grid[i][j], grid[i][j + 1], grid[i + 1][j + 1], grid[i + 1][j]))
    try:
        ret = bmesh.ops.extrude_face_region(bm, geom=list(bm.faces))
        nv = [e for e in ret['geom'] if isinstance(e, bmesh.types.BMVert)]
        bmesh.ops.translate(bm, vec=(0, 0, -0.0012), verts=nv)
        bmesh.ops.recalc_face_normals(bm, faces=list(bm.faces))
    except Exception:
        pass
    me = bpy.data.meshes.new(name + "_mesh")
    bm.to_mesh(me)
    bm.free()
    ob = bpy.data.objects.new(name, me)
    bpy.context.collection.objects.link(ob)
    return finish(ob, name, [M['leg']])


for s in (1, -1):
    side = "L" if s > 0 else "R"
    # Feathered thigh
    ellipsoid("Thigh_Feathered_" + side, (-0.035, s * 0.055, 0.105), (0.032, 0.02, 0.026), M['belly'])
    knee = Vector((-0.03, s * 0.058, 0.095))
    heel = Vector((-0.04, s * 0.060, 0.058))
    F = Vector((-0.02, s * 0.062, 0.010))
    limb("Tibiotarsus_" + side, knee, heel, 0.0075, 0.0065, M['leg'])
    sphere("Heel_Joint_" + side, heel, 0.0068, M['leg'])
    limb("Tarsometatarsus_" + side, heel, F, 0.0058, 0.005, M['leg'])
    sphere("Foot_Base_" + side, F, 0.0055, M['leg'])

    base_yaw = s * math.radians(8)
    toes = {}
    for label, off, L, nph in (('II', -s * 0.50, 0.042, 3),
                               ('III', 0.0, 0.056, 4),
                               ('IV', s * 0.50, 0.052, 5)):
        yaw = base_yaw + off
        toes[label] = (yaw, L)
        for k in range(nph):
            t0, t1 = k / nph, (k + 1) / nph
            p0, p1 = toe_point(F, yaw, L, t0), toe_point(F, yaw, L, t1)
            r0, r1 = 0.0038 * (1 - 0.45 * t0), 0.0038 * (1 - 0.45 * t1)
            limb("Toe_%s_%s_Phalanx%d" % (label, side, k + 1), p0, p1, r0, r1, M['leg'], verts=12)
            if k > 0:
                sphere("Toe_%s_%s_Joint%d" % (label, side, k), p0, r0 * 1.05, M['leg'])
        tip = toe_point(F, yaw, L, 1.0)
        cd = Vector((math.cos(yaw), math.sin(yaw), -0.35)).normalized()
        limb("Claw_%s_%s" % (label, side), tip, tip + cd * 0.0065, 0.0019, 0.0002, M['claw'], verts=10)

    # Hallux (digit I): short, elevated, pointing back and medially
    hy = math.pi + s * math.radians(25)
    h0 = F + Vector((-0.003, 0, 0.003))
    hdir = Vector((math.cos(hy), math.sin(hy), 0.25)).normalized()
    h1 = h0 + hdir * 0.008
    h2 = h1 + hdir * 0.007
    limb("Hallux_%s_Phalanx1" % side, h0, h1, 0.0026, 0.0022, M['leg'], verts=10)
    sphere("Hallux_%s_Joint" % side, h1, 0.0023, M['leg'])
    limb("Hallux_%s_Phalanx2" % side, h1, h2, 0.0022, 0.0016, M['leg'], verts=10)
    limb("Hallux_%s_Claw" % side, h2, h2 + hdir * 0.004, 0.0015, 0.0002, M['claw'], verts=8)

    # Interdigital webbing (II-III and III-IV)
    make_web("Web_II_III_" + side, F, toes['II'], toes['III'])
    make_web("Web_III_IV_" + side, F, toes['III'], toes['IV'])

# ---------------------------------------------------------------------------
# Join all parts into one specimen object
# ---------------------------------------------------------------------------
duck = None
try:
    bpy.ops.object.select_all(action='DESELECT')
    body_obj = bpy.data.objects.get("Body_Torso")
    live = [o for o in bpy.data.objects if o.type == 'MESH']
    for o in live:
        o.select_set(True)
    if body_obj is not None:
        bpy.context.view_layer.objects.active = body_obj
        bpy.ops.object.join()
        duck = bpy.context.view_layer.objects.active
except Exception:
    duck = None

if duck is not None:
    duck.name = "Mallard_Anas_platyrhynchos_Drake"
    try:
        duck.data.name = "Mallard_Anas_platyrhynchos_Mesh"
    except Exception:
        pass
else:
    try:
        root = bpy.data.objects.new("Mallard_Anas_platyrhynchos_Drake", None)
        bpy.context.collection.objects.link(root)
        for o in bpy.data.objects:
            if o.type == 'MESH' and o.parent is None:
                o.parent = root
    except Exception:
        pass

try:
    bpy.ops.object.select_all(action='DESELECT')
except Exception:
    pass