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

random.seed(42)


def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)
    for mesh in bpy.data.meshes:
        bpy.data.meshes.remove(mesh)
    for mat in bpy.data.materials:
        bpy.data.materials.remove(mat)


def safe_socket_set(inputs, name, value):
    try:
        if name in inputs:
            inputs[name].default_value = value
    except Exception:
        pass


def create_material(name, base_color, roughness=0.8, subsurface_weight=0.0):
    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    principled = None
    for node in nodes:
        if node.type == 'BSDF_PRINCIPLED':
            principled = node
            break
    if principled:
        safe_socket_set(principled.inputs, 'Base Color', base_color)
        safe_socket_set(principled.inputs, 'Roughness', roughness)
        safe_socket_set(principled.inputs, 'Subsurface Weight', subsurface_weight)
        safe_socket_set(principled.inputs, 'Subsurface', subsurface_weight)
    return mat


def create_crown():
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=64,
        radius=1.0,
        depth=0.5,
        location=(0, 0, 0.1)
    )
    crown = bpy.context.active_object
    if crown is None:
        return None
    crown.name = "Welwitschia_Crown"
    
    bm = bmesh.new()
    bm.from_mesh(crown.data)
    bm.verts.ensure_lookup_table()
    
    for v in bm.verts:
        xy = Vector((v.co.x, v.co.y, 0))
        r = xy.length
        h_norm = (v.co.z + 0.25) / 0.5
        
        if r > 0.01:
            n1 = noise.noise(v.co * 1.2) * 0.2
            n2 = noise.noise(v.co * 3.5) * 0.08
            n3 = noise.noise(v.co * 7.0) * 0.03
            radial_d = n1 + n2 + n3
            v.co += xy.normalized() * radial_d
        
        if h_norm > 0.6:
            v.co.z += noise.noise(v.co * 2.5) * 0.12
        
        if r < 0.35 and h_norm > 0.4:
            v.co.z -= (0.35 - r) * 0.25
        
        if r > 0.7 and h_norm < 0.3:
            v.co.z += noise.noise(v.co * 4) * 0.04
    
    bm.to_mesh(crown.data)
    bm.free()
    
    mod = crown.modifiers.new("Subsurf", 'SUBSURF')
    mod.levels = 2
    mod.render_levels = 2
    
    mat = create_material("Crown_Mat", (0.20, 0.13, 0.07, 1.0), 0.95)
    crown.data.materials.append(mat)
    
    return crown


def create_taproot():
    bpy.ops.mesh.primitive_cone_add(
        vertices=20,
        radius1=0.35,
        radius2=0.04,
        depth=1.8,
        location=(0, 0, -0.75)
    )
    root = bpy.context.active_object
    if root is None:
        return None
    root.name = "Taproot"
    
    bm = bmesh.new()
    bm.from_mesh(root.data)
    for v in bm.verts:
        if v.co.z < 0:
            xy = Vector((v.co.x, v.co.y, 0))
            if xy.length > 0.01:
                d = noise.noise(v.co * 2.5) * 0.06
                v.co += xy.normalized() * d
    bm.to_mesh(root.data)
    bm.free()
    
    mat = create_material("Root_Mat", (0.32, 0.22, 0.12, 1.0), 0.92)
    root.data.materials.append(mat)
    return root


def create_leaf_ribbon_mesh(base_angle, length, width, segments_l=60, segments_w=8, droop_factor=1.0, wave_amp=1.0):
    verts = []
    faces = []
    
    for i in range(segments_l + 1):
        t = i / segments_l
        dist_from_crown = 0.4 + t * length
        
        taper = 1.0 - 0.35 * t
        w = width * taper
        
        angle_drift = noise.noise(Vector((t * 3, base_angle, 0))) * 0.15
        current_angle = base_angle + angle_drift * t
        
        for j in range(segments_w + 1):
            s = j / segments_w
            lateral = (s - 0.5) * w
            
            base_x = dist_from_crown * math.cos(current_angle)
            base_y = dist_from_crown * math.sin(current_angle)
            
            perp_x = -math.sin(current_angle)
            perp_y = math.cos(current_angle)
            
            x = base_x + lateral * perp_x
            y = base_y + lateral * perp_y
            
            z = 0.35 * (1.0 - t * droop_factor)
            z = max(z, 0.01 + random.uniform(0, 0.02))
            
            z += math.sin(t * math.pi * 5) * 0.025 * wave_amp
            z += math.sin(t * math.pi * 2.3) * 0.015 * wave_amp
            
            edge_f = abs(s - 0.5) * 2
            z -= edge_f * edge_f * 0.015
            
            z += noise.noise(Vector((x, y, t * 10))) * 0.02
            
            verts.append((x, y, z))
    
    for i in range(segments_l):
        for j in range(segments_w):
            v0 = i * (segments_w + 1) + j
            v1 = v0 + 1
            v2 = v0 + (segments_w + 1) + 1
            v3 = v0 + (segments_w + 1)
            faces.append((v0, v1, v2, v3))
    
    mesh = bpy.data.meshes.new("LeafRibbon")
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    return mesh


def create_leaf_assembly(base_angle, name_prefix, mat_base):
    objects = []
    
    main_mesh = create_leaf_ribbon_mesh(
        base_angle=base_angle,
        length=3.5,
        width=0.65,
        segments_l=70,
        segments_w=10,
        droop_factor=0.95,
        wave_amp=1.0
    )
    main_obj = bpy.data.objects.new(f"{name_prefix}_Main", main_mesh)
    bpy.context.collection.objects.link(main_obj)
    
    mod = main_obj.modifiers.new("Subsurf", 'SUBSURF')
    mod.levels = 1
    mod.render_levels = 2
    
    main_mat = create_material(
        f"{name_prefix}_MainMat",
        (mat_base[0], mat_base[1], mat_base[2], 1.0),
        0.72,
        0.08
    )
    main_obj.data.materials.append(main_mat)
    objects.append(main_obj)
    
    num_strips = random.randint(6, 10)
    for i in range(num_strips):
        strip_angle = base_angle + (i - num_strips / 2) * 0.08
        strip_angle += random.uniform(-0.05, 0.05)
        
        strip_length = random.uniform(2.5, 5.0)
        strip_width = random.uniform(0.06, 0.14)
        
        strip_mesh = create_leaf_ribbon_mesh(
            base_angle=strip_angle,
            length=strip_length,
            width=strip_width,
            segments_l=40,
            segments_w=4,
            droop_factor=1.1,
            wave_amp=1.5
        )
        strip_obj = bpy.data.objects.new(f"{name_prefix}_Strip_{i}", strip_mesh)
        bpy.context.collection.objects.link(strip_obj)
        
        start_dist = 3.2 + random.uniform(-0.3, 0.3)
        strip_obj.location = (
            start_dist * math.cos(strip_angle),
            start_dist * math.sin(strip_angle),
            0.03
        )
        
        brown_factor = 0.1 + random.uniform(0, 0.15)
        strip_mat = create_material(
            f"{name_prefix}_StripMat_{i}",
            (mat_base[0] + brown_factor, mat_base[1] + brown_factor * 0.5, mat_base[2] - 0.05, 1.0),
            0.85
        )
        strip_obj.data.materials.append(strip_mat)
        objects.append(strip_obj)
    
    return objects


def create_secondary_leaf_lobes(base_angle, name_prefix, mat_base):
    objects = []
    
    num_lobes = random.randint(2, 4)
    for i in range(num_lobes):
        lobe_angle = base_angle + (i - num_lobes / 2) * 0.25 + random.uniform(-0.1, 0.1)
        
        lobe_mesh = create_leaf_ribbon_mesh(
            base_angle=lobe_angle,
            length=random.uniform(2.0, 3.0),
            width=random.uniform(0.25, 0.4),
            segments_l=45,
            segments_w=6,
            droop_factor=1.0,
            wave_amp=0.8
        )
        lobe_obj = bpy.data.objects.new(f"{name_prefix}_Lobe_{i}", lobe_mesh)
        bpy.context.collection.objects.link(lobe_obj)
        
        color_var = random.uniform(-0.03, 0.03)
        lobe_mat = create_material(
            f"{name_prefix}_LobeMat_{i}",
            (mat_base[0] + color_var, mat_base[1] + color_var, mat_base[2] + color_var, 1.0),
            0.75,
            0.05
        )
        lobe_obj.data.materials.append(lobe_mat)
        objects.append(lobe_obj)
    
    return objects


def create_cone_stalk(position, index, height=0.35):
    bpy.ops.mesh.primitive_cone_add(
        vertices=8,
        radius1=0.022,
        radius2=0.022,
        depth=height,
        location=(position[0], position[1], position[2] + height / 2)
    )
    stalk = bpy.context.active_object
    if stalk is None:
        return None
    stalk.name = f"ConeStalk_{index}"
    
    bm = bmesh.new()
    bm.from_mesh(stalk.data)
    for v in bm.verts:
        lean = (v.co.z + height / 2) * 0.25
        v.co.x += lean * (0.5 if position[0] >= 0 else -0.5)
        v.co.y += lean * (0.3 if position[1] >= 0 else -0.3)
    bm.to_mesh(stalk.data)
    bm.free()
    
    mat = create_material(f"StalkMat_{index}", (0.26, 0.18, 0.10, 1.0), 0.88)
    stalk.data.materials.append(mat)
    return stalk


def create_male_cone(position, index):
    stalk = create_cone_stalk(position, index, 0.3)
    
    cone_x = position[0] + 0.12 * (1 if position[0] >= 0 else -1)
    cone_y = position[1] + 0.08 * (1 if position[1] >= 0 else -1)
    cone_z = position[2] + 0.38
    
    bpy.ops.mesh.primitive_cone_add(
        vertices=16,
        radius1=0.055,
        radius2=0.015,
        depth=0.12,
        location=(cone_x, cone_y, cone_z)
    )
    cone = bpy.context.active_object
    if cone is None:
        return [stalk] if stalk else []
    cone.name = f"MaleCone_{index}"
    
    bm = bmesh.new()
    bm.from_mesh(cone.data)
    for v in bm.verts:
        v.co += Vector((
            noise.noise(v.co * 15) * 0.008,
            noise.noise(v.co * 15 + Vector((1, 0, 0))) * 0.008,
            noise.noise(v.co * 15 + Vector((0, 1, 0))) * 0.005
        ))
    bm.to_mesh(cone.data)
    bm.free()
    
    mat = create_material(f"MaleConeMat_{index}", (0.65, 0.35, 0.22, 1.0), 0.72)
    cone.data.materials.append(mat)
    
    return [stalk, cone] if stalk else [cone]


def create_female_cone(position, index):
    stalk = create_cone_stalk(position, index, 0.25)
    
    cone_x = position[0] + 0.10 * (1 if position[0] >= 0 else -1)
    cone_y = position[1] + 0.06 * (1 if position[1] >= 0 else -1)
    cone_z = position[2] + 0.32
    
    bpy.ops.mesh.primitive_uv_sphere_add(
        segments=20,
        ring_count=14,
        radius=0.085,
        location=(cone_x, cone_y, cone_z)
    )
    cone = bpy.context.active_object
    if cone is None:
        return [stalk] if stalk else []
    cone.name = f"FemaleCone_{index}"
    cone.scale = (1.0, 1.0, 0.65)
    
    bpy.context.view_layer.objects.active = cone
    cone.select_set(True)
    bpy.ops.object.transform_apply(scale=True)
    
    bm = bmesh.new()
    bm.from_mesh(cone.data)
    for v in bm.verts:
        v.co += Vector((
            noise.noise(v.co * 12) * 0.012,
            noise.noise(v.co * 12 + Vector((1, 0, 0))) * 0.012,
            noise.noise(v.co * 12 + Vector((0, 1, 0))) * 0.008
        ))
    bm.to_mesh(cone.data)
    bm.free()
    
    mat = create_material(f"FemaleConeMat_{index}", (0.48, 0.42, 0.28, 1.0), 0.78)
    cone.data.materials.append(mat)
    
    return [stalk, cone] if stalk else [cone]


def create_surface_texture_details(crown):
    num_bumps = 40
    for i in range(num_bumps):
        angle = random.uniform(0, 2 * math.pi)
        r = random.uniform(0.4, 0.95)
        x = r * math.cos(angle)
        y = r * math.sin(angle)
        z = 0.3 + random.uniform(0, 0.15)
        
        bpy.ops.mesh.primitive_uv_sphere_add(
            segments=6,
            ring_count=4,
            radius=random.uniform(0.03, 0.07),
            location=(x, y, z)
        )
        bump = bpy.context.active_object
        if bump:
            bump.name = f"CrownBump_{i}"
            bump.scale = (
                random.uniform(0.8, 1.2),
                random.uniform(0.8, 1.2),
                random.uniform(0.4, 0.8)
            )
            bpy.ops.object.transform_apply(scale=True)
            
            mat = create_material(
                f"BumpMat_{i}",
                (0.22 + random.uniform(-0.03, 0.03), 0.14, 0.08, 1.0),
                0.92
            )
            bump.data.materials.append(mat)
            bump.parent = crown


def create_welwitschia():
    clear_scene()
    
    all_parts = []
    
    crown = create_crown()
    if crown:
        all_parts.append(crown)
        create_surface_texture_details(crown)
    
    root = create_taproot()
    if root:
        all_parts.append(root)
    
    leaf_color_a = (0.30, 0.38, 0.26)
    leaf_color_b = (0.32, 0.36, 0.24)
    
    leaf_a_parts = create_leaf_assembly(0, "LeafA", leaf_color_a)
    all_parts.extend(leaf_a_parts)
    
    leaf_a_lobes = create_secondary_leaf_lobes(0, "LeafA", leaf_color_a)
    all_parts.extend(leaf_a_lobes)
    
    leaf_b_parts = create_leaf_assembly(math.pi, "LeafB", leaf_color_b)
    all_parts.extend(leaf_b_parts)
    
    leaf_b_lobes = create_secondary_leaf_lobes(math.pi, "LeafB", leaf_color_b)
    all_parts.extend(leaf_b_lobes)
    
    cone_positions = [
        (0.22, 0.22, 0.38),
        (-0.18, 0.26, 0.38),
        (0.26, -0.18, 0.38),
        (-0.22, -0.22, 0.38),
        (0.05, 0.30, 0.38),
        (-0.05, -0.30, 0.38),
        (0.30, 0.05, 0.38),
        (-0.30, -0.05, 0.38),
    ]
    
    for i, pos in enumerate(cone_positions):
        if i % 2 == 0:
            cone_parts = create_male_cone(pos, i)
        else:
            cone_parts = create_female_cone(pos, i)
        all_parts.extend([p for p in cone_parts if p])
    
    bpy.ops.object.select_all(action='DESELECT')
    
    if crown:
        for obj in all_parts:
            if obj and obj != crown and obj.name in bpy.data.objects:
                try:
                    obj.select_set(True)
                except ReferenceError:
                    pass
        
        bpy.context.view_layer.objects.active = crown
        if bpy.context.selected_objects:
            bpy.ops.object.parent_set(type='OBJECT', keep_transform=True)
        
        bpy.ops.object.select_all(action='DESELECT')
    
    bpy.ops.object.empty_add(type='PLAIN_AXES', location=(0, 0, 0))
    plant_empty = bpy.context.active_object
    if plant_empty:
        plant_empty.name = "Welwitschia_mirabilis"
        if crown:
            crown.parent = plant_empty
    
    bpy.ops.object.select_all(action='DESELECT')
    if plant_empty:
        plant_empty.select_set(True)
        bpy.context.view_layer.objects.active = plant_empty
    
    print("Welwitschia mirabilis model created successfully.")


if __name__ == "__main__":
    create_welwitschia()