import bpy
import bmesh
import math
import random


def clear_scene():
    bpy.ops.object.select_all(action='DESELECT')
    for obj in list(bpy.data.objects):
        if obj.type == 'MESH':
            bpy.data.objects.remove(obj, do_unlink=True)


def create_material(name, color, roughness=0.5, subsurface=0.0):
    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    
    principled = None
    for node in mat.node_tree.nodes:
        if node.type == 'BSDF_PRINCIPLED':
            principled = node
            break
    
    if principled:
        try:
            principled.inputs['Base Color'].default_value = color
        except Exception:
            pass
        try:
            principled.inputs['Roughness'].default_value = roughness
        except Exception:
            pass
        try:
            principled.inputs['Subsurface Weight'].default_value = subsurface
        except Exception:
            try:
                principled.inputs['Subsurface'].default_value = subsurface
            except Exception:
                pass
    
    return mat


def create_cap():
    bpy.ops.mesh.primitive_uv_sphere_add(
        segments=64, ring_count=32, radius=1.0, location=(0, 0, 2.5)
    )
    obj = bpy.context.active_object
    if obj is None:
        return None
    obj.name = "Cap"
    
    bm = bmesh.new()
    bm.from_mesh(obj.data)
    
    verts_to_delete = [v for v in bm.verts if v.co.z < 0]
    bmesh.ops.delete(bm, geom=verts_to_delete, context='VERTS')
    
    for v in bm.verts:
        h = v.co.z
        r = math.sqrt(v.co.x**2 + v.co.y**2)
        if r > 0.001:
            new_z = h * 0.38 + 0.02
            expansion = 1.0 + 0.22 * (1 - h)
            angle = math.atan2(v.co.y, v.co.x)
            v.co.x = r * expansion * math.cos(angle)
            v.co.y = r * expansion * math.sin(angle)
            v.co.z = new_z
    
    bm.to_mesh(obj.data)
    bm.free()
    
    return bpy.data.objects.get("Cap")


def create_cap_flesh():
    bpy.ops.mesh.primitive_cone_add(
        vertices=64, radius1=1.2, radius2=0.14, depth=0.38, location=(0, 0, 2.33)
    )
    obj = bpy.context.active_object
    if obj:
        obj.name = "CapFlesh"
    return bpy.data.objects.get("CapFlesh")


def create_gills():
    gill_parts = []
    num_gills = 72
    
    for i in range(num_gills):
        angle = 2 * math.pi * i / num_gills
        
        bpy.ops.mesh.primitive_plane_add(size=1.0, location=(0, 0, 2.18))
        gill = bpy.context.active_object
        if gill is None:
            continue
        
        gill.name = f"Gill{i:03d}"
        gill.rotation_euler = (math.pi / 2, 0, angle)
        gill.scale = (0.008, 0.48, 0.14)
        
        offset = 0.58
        gill.location.x = offset * math.cos(angle)
        gill.location.y = offset * math.sin(angle)
        
        gill_parts.append(gill.name)
    
    bpy.ops.object.select_all(action='DESELECT')
    first = None
    for name in gill_parts:
        obj = bpy.data.objects.get(name)
        if obj:
            obj.select_set(True)
            if first is None:
                first = obj
    
    if first:
        bpy.context.view_layer.objects.active = first
        bpy.ops.object.join()
        result = bpy.context.active_object
        if result:
            result.name = "Gills"
    
    return bpy.data.objects.get("Gills")


def create_warts():
    random.seed(42)
    wart_parts = []
    
    for i in range(32):
        theta = random.uniform(0, 2 * math.pi)
        phi = random.uniform(0.18, 0.88)
        
        cap_r = 1.22 * math.sqrt(1 - phi**2) * (1 + 0.22 * (1 - phi))
        x = cap_r * math.cos(theta)
        y = cap_r * math.sin(theta)
        z = 2.5 + 0.38 * phi + 0.06
        
        size = random.uniform(0.035, 0.095)
        
        bpy.ops.mesh.primitive_uv_sphere_add(
            segments=10, ring_count=8, radius=size, location=(x, y, z)
        )
        wart = bpy.context.active_object
        if wart:
            wart.name = f"Wart{i:03d}"
            wart.scale.z = random.uniform(0.28, 0.42)
            wart_parts.append(wart.name)
    
    bpy.ops.object.select_all(action='DESELECT')
    first = None
    for name in wart_parts:
        obj = bpy.data.objects.get(name)
        if obj:
            obj.select_set(True)
            if first is None:
                first = obj
    
    if first:
        bpy.context.view_layer.objects.active = first
        bpy.ops.object.join()
        result = bpy.context.active_object
        if result:
            result.name = "Warts"
    
    return bpy.data.objects.get("Warts")


def create_stem():
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=32, radius=0.13, depth=2.4, location=(0, 0, 1.2)
    )
    obj = bpy.context.active_object
    if obj is None:
        return None
    obj.name = "Stem"
    
    bm = bmesh.new()
    bm.from_mesh(obj.data)
    
    for v in bm.verts:
        z_norm = (v.co.z + 1.2) / 2.4
        r = math.sqrt(v.co.x**2 + v.co.y**2)
        
        if r > 0.001:
            if z_norm < 0.18:
                scale = 1.0 + 1.8 * (0.18 - z_norm)
            elif z_norm > 0.92:
                scale = 1.0 + 0.25 * (z_norm - 0.92)
            else:
                scale = 1.0 - 0.12 * math.sin((z_norm - 0.18) * math.pi / 0.74)
            
            angle = math.atan2(v.co.y, v.co.x)
            irregular = 1.0 + 0.025 * math.sin(angle * 7 + z_norm * 4)
            v.co.x = r * scale * irregular * math.cos(angle)
            v.co.y = r * scale * irregular * math.sin(angle)
    
    bm.to_mesh(obj.data)
    bm.free()
    
    return bpy.data.objects.get("Stem")


def create_ring():
    bpy.ops.mesh.primitive_torus_add(
        major_radius=0.19, minor_radius=0.022,
        major_segments=36, minor_segments=12,
        location=(0, 0, 1.92)
    )
    obj = bpy.context.active_object
    if obj is None:
        return None
    obj.name = "Ring"
    
    bm = bmesh.new()
    bm.from_mesh(obj.data)
    
    for v in bm.verts:
        xy_dist = math.sqrt(v.co.x**2 + v.co.y**2)
        angle = math.atan2(v.co.y, v.co.x)
        droop = 0.035 * max(0, (xy_dist - 0.19) / 0.022)
        v.co.z -= droop
        v.co.z += 0.006 * math.sin(angle * 14)
        
        if xy_dist > 0.19:
            v.co.x *= 1.15
            v.co.y *= 1.15
    
    bm.to_mesh(obj.data)
    bm.free()
    
    ring_obj = bpy.data.objects.get("Ring")
    if ring_obj:
        ring_obj.scale = (1.0, 1.0, 0.55)
        bpy.context.view_layer.objects.active = ring_obj
        bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    
    return bpy.data.objects.get("Ring")


def create_volva():
    parts = []
    
    for i in range(4):
        bpy.ops.mesh.primitive_torus_add(
            major_radius=0.32 - i * 0.025,
            minor_radius=0.018,
            major_segments=28, minor_segments=10,
            location=(0, 0, 0.08 + i * 0.055)
        )
        obj = bpy.context.active_object
        if obj:
            obj.name = f"VolvaRing{i}"
            obj.scale.z = 0.65 + i * 0.08
            parts.append(obj.name)
    
    bpy.ops.mesh.primitive_cone_add(
        vertices=32, radius1=0.36, radius2=0.30, depth=0.12, location=(0, 0, 0.06)
    )
    cup = bpy.context.active_object
    if cup:
        cup.name = "VolvaCup"
        parts.append(cup.name)
    
    bpy.ops.object.select_all(action='DESELECT')
    first = None
    for name in parts:
        obj = bpy.data.objects.get(name)
        if obj:
            obj.select_set(True)
            if first is None:
                first = obj
    
    if first:
        bpy.context.view_layer.objects.active = first
        bpy.ops.object.join()
        result = bpy.context.active_object
        if result:
            result.name = "Volva"
    
    return bpy.data.objects.get("Volva")


def create_stem_fibers():
    fiber_parts = []
    num_fibers = 24
    
    for i in range(num_fibers):
        angle = 2 * math.pi * i / num_fibers
        radius = 0.135
        
        bpy.ops.mesh.primitive_cylinder_add(
            vertices=8, radius=0.004, depth=1.6, location=(0, 0, 1.3)
        )
        fiber = bpy.context.active_object
        if fiber:
            fiber.name = f"Fiber{i:03d}"
            fiber.location.x = radius * math.cos(angle)
            fiber.location.y = radius * math.sin(angle)
            fiber_parts.append(fiber.name)
    
    bpy.ops.object.select_all(action='DESELECT')
    first = None
    for name in fiber_parts:
        obj = bpy.data.objects.get(name)
        if obj:
            obj.select_set(True)
            if first is None:
                first = obj
    
    if first:
        bpy.context.view_layer.objects.active = first
        bpy.ops.object.join()
        result = bpy.context.active_object
        if result:
            result.name = "StemFibers"
    
    return bpy.data.objects.get("StemFibers")


def create_cap_margin():
    bpy.ops.mesh.primitive_torus_add(
        major_radius=1.2, minor_radius=0.035,
        major_segments=64, minor_segments=12,
        location=(0, 0, 2.52)
    )
    obj = bpy.context.active_object
    if obj:
        obj.name = "CapMargin"
        obj.scale.z = 0.4
        bpy.context.view_layer.objects.active = obj
        bpy.ops.object.transform_apply(location=False, rotation=False, scale=True)
    return bpy.data.objects.get("CapMargin")


def main():
    clear_scene()
    
    red_mat = create_material("RedCap", (0.82, 0.06, 0.02, 1.0), 0.32, 0.14)
    white_mat = create_material("White", (0.96, 0.93, 0.87, 1.0), 0.48, 0.18)
    gill_mat = create_material("GillWhite", (0.98, 0.95, 0.89, 1.0), 0.58, 0.1)
    stem_mat = create_material("StemWhite", (0.94, 0.91, 0.84, 1.0), 0.52, 0.16)
    
    create_cap()
    create_cap_margin()
    create_cap_flesh()
    create_gills()
    create_warts()
    create_stem()
    create_stem_fibers()
    create_ring()
    create_volva()
    
    assignments = [
        ("Cap", red_mat),
        ("CapMargin", red_mat),
        ("CapFlesh", white_mat),
        ("Gills", gill_mat),
        ("Warts", white_mat),
        ("Stem", stem_mat),
        ("StemFibers", stem_mat),
        ("Ring", white_mat),
        ("Volva", white_mat),
    ]
    
    for name, mat in assignments:
        obj = bpy.data.objects.get(name)
        if obj and obj.data:
            obj.data.materials.clear()
            obj.data.materials.append(mat)
    
    part_names = [
        "Cap", "CapMargin", "CapFlesh", "Gills", "Warts",
        "Stem", "StemFibers", "Ring", "Volva"
    ]
    
    bpy.ops.object.select_all(action='DESELECT')
    first = None
    for name in part_names:
        obj = bpy.data.objects.get(name)
        if obj:
            obj.select_set(True)
            if first is None:
                first = obj
    
    if first:
        bpy.context.view_layer.objects.active = first
        bpy.ops.object.join()
        result = bpy.context.active_object
        if result:
            result.name = "Amanita_muscaria"
            bpy.ops.object.shade_smooth()
            bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS')
    
    bpy.ops.object.select_all(action='DESELECT')
    mushroom = bpy.data.objects.get("Amanita_muscaria")
    if mushroom:
        mushroom.select_set(True)
        bpy.context.view_layer.objects.active = mushroom


if __name__ == "__main__":
    main()