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

random.seed(42)

# ---------------------------------------------------------------------------
# Clean start (allowed: we only clear at the BEGINNING, never at the end)
# ---------------------------------------------------------------------------
for o in list(bpy.data.objects):
    bpy.data.objects.remove(o, do_unlink=True)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def new_obj_from_bmesh(name, bm):
    me = bpy.data.meshes.new(name + "Mesh")
    bm.to_mesh(me)
    bm.free()
    obj = bpy.data.objects.new(name, me)
    bpy.context.collection.objects.link(obj)
    return obj


def make_material(name, base_color, roughness=0.7, subsurface=0.0, sss_color=None):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    bsdf = nodes.get("Principled BSDF")
    if bsdf is None:
        for n in nodes:
            if n.type == 'BSDF_PRINCIPLED':
                bsdf = n
                break
    if bsdf is not None:
        inp = bsdf.inputs
        # Defensive: guard every socket assignment individually.
        try:
            inp["Base Color"].default_value = base_color
        except Exception:
            pass
        try:
            inp["Roughness"].default_value = roughness
        except Exception:
            pass
        # Blender 4.x renamed 'Subsurface' -> 'Subsurface Weight'
        try:
            if "Subsurface Weight" in inp:
                inp["Subsurface Weight"].default_value = subsurface
            elif "Subsurface" in inp:
                inp["Subsurface"].default_value = subsurface
        except Exception:
            pass
        try:
            if sss_color is not None and "Subsurface Color" in inp:
                inp["Subsurface Color"].default_value = sss_color
        except Exception:
            pass
        try:
            if "Specular IOR Level" in inp:
                inp["Specular IOR Level"].default_value = 0.2
            elif "Specular" in inp:
                inp["Specular"].default_value = 0.2
        except Exception:
            pass
    return mat


def assign_material(obj, mat):
    if obj is None or obj.data is None:
        return
    obj.data.materials.clear()
    obj.data.materials.append(mat)


def apply_solidify(obj, thickness):
    if obj is None:
        return
    mod = obj.modifiers.new("Solid", 'SOLIDIFY')
    mod.thickness = thickness
    mod.offset = 0.0
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    try:
        bpy.ops.object.modifier_apply(modifier=mod.name)
    except Exception:
        pass


def shade_smooth(obj):
    if obj is None:
        return
    for p in obj.data.polygons:
        p.use_smooth = True


# ---------------------------------------------------------------------------
# Materials
# ---------------------------------------------------------------------------
mat_leaf = make_material("Leaf", (0.10, 0.22, 0.06, 1.0), roughness=0.55,
                         subsurface=0.15, sss_color=(0.15, 0.35, 0.08, 1.0))
mat_leaf_tip = make_material("LeafTip", (0.30, 0.20, 0.05, 1.0), roughness=0.7)
mat_wood = make_material("Crown", (0.22, 0.13, 0.07, 1.0), roughness=0.9)
mat_root = make_material("Root", (0.30, 0.20, 0.12, 1.0), roughness=0.95)
mat_cone = make_material("Cone", (0.45, 0.10, 0.06, 1.0), roughness=0.6)
mat_stalk = make_material("Stalk", (0.30, 0.16, 0.09, 1.0), roughness=0.85)


# ---------------------------------------------------------------------------
# 1) CROWN  (short woody obconical stem with a central concave groove)
# ---------------------------------------------------------------------------
bpy.ops.mesh.primitive_ico_sphere_add(subdivisions=4, radius=0.85,
                                       location=(0.0, 0.0, 0.0))
crown = bpy.context.active_object
crown.name = "Crown"
for v in crown.data.vertices:
    v.co.z *= 0.42
    r = math.hypot(v.co.x, v.co.y)
    # carve the concave saddle/groove into the top of the crown
    if v.co.z > 0.02:
        depression = max(0.0, (0.60 - r)) * 0.85
        v.co.z -= depression
        # slight elongation of the groove along X
        v.co.z -= max(0.0, (0.35 - abs(v.co.y))) * 0.15
    # raise slightly so the base sits at ground level
    v.co.z += 0.20
shade_smooth(crown)
assign_material(crown, mat_wood)


# ---------------------------------------------------------------------------
# 2) TAPROOT  (long tapering root + a few lateral rootlets)
# ---------------------------------------------------------------------------
bm = bmesh.new()
# Main taproot: wide at top (under crown), tapering downward.
bmesh.ops.create_cone(
    bm, cap_ends=True, cap_tris=False, segments=24,
    radius1=0.55, radius2=0.10, depth=3.2,
    matrix=Matrix.Translation((0.0, 0.0, -1.55))
)
# gentle bend + surface wobble on the taproot
for v in bm.verts:
    depth_f = (-v.co.z)
    v.co.x += 0.10 * math.sin(depth_f * 1.1)
    v.co.y += 0.06 * math.cos(depth_f * 0.9)
    r = math.hypot(v.co.x, v.co.y)
    if r > 0.02:
        wob = 1.0 + 0.06 * math.sin(v.co.z * 6.0 + math.atan2(v.co.y, v.co.x) * 4.0)
        v.co.x *= wob
        v.co.y *= wob

# Lateral rootlets
for i in range(5):
    ang = random.uniform(0, math.tau)
    down = random.uniform(-2.4, -0.6)
    length = random.uniform(0.7, 1.4)
    rad = random.uniform(0.05, 0.11)
    dirv = Vector((math.cos(ang), math.sin(ang), -0.6)).normalized()
    rot = dirv.to_track_quat('Z', 'Y').to_matrix().to_4x4()
    loc = Matrix.Translation((0.20 * math.cos(ang), 0.20 * math.sin(ang),
                              down)) @ rot @ Matrix.Translation((0, 0, length * 0.5))
    bmesh.ops.create_cone(
        bm, cap_ends=True, cap_tris=False, segments=10,
        radius1=rad, radius2=0.01, depth=length,
        matrix=loc
    )
root = new_obj_from_bmesh("Taproot", bm)
shade_smooth(root)
assign_material(root, mat_root)


# ---------------------------------------------------------------------------
# 3) LEAVES  (exactly TWO strap leaves, each frayed into several ribbons)
# ---------------------------------------------------------------------------
def build_leaf(name, direction):
    """direction = +1 or -1 along the Y axis."""
    bm = bmesh.new()
    n_ribbons = 7
    base_z = 0.34
    total_len = 6.5
    for r in range(n_ribbons):
        spread = r - (n_ribbons - 1) / 2.0
        lateral0 = spread * 0.11
        fan = spread * 0.13
        ribbon_len = total_len * (0.72 + 0.28 * random.random())
        width = 0.13 * (1.0 - 0.15 * abs(spread) / n_ribbons)
        segs = 46
        freq = 2.5 + random.random() * 2.5
        amp = 0.12 + 0.10 * random.random()
        phase = random.random() * math.tau
        xwave_f = 1.5 + random.random() * 2.0
        left = []
        right = []
        for i in range(segs + 1):
            u = i / segs
            dist = u * ribbon_len
            # planar path (fans outward from the groove)
            x = lateral0 + dist * math.sin(fan)
            y = direction * dist * math.cos(fan)
            # lateral undulation of the ribbon
            x += 0.18 * math.sin(u * xwave_f * math.pi + phase) * u
            # vertical droop from the crown down to the ground, undulating
            z = base_z * (1.0 - u) ** 0.7
            z += amp * math.sin(u * freq * math.pi + phase) * u
            if z < 0.012:
                z = 0.012
            w = width * (1.0 - 0.45 * u)
            # ribbon runs mainly along Y, so width spreads along X
            perp = Vector((1.0, 0.0, 0.0))
            cx = x
            left.append(bm.verts.new((cx - w, y, z)))
            right.append(bm.verts.new((cx + w, y, z)))
        for i in range(segs):
            bm.faces.new((left[i], right[i], right[i + 1], left[i + 1]))
    obj = new_obj_from_bmesh(name, bm)
    shade_smooth(obj)
    assign_material(obj, mat_leaf)
    apply_solidify(obj, 0.018)
    return obj


leaf_a = build_leaf("Leaf_A", +1.0)
leaf_b = build_leaf("Leaf_B", -1.0)


# ---------------------------------------------------------------------------
# 4) REPRODUCTIVE CONES  (branched peduncles bearing several cones on the crown)
# ---------------------------------------------------------------------------
bm = bmesh.new()
n_clusters = 4
for c in range(n_clusters):
    base_ang = (c / n_clusters) * math.tau + random.uniform(-0.2, 0.2)
    bx = 0.45 * math.cos(base_ang)
    by = 0.45 * math.sin(base_ang)
    bz = 0.18
    n_cones = random.randint(2, 3)
    for k in range(n_cones):
        a2 = base_ang + (k - n_cones / 2.0) * 0.28
        ex = bx + 0.18 * math.cos(a2)
        ey = by + 0.18 * math.sin(a2)
        h = random.uniform(0.30, 0.48)
        tipx = ex + 0.06 * math.cos(a2)
        tipy = ey + 0.06 * math.sin(a2)
        tipz = bz + h
        # peduncle (thin stalk)
        mid = Vector(((bx + tipx) / 2, (by + tipy) / 2, (bz + tipz) / 2))
        dirv = Vector((tipx - bx, tipy - by, tipz - bz))
        stalk_len = dirv.length
        dirv.normalize()
        rot = dirv.to_track_quat('Z', 'Y').to_matrix().to_4x4()
        mtx = Matrix.Translation(mid) @ rot
        bmesh.ops.create_cone(
            bm, cap_ends=True, cap_tris=False, segments=8,
            radius1=0.022, radius2=0.016, depth=stalk_len,
            matrix=mtx
        )
        # the cone (elongated ribbed ovoid) at the tip
        ret = bmesh.ops.create_icosphere(bm, subdivisions=2, radius=0.06,
                                          matrix=Matrix.Translation((0, 0, 0)))
        cverts = ret["verts"]
        for v in cverts:
            v.co.z *= 2.0            # elongate
            r = math.hypot(v.co.x, v.co.y)
            if r > 0.001:
                rib = 1.0 + 0.12 * math.sin(math.atan2(v.co.y, v.co.x) * 8.0)
                v.co.x *= rib
                v.co.y *= rib
            v.co += Vector((tipx, tipy, tipz))
repro = new_obj_from_bmesh("Cones", bm)
shade_smooth(repro)
assign_material(repro, mat_cone)


# ---------------------------------------------------------------------------
# 5) JOIN EVERYTHING INTO A SINGLE WHOLE SPECIMEN
# ---------------------------------------------------------------------------
parts = [crown, root, leaf_a, leaf_b, repro]
parts = [p for p in parts if p is not None and p.name in bpy.data.objects]

bpy.ops.object.select_all(action='DESELECT')
for p in parts:
    p.select_set(True)
bpy.context.view_layer.objects.active = parts[0]
if len(parts) > 1:
    try:
        bpy.ops.object.join()
    except Exception:
        pass

# Re-fetch the surviving (joined) object; the join may have invalidated refs.
final = bpy.context.view_layer.objects.active
if final is not None:
    final.name = "Welwitschia_mirabilis"
    final.location = (0.0, 0.0, 0.0)