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

random.seed(2026)

def safe_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 create_material_safe(name, base_color, roughness=0.8, subsurface_weight=0.0, specular=0.5):
    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    bsdf = nodes.get("Principled BSDF")
    if bsdf is not None:
        try:
            bsdf.inputs["Base Color"].default_value = base_color
        except Exception:
            pass
        try:
            bsdf.inputs["Roughness"].default_value = float(roughness)
        except Exception:
            pass
        try:
            bsdf.inputs["Subsurface Weight"].default_value = float(subsurface_weight)
        except Exception:
            pass
        try:
            bsdf.inputs["Specular IOR Level"].default_value = float(specular)
        except Exception:
            try:
                bsdf.inputs["Specular"].default_value = float(specular)
            except Exception:
                pass
    return mat

def add_noise_to_verts(bm, scale=0.02):
    for v in bm.verts:
        v.co.x += random.uniform(-scale, scale)
        v.co.y += random.uniform(-scale, scale)
        v.co.z += random.uniform(-scale * 0.5, scale * 0.5)

def create_caudex():
    bpy.ops.mesh.primitive_cylinder_add(
        vertices=48,
        radius=0.55,
        depth=0.35,
        location=(0, 0, 0.175),
        end_fill_type='NGON'
    )
    caudex = bpy.context.active_object
    if caudex is None:
        return None
    caudex.name = "Welwitschia_Caudex"
    
    bpy.ops.object.mode_set(mode='EDIT')
    bm = bmesh.from_edit_mesh(caudex.data)
    bm.verts.ensure_lookup_table()
    
    for v in bm.verts:
        dist = math.sqrt(v.co.x**2 + v.co.y**2)
        if v.co.z > 0.25:
            normalized_dist = dist / 0.55
            depression = 0.08 * (1.0 - normalized_dist**1.5)
            v.co.z -= depression
            angle = math.atan2(v.co.y, v.co.x)
            ridge = 0.02 * math.sin(angle * 6) * (1.0 - normalized_dist)
            v.co.z += ridge
        if v.co.z < 0.1:
            scale_factor = 1.0 + 0.25 * (0.175 - v.co.z) / 0.175
            angle = math.atan2(v.co.y, v.co.x)
            irregular = 1.0 + 0.08 * math.sin(angle * 5 + random.uniform(0, 0.5))
            v.co.x *= scale_factor * irregular
            v.co.y *= scale_factor * irregular
        v.co.x += random.uniform(-0.015, 0.015)
        v.co.y += random.uniform(-0.015, 0.015)
        v.co.z += random.uniform(-0.008, 0.008)
    
    bmesh.update_edit_mesh(caudex.data)
    bpy.ops.object.mode_set(mode='OBJECT')
    
    caudex = bpy.data.objects.get("Welwitschia_Caudex")
    if caudex is None:
        return None
    
    try:
        subsurf = caudex.modifiers.new(name="Subdivision", type='SUBSURF')
        subsurf.levels = 2
        subsurf.render_levels = 2
    except Exception:
        pass
    
    bark_mat = create_material_safe(
        "Caudex_Bark",
        (0.18, 0.12, 0.07, 1.0),
        roughness=0.95,
        subsurface_weight=0.02
    )
    caudex.data.materials.append(bark_mat)
    
    return caudex

def create_leaf_segment(length, base_width, tip_width, segments_count=40):
    mesh = bpy.data.meshes.new("LeafSegment_Mesh")
    obj = bpy.data.objects.new("LeafSegment", mesh)
    bpy.context.collection.objects.link(obj)
    
    bm = bmesh.new()
    
    verts_left = []
    verts_right = []
    
    for i in range(segments_count + 1):
        t = i / segments_count
        x = t * length
        width = base_width + (tip_width - base_width) * t
        droop = -0.4 * (t ** 1.8) * length
        curl_z = 0.08 * math.sin(t * math.pi * 1.5)
        wave_y = 0.03 * math.sin(t * math.pi * 3)
        twist = t * 0.15
        
        for side in [-1, 1]:
            y_base = side * width * 0.5
            y = y_base + wave_y * side
            z = droop + curl_z + y_base * twist * 0.3
            v = bm.verts.new((x, y, z))
            if side == -1:
                verts_left.append(v)
            else:
                verts_right.append(v)
    
    bm.verts.ensure_lookup_table()
    
    for i in range(segments_count):
        try:
            face_verts = [
                verts_left[i],
                verts_right[i],
                verts_right[i + 1],
                verts_left[i + 1]
            ]
            bm.faces.new(face_verts)
        except Exception:
            pass
    
    bm.to_mesh(mesh)
    bm.free()
    
    return obj

def create_leaf_strip(base_angle, strip_offset, length, width, age_damage=0.5):
    angle = base_angle + strip_offset
    
    tip_width_factor = 0.2 + random.uniform(0, 0.3) * (1.0 - age_damage)
    leaf = create_leaf_segment(
        length=length,
        base_width=width,
        tip_width=width * tip_width_factor,
        segments_count=35
    )
    
    if leaf is None:
        return None
    
    base_radius = 0.5
    base_x = base_radius * math.cos(angle)
    base_y = base_radius * math.sin(angle)
    base_z = 0.28
    
    leaf.location = (base_x, base_y, base_z)
    leaf.rotation_euler = Euler((
        random.uniform(-0.15, 0.05),
        random.uniform(-0.08, 0.08),
        angle
    ), 'XYZ')
    
    return leaf

def create_welwitschia_leaves():
    leaves = []
    
    for leaf_idx in range(2):
        base_angle = leaf_idx * math.pi
        num_strips = random.randint(5, 9)
        spread = 0.6
        
        for strip_idx in range(num_strips):
            offset = (strip_idx - (num_strips - 1) / 2) * (spread / num_strips)
            length = random.uniform(1.8, 2.8)
            width = random.uniform(0.06, 0.12)
            age = random.uniform(0.3, 1.0)
            
            leaf = create_leaf_strip(base_angle, offset, length, width, age)
            if leaf is not None:
                leaf.name = f"Leaf{leaf_idx}_Strip{strip_idx}"
                leaves.append(leaf)
    
    leaf_mat = create_material_safe(
        "Welwitschia_Leaf",
        (0.15, 0.28, 0.12, 1.0),
        roughness=0.65,
        subsurface_weight=0.15
    )
    
    dead_tip_mat = create_material_safe(
        "Leaf_DeadTip",
        (0.35, 0.28, 0.18, 1.0),
        roughness=0.85,
        subsurface_weight=0.0
    )
    
    for leaf in leaves:
        if leaf is not None and leaf.data is not None:
            leaf.data.materials.append(leaf_mat)
    
    return leaves

def create_strobilus(is_male=True):
    mesh = bpy.data.meshes.new("Strobilus_Mesh")
    obj = bpy.data.objects.new("Strobilus", mesh)
    bpy.context.collection.objects.link(obj)
    
    bm = bmesh.new()
    
    if is_male:
        height = random.uniform(0.04, 0.06)
        radius1 = random.uniform(0.012, 0.018)
        radius2 = radius1 * 0.7
    else:
        height = random.uniform(0.06, 0.1)
        radius1 = random.uniform(0.025, 0.04)
        radius2 = radius1 * 0.85
    
    bmesh.ops.create_cone(
        bm,
        segments=12,
        radius1=radius1,
        radius2=radius2,
        depth=height,
        cap_ends=True,
        cap_tris=False
    )
    
    for v in bm.verts:
        if abs(v.co.z) < height * 0.4:
            angle = math.atan2(v.co.y, v.co.x)
            bump = 0.003 * math.sin(angle * 8)
            dist = math.sqrt(v.co.x**2 + v.co.y**2)
            if dist > 0:
                v.co.x += bump * v.co.x / dist
                v.co.y += bump * v.co.y / dist
    
    bm.to_mesh(mesh)
    bm.free()
    
    return obj

def create_cone_stem(height, radius=0.006):
    mesh = bpy.data.meshes.new("ConeStem_Mesh")
    obj = bpy.data.objects.new("ConeStem", mesh)
    bpy.context.collection.objects.link(obj)
    
    bm = bmesh.new()
    bmesh.ops.create_cone(
        bm,
        segments=8,
        radius1=radius,
        radius2=radius * 0.8,
        depth=height,
        cap_ends=True,
        cap_tris=False
    )
    bm.to_mesh(mesh)
    bm.free()
    
    return obj

def create_reproductive_structures():
    structures = []
    is_male = random.choice([True, False])
    
    num_stalks = random.randint(2, 5)
    
    cone_mat = create_material_safe(
        "Strobilus_Mat",
        (0.55, 0.25, 0.15, 1.0) if is_male else (0.45, 0.32, 0.18, 1.0),
        roughness=0.8
    )
    
    stem_mat = create_material_safe(
        "ConeStem_Mat",
        (0.22, 0.18, 0.1, 1.0),
        roughness=0.85
    )
    
    for stalk_idx in range(num_stalks):
        angle = random.uniform(0, 2 * math.pi)
        radius = random.uniform(0.15, 0.35)
        
        main_stem_height = random.uniform(0.12, 0.25)
        main_stem = create_cone_stem(main_stem_height, 0.008)
        if main_stem is None:
            continue
            
        main_stem.name = f"MainStem_{stalk_idx}"
        base_x = radius * math.cos(angle)
        base_y = radius * math.sin(angle)
        base_z = 0.3
        main_stem.location = (base_x, base_y, base_z + main_stem_height * 0.5)
        main_stem.rotation_euler.x = random.uniform(-0.15, 0.15)
        main_stem.rotation_euler.y = random.uniform(-0.15, 0.15)
        main_stem.data.materials.append(stem_mat)
        structures.append(main_stem)
        
        num_cones = random.randint(2, 5)
        for cone_idx in range(num_cones):
            branch_angle = random.uniform(0, 2 * math.pi)
            branch_tilt = random.uniform(0.3, 0.8)
            branch_height = random.uniform(0.03, 0.08)
            
            branch = create_cone_stem(branch_height, 0.005)
            if branch is None:
                continue
            branch.name = f"Branch_{stalk_idx}_{cone_idx}"
            
            branch_x = base_x + 0.02 * math.cos(branch_angle)
            branch_y = base_y + 0.02 * math.sin(branch_angle)
            branch_z = base_z + main_stem_height * (0.5 + 0.4 * cone_idx / num_cones)
            
            branch.location = (branch_x, branch_y, branch_z)
            branch.rotation_euler = Euler((branch_tilt, 0, branch_angle), 'XYZ')
            branch.data.materials.append(stem_mat)
            structures.append(branch)
            
            cone = create_strobilus(is_male)
            if cone is None:
                continue
            cone.name = f"Cone_{stalk_idx}_{cone_idx}"
            
            cone_offset = branch_height * 0.5
            cone_x = branch_x + cone_offset * math.sin(branch_tilt) * math.cos(branch_angle)
            cone_y = branch_y + cone_offset * math.sin(branch_tilt) * math.sin(branch_angle)
            cone_z = branch_z + cone_offset * math.cos(branch_tilt)
            
            cone.location = (cone_x, cone_y, cone_z)
            cone.rotation_euler = Euler((branch_tilt, 0, branch_angle), 'XYZ')
            cone.data.materials.append(cone_mat)
            structures.append(cone)
    
    return structures

def create_taproot():
    roots = []
    
    mesh = bpy.data.meshes.new("Taproot_Mesh")
    taproot = bpy.data.objects.new("Taproot", mesh)
    bpy.context.collection.objects.link(taproot)
    
    bm = bmesh.new()
    
    segments = 20
    rings = 12
    top_radius = 0.35
    bottom_radius = 0.06
    depth = 1.2
    
    vert_rings = []
    for ring_idx in range(rings + 1):
        t = ring_idx / rings
        z = -t * depth
        radius = top_radius + (bottom_radius - top_radius) * (t ** 0.7)
        
        ring_verts = []
        for seg_idx in range(segments):
            angle = 2 * math.pi * seg_idx / segments
            wobble = 0.03 * math.sin(angle * 3 + t * 5)
            r = radius + wobble
            x = r * math.cos(angle)
            y = r * math.sin(angle)
            z_offset = 0.02 * math.sin(angle * 4) * (1 - t)
            v = bm.verts.new((x, y, z + z_offset))
            ring_verts.append(v)
        vert_rings.append(ring_verts)
    
    bm.verts.ensure_lookup_table()
    
    for ring_idx in range(rings):
        for seg_idx in range(segments):
            next_seg = (seg_idx + 1) % segments
            try:
                v1 = vert_rings[ring_idx][seg_idx]
                v2 = vert_rings[ring_idx][next_seg]
                v3 = vert_rings[ring_idx + 1][next_seg]
                v4 = vert_rings[ring_idx + 1][seg_idx]
                bm.faces.new([v1, v2, v3, v4])
            except Exception:
                pass
    
    try:
        top_center = bm.verts.new((0, 0, 0))
        for seg_idx in range(segments):
            next_seg = (seg_idx + 1) % segments
            bm.faces.new([top_center, vert_rings[0][seg_idx], vert_rings[0][next_seg]])
    except Exception:
        pass
    
    try:
        bottom_center = bm.verts.new((0, 0, -depth))
        for seg_idx in range(segments):
            next_seg = (seg_idx + 1) % segments
            bm.faces.new([bottom_center, vert_rings[-1][next_seg], vert_rings[-1][seg_idx]])
    except Exception:
        pass
    
    bm.to_mesh(mesh)
    bm.free()
    
    root_mat = create_material_safe(
        "Root_Mat",
        (0.25, 0.15, 0.08, 1.0),
        roughness=0.92
    )
    taproot.data.materials.append(root_mat)
    roots.append(taproot)
    
    num_lateral = random.randint(4, 8)
    for lat_idx in range(num_lateral):
        lat_mesh = bpy.data.meshes.new(f"LateralRoot_{lat_idx}_Mesh")
        lat_root = bpy.data.objects.new(f"LateralRoot_{lat_idx}", lat_mesh)
        bpy.context.collection.objects.link(lat_root)
        
        bm = bmesh.new()
        
        lat_angle = 2 * math.pi * lat_idx / num_lateral + random.uniform(-0.2, 0.2)
        lat_length = random.uniform(0.3, 0.6)
        lat_start_radius = random.uniform(0.025, 0.04)
        lat_depth = random.uniform(0.15, 0.4)
        
        bmesh.ops.create_cone(
            bm,
            segments=8,
            radius1=lat_start_radius,
            radius2=lat_start_radius * 0.2,
            depth=lat_length,
            cap_ends=True,
            cap_tris=False
        )
        
        bm.to_mesh(lat_mesh)
        bm.free()
        
        start_x = 0.3 * math.cos(lat_angle)
        start_y = 0.3 * math.sin(lat_angle)
        start_z = -lat_depth
        
        lat_root.location = (
            start_x + lat_length * 0.4 * math.cos(lat_angle),
            start_y + lat_length * 0.4 * math.sin(lat_angle),
            start_z - lat_length * 0.2
        )
        lat_root.rotation_euler = Euler((
            math.pi * 0.5 + random.uniform(0.2, 0.5),
            0,
            lat_angle
        ), 'XYZ')
        
        lat_root.data.materials.append(root_mat)
        roots.append(lat_root)
    
    return roots

def create_ground_debris():
    debris = []
    
    num_dead_tips = random.randint(3, 6)
    
    debris_mat = create_material_safe(
        "Debris_Mat",
        (0.3, 0.22, 0.12, 1.0),
        roughness=0.9
    )
    
    for i in range(num_dead_tips):
        mesh = bpy.data.meshes.new(f"Debris_{i}_Mesh")
        obj = bpy.data.objects.new(f"Debris_{i}", mesh)
        bpy.context.collection.objects.link(obj)
        
        bm = bmesh.new()
        
        length = random.uniform(0.15, 0.35)
        width = random.uniform(0.02, 0.05)
        
        verts = []
        segs = 8
        for j in range(segs + 1):
            t = j / segs
            x = t * length
            curl = 0.05 * math.sin(t * math.pi)
            for side in [-1, 1]:
                w = width * (1 - t * 0.7)
                v = bm.verts.new((x, side * w * 0.5, curl))
                verts.append(v)
        
        bm.verts.ensure_lookup_table()
        
        for j in range(segs):
            try:
                bm.faces.new([verts[j*2], verts[j*2+1], verts[j*2+3], verts[j*2+2]])
            except Exception:
                pass
        
        bm.to_mesh(mesh)
        bm.free()
        
        angle = random.uniform(0, 2 * math.pi)
        dist = random.uniform(0.6, 1.2)
        obj.location = (
            dist * math.cos(angle),
            dist * math.sin(angle),
            0.01
        )
        obj.rotation_euler = Euler((
            random.uniform(-0.1, 0.1),
            random.uniform(-0.1, 0.1),
            random.uniform(0, 2 * math.pi)
        ), 'XYZ')
        
        obj.data.materials.append(debris_mat)
        debris.append(obj)
    
    return debris

def create_welwitschia_mirabilis():
    safe_clear_scene()
    
    all_parts = []
    
    caudex = create_caudex()
    if caudex is not None:
        all_parts.append(caudex)
    
    leaves = create_welwitschia_leaves()
    all_parts.extend([l for l in leaves if l is not None])
    
    repro = create_reproductive_structures()
    all_parts.extend([r for r in repro if r is not None])
    
    roots = create_taproot()
    all_parts.extend([r for r in roots if r is not None])
    
    debris = create_ground_debris()
    all_parts.extend([d for d in debris if d is not None])
    
    caudex = bpy.data.objects.get("Welwitschia_Caudex")
    if caudex is not None:
        for part in all_parts:
            if part is not None and part != caudex:
                try:
                    part.parent = caudex
                    part.matrix_parent_inverse = caudex.matrix_world.inverted()
                except Exception:
                    pass
    
    bpy.ops.object.select_all(action='DESELECT')
    for part in all_parts:
        if part is not None:
            try:
                part.select_set(True)
            except Exception:
                pass
    
    if caudex is not None:
        try:
            bpy.context.view_layer.objects.active = caudex
        except Exception:
            pass
    
    return all_parts

def setup_scene():
    try:
        bpy.ops.object.light_add(type='SUN', location=(5, -5, 10))
        sun = bpy.context.active_object
        if sun is not None:
            sun.name = "Sun_Light"
            sun.data.energy = 3.0
            sun.rotation_euler = Euler((math.radians(45), math.radians(15), math.radians(45)), 'XYZ')
    except Exception:
        pass
    
    try:
        bpy.ops.object.camera_add(location=(3.5, -3.5, 1.8))
        camera = bpy.context.active_object
        if camera is not None:
            camera.name = "Camera"
            camera.rotation_euler = Euler((math.radians(72), 0, math.radians(45)), 'XYZ')
            bpy.context.scene.camera = camera
    except Exception:
        pass

if __name__ == "__main__":
    create_welwitschia_mirabilis()
    setup_scene()