import bpy
import bmesh
import mathutils
import random
import math

class SolidTreeGenerator:
    def __init__(self, levels=5, length=3.0, radius=0.3):
        self.levels = levels
        self.base_length = length
        self.base_radius = radius
        
        # Хранилище для данных: (start_pos, end_pos, radius_start, radius_end)
        self.branches_data = []

    def calculate_tree(self, start_pos, direction, length, radius, level):
        if level <= 0 or length < 0.1:
            return

        end_pos = start_pos + direction * length
        # Сохраняем данные сегмента
        self.branches_data.append({
            'start': start_pos.copy(),
            'end': end_pos.copy(),
            'r_start': radius,
            'r_end': radius * 0.7
        })

        # 1. Основной ствол (продолжение)
        trunk_dir = (direction + self.get_random_vector(0.1)).normalized()
        self.calculate_tree(end_pos, trunk_dir, length * 0.8, radius * 0.7, level - 1)

        # 2. Боковые ветки (ветвление)
        if level > 1:
            num_sides = random.randint(2, 3) # Минимум 2 ветки для видимости
            for _ in range(num_sides):
                # Создаем вектор, сильно отклоненный от ствола (30-60 градусов)
                axis = self.get_random_vector(1.0).normalized()
                angle = math.radians(random.uniform(30, 60))
                
                side_dir = direction.copy()
                side_dir.rotate(mathutils.Quaternion(axis, angle))
                
                # Боковые ветки короче
                self.calculate_tree(end_pos, side_dir, length * 0.6, radius * 0.5, level - 1)

    def get_random_vector(self, intensity):
        return mathutils.Vector((
            random.uniform(-intensity, intensity),
            random.uniform(-intensity, intensity),
            random.uniform(-intensity, intensity)
        ))

    def build_mesh(self):
        mesh = bpy.data.meshes.new("TreeMesh")
        obj = bpy.data.objects.new("Tree", mesh)
        bpy.context.collection.objects.link(obj)
        
        bm = bmesh.new()
        skin_layer = bm.verts.layers.skin.verify()
        
        # Словарь для предотвращения дублирования вершин в одной точке
        # Ключ - кортеж координат, Значение - объект вершины BMesh
        vert_map = {}

        for b in self.branches_data:
            # Превращаем координаты в кортежи для словаря
            s_key = tuple(round(v, 4) for v in b['start'])
            e_key = tuple(round(v, 4) for v in b['end'])
            
            # Получаем или создаем начальную вершину
            if s_key not in vert_map:
                v_start = bm.verts.new(b['start'])
                v_start[skin_layer].radius = (b['r_start'], b['r_start'])
                vert_map[s_key] = v_start
            else:
                v_start = vert_map[s_key]
                
            # Получаем или создаем конечную вершину
            if e_key not in vert_map:
                v_end = bm.verts.new(b['end'])
                v_end[skin_layer].radius = (b['r_end'], b['r_end'])
                vert_map[e_key] = v_end
            else:
                v_end = vert_map[e_key]
            
            # Создаем ребро, если его еще нет
            if not bm.edges.get((v_start, v_end)):
                bm.edges.new((v_start, v_end))

        # Находим корень (самую нижнюю точку) и помечаем его
        root_v = min(bm.verts, key=lambda v: v.co.z)
        root_v[skin_layer].use_root = True

        bm.to_mesh(mesh)
        bm.free()
        
        # Модификаторы
        obj.modifiers.new(name="Skin", type='SKIN')
        sub = obj.modifiers.new(name="Subdiv", type='SUBSURF')
        sub.levels = 1 # Для начала 1, чтобы не тормозило

# Очистка сцены
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()

# Запуск
generator = SolidTreeGenerator(levels=5, length=3.0, radius=0.4)
generator.calculate_tree(mathutils.Vector((0,0,0)), mathutils.Vector((0,0,1)), 3.0, 0.4, 5)
generator.build_mesh()