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

random.seed(42)


def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)
    for block in bpy.data.meshes:
        if block.users == 0:
            bpy.data.meshes.remove(block)
    for block in bpy.data.materials:
        if block.users == 0:
            bpy.data.materials.remove(block)
    for block in bpy.data.curves:
        if block.users == 0:
            bpy.data.curves.remove(block)


def create_material(name, base_color, roughness=0.5, subsurface_weight=0.0):
    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    principled = nodes.get("Principled BSDF")
    if principled:
        try:
            principled.inputs["Base Color"].default_value = base_color
        except Exception:
            pass
        try:
            principled.inputs["Roughness"].default_value = float(roughness)
        except Exception:
            pass
        try:
            principled.inputs["Subsurface Weight"].default_value = float(subsurface_weight)
        except Exception:
            try:
                principled.inputs["Subsurface"].default_value = float(subsurface_weight)
            except Exception:
                pass
    return mat


def create_stalk(base_radius=0.035, height=2.2, num_segments=10):
    stalk_parts = []
    segment_height = height / num_segments
    
    for i in range(num_segments):
        taper = 1.0 - (i / num_segments) * 0.35
        r_bottom = base_radius * taper
        r_top = base_radius * (1.0 - ((i + 1) / num_segments) * 0.35)
        z_pos = i * segment_height + segment_height / 2
        
        bpy.ops.mesh.primitive_cone_add(
            vertices=16,
            radius1=r_bottom,
            radius2=r_top,
            depth=segment_height * 0.95,
            location=(0, 0, z_pos),
            end_fill_type='NGON'
        )
        obj = bpy.context.active_object
        if obj:
            obj.name = f"Stalk_{i}"
            stalk_parts.append(obj.name)
        
        if i < num_segments - 1:
            node_z = (i + 1) * segment_height
            bpy.ops.mesh.primitive_torus_add(
                major_radius=r_top * 1.2,
                minor_radius=r_top * 0.18,
                location=(0, 0, node_z)
            )
            node_obj = bpy.context.active_object
            if node_obj:
                node_obj.name = f"Node_{i}"
                stalk_parts.append(node_obj.name)
    
    return stalk_parts, segment_height


def create_leaf_mesh(length, width, droop, segments_u=15, segments_v=4):
    mesh = bpy.data.meshes.new("LeafMesh")
    bm = bmesh.new()
    
    verts = []
    for i in range(segments_u + 1):
        t = i / segments_u
        x = t * length
        z = -droop * (t ** 1.8)
        z += math.sin(t * math.pi * 3) * 0.008
        
        w = width * math.sin(t * math.pi) * 0.85 + width * 0.15
        if t < 0.08:
            w *= t / 0.08
        if t > 0.92:
            w *= (1 - t) / 0.08
        
        for j in range(segments_v + 1):
            s = (j / segments_v - 0.5) * 2
            y = s * w / 2
            local_z = z - abs(s) * 0.008
            v = bm.verts.new((x, y, local_z))
            verts.append(v)
    
    bm.verts.ensure_lookup_table()
    
    for i in range(segments_u):
        for j in range(segments_v):
            idx = i * (segments_v + 1) + j
            v1 = verts[idx]
            v2 = verts[idx + 1]
            v3 = verts[idx + segments_v + 2]
            v4 = verts[idx + segments_v + 1]
            try:
                bm.faces.new([v1, v2, v3, v4])
            except Exception:
                pass
    
    bm.to_mesh(mesh)
    bm.free()
    return mesh


def create_leaves(stalk_height, num_leaves=12):
    leaf_objects = []
    
    for i in range(num_leaves):
        t = (i + 2) / (num_leaves + 4)
        z = t * stalk_height
        angle = (i % 2) * 180 + random.uniform(-25, 25)
        size_factor = 1.0 - abs(t - 0.45) * 1.4
        size_factor = max(0.35, min(1.0, size_factor))
        
        leaf_length = 0.55 + size_factor * 0.55
        leaf_width = 0.055 + size_factor * 0.045
        droop = 0.15 + (1 - t) * 0.35 + size_factor * 0.1
        
        mesh = create_leaf_mesh(leaf_length, leaf_width, droop)
        leaf_obj = bpy.data.objects.new(f"Leaf_{i}", mesh)
        bpy.context.collection.objects.link(leaf_obj)
        
        offset = 0.04
        x_off = math.cos(math.radians(angle)) * offset
        y_off = math.sin(math.radians(angle)) * offset
        
        leaf_obj.location = (x_off, y_off, z)
        leaf_obj.rotation_euler = (0, 0, math.radians(angle))
        leaf_objects.append(leaf_obj.name)
    
    return leaf_objects


def create_tassel(location, height=0.28):
    tassel_parts = []
    
    bpy.ops.mesh.primitive_cone_add(
        vertices=8,
        radius1=0.012,
        radius2=0.002,
        depth=height * 0.7,
        location=(location[0], location[1], location[2] + height * 0.35),
        end_fill_type='NGON'
    )
    central = bpy.context.active_object
    if central:
        central.name = "TasselCentral"
        tassel_parts.append(central.name)
    
    num_branches = 18
    for i in range(num_branches):
        angle = (i / num_branches) * math.pi * 2 + random.uniform(-0.15, 0.15)
        tier = i % 4
        z_offset = tier * 0.04 - 0.02
        branch_len = 0.10 + random.uniform(-0.025, 0.025)
        
        dx = math.cos(angle) * branch_len
        dy = math.sin(angle) * branch_len
        dz = random.uniform(0.01, 0.05)
        
        mid_x = location[0] + dx / 2
        mid_y = location[1] + dy / 2
        mid_z = location[2] + z_offset + dz / 2
        
        bpy.ops.mesh.primitive_cone_add(
            vertices=6,
            radius1=0.0025,
            radius2=0.0008,
            depth=branch_len,
            location=(mid_x, mid_y, mid_z),
            end_fill_type='NGON'
        )
        branch = bpy.context.active_object
        if branch:
            branch.rotation_euler = (
                math.radians(55 + random.uniform(-8, 8)),
                0,
                angle
            )
            branch.name = f"TasselBranch_{i}"
            tassel_parts.append(branch.name)
            
            end_x = location[0] + dx
            end_y = location[1] + dy
            end_z = location[2] + z_offset + dz
            
            for j in range(2):
                ax = end_x + random.uniform(-0.006, 0.006)
                ay = end_y + random.uniform(-0.006, 0.006)
                az = end_z + random.uniform(-0.008, 0.008)
                
                bpy.ops.mesh.primitive_uv_sphere_add(
                    radius=0.0035,
                    segments=6,
                    ring_count=4,
                    location=(ax, ay, az)
                )
                anther = bpy.context.active_object
                if anther:
                    anther.scale = (1.0, 0.6, 0.6)
                    anther.name = f"Anther_{i}_{j}"
                    tassel_parts.append(anther.name)
    
    return tassel_parts


def create_cob_with_kernels(location, length=0.16, radius=0.028, num_rows=12, kernels_per_row=22):
    parts = []
    
    bpy.ops.mesh.primitive_cone_add(
        vertices=20,
        radius1=radius * 0.9,
        radius2=radius * 0.9,
        depth=length,
        location=location,
        end_fill_type='NGON'
    )
    cob = bpy.context.active_object
    if cob:
        cob.name = "Cob"
        parts.append(cob.name)
    
    for row in range(num_rows):
        row_angle = (row / num_rows) * 2 * math.pi
        
        for k in range(kernels_per_row):
            t = (k + 0.5) / kernels_per_row
            z = location[2] - length / 2 + t * length * 0.88 + length * 0.06
            
            offset_angle = row_angle + (k % 2) * (math.pi / num_rows)
            
            kx = location[0] + math.cos(offset_angle) * (radius * 0.95)
            ky = location[1] + math.sin(offset_angle) * (radius * 0.95)
            
            bpy.ops.mesh.primitive_uv_sphere_add(
                radius=0.0065,
                segments=5,
                ring_count=3,
                location=(kx, ky, z)
            )
            kernel = bpy.context.active_object
            if kernel:
                kernel.scale = (0.55, 0.75, 1.0)
                kernel.rotation_euler = (0, 0, offset_angle)
                kernel.name = f"Kernel_{row}_{k}"
                parts.append(kernel.name)
    
    return parts


def create_husk_leaf(layer, location, length, num_layers):
    mesh = bpy.data.meshes.new(f"HuskMesh_{layer}")
    bm = bmesh.new()
    
    segments_u = 12
    segments_v = 10
    
    angle_offset = layer * 0.25
    openness = (layer / num_layers) * 0.4
    husk_length = length * (1.05 + layer * 0.08)
    wrap_angle = math.pi * (0.9 - layer * 0.1)
    
    verts = []
    for i in range(segments_u + 1):
        u = i / segments_u
        z = u * husk_length
        
        for j in range(segments_v + 1):
            v = j / segments_v - 0.5
            angle = v * wrap_angle + angle_offset
            
            r = 0.032 + layer * 0.006
            if u < 0.15:
                r *= u / 0.15
            if u > 0.75:
                r += (u - 0.75) * openness * 0.12
            
            x = math.cos(angle) * r
            y = math.sin(angle) * r
            
            vert = bm.verts.new((x, y, z))
            verts.append(vert)
    
    bm.verts.ensure_lookup_table()
    
    for i in range(segments_u):
        for j in range(segments_v):
            idx = i * (segments_v + 1) + j
            try:
                bm.faces.new([
                    verts[idx],
                    verts[idx + 1],
                    verts[idx + segments_v + 2],
                    verts[idx + segments_v + 1]
                ])
            except Exception:
                pass
    
    bm.to_mesh(mesh)
    bm.free()
    
    obj = bpy.data.objects.new(f"Husk_{layer}", mesh)
    bpy.context.collection.objects.link(obj)
    obj.location = (location[0], location[1], location[2] - length / 2)
    
    return obj.name


def create_husk(location, length=0.20, num_layers=5):
    parts = []
    for layer in range(num_layers):
        name = create_husk_leaf(layer, location, length, num_layers)
        parts.append(name)
    return parts


def create_silk(location, num_strands=35):
    parts = []
    
    for i in range(num_strands):
        angle = random.uniform(0, 2 * math.pi)
        r = random.uniform(0, 0.022)
        
        start_x = location[0] + math.cos(angle) * r
        start_y = location[1] + math.sin(angle) * r
        start_z = location[2]
        
        curve_data = bpy.data.curves.new(f'SilkCurve_{i}', type='CURVE')
        curve_data.dimensions = '3D'
        
        spline = curve_data.splines.new('BEZIER')
        spline.bezier_points.add(2)
        
        strand_length = random.uniform(0.06, 0.12)
        curl_x = random.uniform(-0.04, 0.04)
        curl_y = random.uniform(-0.04, 0.04)
        
        points = [
            (start_x, start_y, start_z),
            (start_x + curl_x * 0.5, start_y + curl_y * 0.5, start_z + strand_length * 0.4),
            (start_x + curl_x, start_y + curl_y, start_z + strand_length)
        ]
        
        for j, point in enumerate(points):
            bp = spline.bezier_points[j]
            bp.co = point
            bp.handle_left_type = 'AUTO'
            bp.handle_right_type = 'AUTO'
        
        curve_data.bevel_depth = 0.0008
        curve_data.bevel_resolution = 1
        
        silk_obj = bpy.data.objects.new(f'Silk_{i}', curve_data)
        bpy.context.collection.objects.link(silk_obj)
        parts.append(silk_obj.name)
    
    return parts


def create_ear(stalk_height):
    ear_height = stalk_height * 0.55
    ear_location = (0.055, 0, ear_height)
    
    all_parts = []
    
    cob_parts = create_cob_with_kernels(ear_location, length=0.16, num_rows=10, kernels_per_row=18)
    all_parts.extend(cob_parts)
    
    husk_parts = create_husk(ear_location, length=0.19, num_layers=5)
    all_parts.extend(husk_parts)
    
    silk_location = (ear_location[0], ear_location[1], ear_location[2] + 0.08)
    silk_parts = create_silk(silk_location, num_strands=30)
    all_parts.extend(silk_parts)
    
    return all_parts


def create_root_curve(name, points, bevel_depth):
    curve_data = bpy.data.curves.new(name, type='CURVE')
    curve_data.dimensions = '3D'
    
    spline = curve_data.splines.new('BEZIER')
    spline.bezier_points.add(len(points) - 1)
    
    for j, point in enumerate(points):
        bp = spline.bezier_points[j]
        bp.co = point
        bp.handle_left_type = 'AUTO'
        bp.handle_right_type = 'AUTO'
    
    curve_data.bevel_depth = bevel_depth
    curve_data.bevel_resolution = 2
    
    root_obj = bpy.data.objects.new(name, curve_data)
    bpy.context.collection.objects.link(root_obj)
    return root_obj.name


def create_roots(depth=0.28, spread=0.18):
    parts = []
    
    num_roots = 20
    for i in range(num_roots):
        angle = (i / num_roots) * 2 * math.pi + random.uniform(-0.25, 0.25)
        root_length = depth * random.uniform(0.55, 1.0)
        root_spread = spread * random.uniform(0.4, 1.0)
        
        end_angle = angle + random.uniform(-0.4, 0.4)
        end_x = math.cos(end_angle) * root_spread
        end_y = math.sin(end_angle) * root_spread
        end_z = -root_length
        
        points = [
            Vector((0, 0, -0.01)),
            Vector((end_x * 0.35, end_y * 0.35, -root_length * 0.45)),
            Vector((end_x, end_y, end_z))
        ]
        
        bevel = 0.006 - (i % 5) * 0.0008
        name = create_root_curve(f'Root_{i}', points, bevel)
        parts.append(name)
    
    num_brace = 6
    for i in range(num_brace):
        angle = (i / num_brace) * 2 * math.pi + random.uniform(-0.2, 0.2)
        start_z = 0.08 + random.uniform(0, 0.04)
        end_spread = 0.10 + random.uniform(0, 0.04)
        
        points = [
            Vector((math.cos(angle) * 0.038, math.sin(angle) * 0.038, start_z)),
            Vector((math.cos(angle) * 0.07, math.sin(angle) * 0.07, start_z * 0.35)),
            Vector((math.cos(angle) * end_spread, math.sin(angle) * end_spread, 0))
        ]
        
        name = create_root_curve(f'BraceRoot_{i}', points, 0.005)
        parts.append(name)
    
    return parts


def apply_materials():
    stalk_mat = create_material("StalkMaterial", (0.18, 0.32, 0.08, 1.0), roughness=0.75)
    leaf_mat = create_material("LeafMaterial", (0.12, 0.38, 0.06, 1.0), roughness=0.65, subsurface_weight=0.08)
    kernel_mat = create_material("KernelMaterial", (0.92, 0.72, 0.18, 1.0), roughness=0.35)
    husk_mat = create_material("HuskMaterial", (0.38, 0.48, 0.18, 1.0), roughness=0.72)
    silk_mat = create_material("SilkMaterial", (0.82, 0.58, 0.28, 1.0), roughness=0.85)
    tassel_mat = create_material("TasselMaterial", (0.58, 0.48, 0.18, 1.0), roughness=0.65)
    root_mat = create_material("RootMaterial", (0.38, 0.28, 0.18, 1.0), roughness=0.92)
    cob_mat = create_material("CobMaterial", (0.88, 0.82, 0.65, 1.0), roughness=0.55)
    
    for obj in bpy.data.objects:
        if obj.type not in ['MESH', 'CURVE']:
            continue
        
        name = obj.name.lower()
        
        if hasattr(obj.data, 'materials'):
            obj.data.materials.clear()
            
            mat = None
            if 'stalk' in name or 'node' in name:
                mat = stalk_mat
            elif 'leaf' in name:
                mat = leaf_mat
            elif 'kernel' in name:
                mat = kernel_mat
            elif 'husk' in name:
                mat = husk_mat
            elif 'silk' in name:
                mat = silk_mat
            elif 'tassel' in name or 'anther' in name:
                mat = tassel_mat
            elif 'root' in name or 'brace' in name:
                mat = root_mat
            elif 'cob' in name:
                mat = cob_mat
            
            if mat:
                obj.data.materials.append(mat)


def join_all_parts():
    bpy.ops.object.select_all(action='DESELECT')
    
    curve_objects = [obj for obj in bpy.data.objects if obj.type == 'CURVE']
    for obj in curve_objects:
        obj.select_set(True)
        bpy.context.view_layer.objects.active = obj
    
    if curve_objects:
        bpy.ops.object.convert(target='MESH')
    
    bpy.ops.object.select_all(action='DESELECT')
    
    all_meshes = [obj for obj in bpy.data.objects if obj.type == 'MESH']
    
    if all_meshes:
        for obj in all_meshes:
            obj.select_set(True)
        bpy.context.view_layer.objects.active = all_meshes[0]
        bpy.ops.object.join()
        
        final_obj = bpy.context.active_object
        if final_obj:
            final_obj.name = "ZeaMays_Corn"
            bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='MEDIAN')


def main():
    clear_scene()
    
    stalk_height = 2.2
    stalk_parts, segment_height = create_stalk(
        base_radius=0.035,
        height=stalk_height,
        num_segments=10
    )
    
    leaf_parts = create_leaves(stalk_height, num_leaves=12)
    
    tassel_parts = create_tassel(
        location=(0, 0, stalk_height),
        height=0.26
    )
    
    ear_parts = create_ear(stalk_height)
    
    root_parts = create_roots(depth=0.26, spread=0.16)
    
    apply_materials()
    
    join_all_parts()


if __name__ == "__main__":
    main()