---

# -*- coding: utf-8 -*-
"""
КОНВЕЙЕР v4: в Blender работает как раньше, в Unity едет через Shape Keys (Blend Shapes).
"""
import math
import bpy
import bmesh
from mathutils import Vector, Euler

# ============================  НАСТРОЙКИ  ============================
BED_NAME = ""
BED_LENGTH = 4.5
BED_WIDTH  = 1.5
BED_THICK  = 0.10
BELT_THICK = 0.012
CLEARANCE  = 0.001
MIN_RADIUS = 0.01

RIB_HEIGHT     = 0.007
RIB_THICK_REL  = 0.16
RIB_PITCH_REL  = 0.24
RIB_MARGIN_REL = 0.015
RIB_ANGLE_DEG  = 45.0
RIB_ROUND      = 0.0025
DIRECTION      = 1

MAX_TILE_VERTS = 30000

# --- АНИМАЦИЯ / ЭКСПОРТ ---
PERIOD_FRAMES  = 48    # кадров на один шаг
LOOPS          = 1     # один оборот = 1 цикл (в Unity зациклишь)
BAKE_SHAPEKEYS = True  # True = запечь движение в Shape Keys для Unity
SK_STEP        = 2     # каждый N-й кадр сохраняем как шейпкей (2 = через кадр)

# =====================================================================

K = 0.5522847498
GENERATED = ("Conveyor_Root", "Belt_Path", "Belt_Plate", "Belt_Ribs")


def make_material(name, color, roughness):
    mat = bpy.data.materials.new(name)
    try:
        mat.use_nodes = True
        bsdf = mat.node_tree.nodes.get("Principled BSDF")
        if bsdf is not None:
            bsdf.inputs["Base Color"].default_value = (color[0], color[1], color[2], 1.0)
            bsdf.inputs["Roughness"].default_value = roughness
    except Exception as err:
        print("материал:", err)
    mat.diffuse_color = (color[0], color[1], color[2], 1.0)
    return mat


def finish_mesh(bm, me, material):
    for f in bm.faces:
        f.smooth = True
    bm.to_mesh(me)
    bm.free()
    if hasattr(me, "use_auto_smooth"):
        me.use_auto_smooth = True
    me.materials.append(material)
    me.update()


def pick_bed():
    if BED_NAME:
        ob = bpy.data.objects.get(BED_NAME)
        if ob is None or ob.type != 'MESH':
            return None
        return ob
    ob = bpy.context.active_object
    if ob is not None and ob.type == 'MESH' and ob.name not in GENERATED and ob.select_get():
        return ob
    return None


def fit_from_bed(ob):
    d = ob.dimensions
    lx, ly, lz = d.x, d.y, d.z
    mw = ob.matrix_world
    corners = [Vector(c) for c in ob.bound_box]
    center_local = sum(corners, Vector((0.0, 0.0, 0.0))) / 8.0
    center = mw @ center_local
    _, quat, _ = mw.decompose()
    if lx >= ly:
        length, width = lx, ly
    else:
        length, width = ly, lx
        quat = quat @ Euler((0.0, 0.0, math.radians(90.0)), 'XYZ').to_quaternion()
    return center, quat, length, width, lz


def build_path(L, R):
    cu = bpy.data.curves.new("Belt_Path", 'CURVE')
    cu.dimensions = '2D'
    cu.resolution_u = 64
    sp = cu.splines.new('BEZIER')
    sp.bezier_points.add(5)
    sp.use_cyclic_u = True

    s = L / 3.0
    kr = K * R
    pts = [
        ((-L / 2, R),      (-L / 2 - kr, R),  (-L / 2 + s, R)),
        ((L / 2, R),       (L / 2 - s, R),    (L / 2 + kr, R)),
        ((L / 2 + R, 0),   (L / 2 + R, kr),   (L / 2 + R, -kr)),
        ((L / 2, -R),      (L / 2 + kr, -R),  (L / 2 - s, -R)),
        ((-L / 2, -R),     (-L / 2 + s, -R),  (-L / 2 - kr, -R)),
        ((-L / 2 - R, 0),  (-L / 2 - R, -kr), (-L / 2 - R, kr)),
    ]
    for bp, (co, hl, hr) in zip(sp.bezier_points, pts):
        bp.handle_left_type = 'FREE'
        bp.handle_right_type = 'FREE'
        bp.co = (co[0], co[1], 0.0)
        bp.handle_left = (hl[0], hl[1], 0.0)
        bp.handle_right = (hr[0], hr[1], 0.0)

    obj = bpy.data.objects.new("Belt_Path", cu)
    obj.rotation_euler = (math.radians(90.0), 0.0, 0.0)
    return obj


def measure_curve(obj, fallback):
    try:
        bpy.context.view_layer.update()
        dg = bpy.context.evaluated_depsgraph_get()
        ev = obj.evaluated_get(dg)
        me = ev.to_mesh()
        total = 0.0
        if me is not None:
            for e in me.edges:
                a = me.vertices[e.vertices[0]].co
                b = me.vertices[e.vertices[1]].co
                total += (a - b).length
        ev.to_mesh_clear()
        if total > 0.8 * fallback:
            return total
    except Exception as err:
        print("measure_curve:", err)
    return fallback


def build_plate(length, W, TH, material, PLATE_SEG):
    me = bpy.data.meshes.new("Belt_Plate")
    bm = bmesh.new()
    n = max(2, int(math.ceil(length / PLATE_SEG)))
    hw = W / 2.0
    profile = [(-hw, 0.0), (hw, 0.0), (hw, TH), (-hw, TH)]
    rings = []
    for i in range(n + 1):
        x = length * i / n
        rings.append([bm.verts.new((x, y, z)) for (y, z) in profile])
    for i in range(n):
        a, b = rings[i], rings[i + 1]
        for k in range(4):
            bm.faces.new((a[k], a[(k + 1) % 4], b[(k + 1) % 4], b[k]))
    finish_mesh(bm, me, material)
    return me


def build_rib_tile(W, TH, pitch, seg, material):
    me = bpy.data.meshes.new("Belt_Ribs")
    bm = bmesh.new()

    hw = W / 2.0 - RIB_MARGIN_REL * W
    a = hw * math.tan(math.radians(RIB_ANGLE_DEG))
    t = RIB_THICK_REL * W
    span = a + t
    x_off = 0.0 if span >= pitch else (pitch - span) / 2.0

    embed = 0.3 * TH
    z0 = TH - embed
    z1 = TH + RIB_HEIGHT

    while True:
        ny = max(4, int(math.ceil(a / seg)))
        nx = max(2, int(math.ceil(t / seg)))
        if (2 * ny + 1) * (2 * nx + 2) <= MAX_TILE_VERTS:
            break
        seg *= 1.15

    rows = []
    for j in range(-ny, ny + 1):
        y = hw * j / ny
        d = a * abs(j) / ny
        xl = span - d
        xt = xl - t
        if DIRECTION < 0:
            xl, xt = span - xl, span - xt
        xl += x_off
        xt += x_off
        ring = []
        for i in range(nx + 1):
            ring.append(bm.verts.new((xl + (xt - xl) * i / nx, y, z0)))
        for i in range(nx, -1, -1):
            ring.append(bm.verts.new((xl + (xt - xl) * i / nx, y, z1)))
        rows.append(ring)

    M = len(rows[0])
    for j in range(len(rows) - 1):
        r0, r1 = rows[j], rows[j + 1]
        for k in range(M):
            bm.faces.new((r0[k], r0[(k + 1) % M], r1[(k + 1) % M], r1[k]))
    bm.faces.new(rows[0])
    bm.faces.new(rows[-1][::-1])
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    finish_mesh(bm, me, material)
    return me, seg


def add_mods(obj, path_obj, bevel_width, bevel_segments, pitch=None, count=None):
    if bevel_width > 0:
        bev = obj.modifiers.new("Bevel", 'BEVEL')
        bev.width = bevel_width
        bev.segments = bevel_segments
        bev.limit_method = 'ANGLE'
        bev.angle_limit = math.radians(35.0)
        try:
            bev.profile = 0.5
        except Exception:
            pass

    if pitch is not None:
        arr = obj.modifiers.new("Array", 'ARRAY')
        arr.fit_type = 'FIXED_COUNT'
        arr.count = count
        arr.use_relative_offset = False
        arr.use_constant_offset = True
        arr.constant_offset_displace = (pitch, 0.0, 0.0)

    cur = obj.modifiers.new("Curve", 'CURVE')
    cur.object = path_obj
    cur.deform_axis = 'POS_X'

    try:
        wn = obj.modifiers.new("WeightedNormal", 'WEIGHTED_NORMAL')
        wn.keep_sharp = True
        wn.weight = 50
    except Exception as err:
        print("WeightedNormal:", err)


# =====================  СБОРКА  =====================
scene = bpy.context.scene

bed = pick_bed()

if bed is not None:
    center, quat, bed_len, W, bed_thk = fit_from_bed(bed)
    print("Подгоняю под %r: длина %.3f, ширина %.3f, толщина %.3f" % (bed.name, bed_len, W, bed_thk))
    if bed_thk <= 0.0 or bed_thk > 0.35 * min(bed_len, W):
        bed_thk = BED_THICK
else:
    center, quat = None, None
    bed_len, W, bed_thk = BED_LENGTH, BED_WIDTH, BED_THICK

R = max(bed_thk / 2.0 + CLEARANCE, MIN_RADIUS)
TH = max(BELT_THICK, 0.45 * R)
L = bed_len - 2.0 * R
if L < 0.1:
    L = 0.1

PLATE_SEG = min(0.03, max(0.003, R * 0.2))

for nm in GENERATED:
    old = bpy.data.objects.get(nm)
    if old is not None:
        bpy.data.objects.remove(old, do_unlink=True)

coll = bpy.data.collections.get("Conveyor")
if coll is None:
    coll = bpy.data.collections.new("Conveyor")
if coll.name not in scene.collection.children:
    scene.collection.children.link(coll)

root = bpy.data.objects.new("Conveyor_Root", None)
root.empty_display_type = 'PLAIN_AXES'
coll.objects.link(root)
if center is not None:
    root.rotation_mode = 'QUATERNION'
    root.rotation_quaternion = quat
    root.location = center
else:
    root.location = (0.0, 0.0, R + TH)

path_obj = build_path(L, R)
coll.objects.link(path_obj)
path_obj.parent = root

ideal_len = 2 * L + 2 * math.pi * R
path_len = measure_curve(path_obj, ideal_len)

want_pitch = RIB_PITCH_REL * W
count = max(1, int(round(path_len / want_pitch)))
pitch = path_len / count

mat_plate = make_material("Belt_Rubber", (0.045, 0.045, 0.05), 0.65)
mat_rib = make_material("Belt_Ribs_Mat", (0.16, 0.16, 0.18), 0.55)

plate_obj = bpy.data.objects.new("Belt_Plate", build_plate(path_len + 0.01, W, TH, mat_plate, PLATE_SEG))
rib_mesh, seg_used = build_rib_tile(W, TH, pitch, PLATE_SEG, mat_rib)
ribs_obj = bpy.data.objects.new("Belt_Ribs", rib_mesh)

for o in (plate_obj, ribs_obj):
    coll.objects.link(o)
    o.parent = root

add_mods(plate_obj, path_obj, min(0.003, TH * 0.25), 2)
rib_round = min(RIB_ROUND, RIB_HEIGHT * 0.45)
add_mods(ribs_obj, path_obj, rib_round, 3, pitch=pitch, count=count)

# =====================  АНИМАЦИЯ (для Blender)  =====================
prefs = bpy.context.preferences.edit
old_interp = prefs.keyframe_new_interpolation_type
prefs.keyframe_new_interpolation_type = 'LINEAR'

last = 1 + PERIOD_FRAMES * LOOPS
for o in (plate_obj, ribs_obj):
    o.location.x = 0.0
    o.keyframe_insert(data_path="location", index=0, frame=1)
    o.location.x = pitch * LOOPS
    o.keyframe_insert(data_path="location", index=0, frame=last)
    o.location.x = 0.0

prefs.keyframe_new_interpolation_type = old_interp

scene.frame_start = 1
scene.frame_end = PERIOD_FRAMES * LOOPS

# =====================  ЗАПЕКАНИЕ В SHAPE KEYS  =====================
def bake_to_shapekeys(obj):
    """
    Проходит по кадрам анимации, для каждого кадра берёт деформированную
    Curve-модификатором геометрию и сохраняет как Shape Key.
    В Unity это выглядит как Blend Shapes, которые проигрываются в цикле.
    """
    scene.frame_set(1)
    bpy.context.view_layer.update()

    # Базовая форма — на кадре 1
    dg = bpy.context.evaluated_depsgraph_get()
    ev = obj.evaluated_get(dg)
    base_mesh = ev.to_mesh()
    base_coords = [v.co.copy() for v in base_mesh.vertices]
    ev.to_mesh_clear()

    # Создаём базовый shape key + Basis
    if obj.data.shape_keys is None:
        obj.shape_key_add(name="Basis", from_mix=False)
        basis = obj.data.shape_keys.key_blocks["Basis"]
        for i, co in enumerate(base_coords):
            basis.data[i].co = co
    else:
        # перезаписываем базис
        basis = obj.data.shape_keys.key_blocks.get("Basis")
        if basis is None:
            obj.shape_key_add(name="Basis", from_mix=False)
            basis = obj.data.shape_keys.key_blocks["Basis"]
        for i, co in enumerate(base_coords):
            basis.data[i].co = co

    # Проходим кадры
    frames = list(range(1, last + 1, SK_STEP))
    if frames[-1] != last:
        frames.append(last)

    for f in frames:
        scene.frame_set(f)
        bpy.context.view_layer.update()
        dg = bpy.context.evaluated_depsgraph_get()
        ev = obj.evaluated_get(dg)
        me = ev.to_mesh()

        sk = obj.shape_key_add(name="Frame_%03d" % f, from_mix=False)
        for i, v in enumerate(me.vertices):
            sk.data[i].co = v.co.copy()
        ev.to_mesh_clear()

    # Отключаем модификаторы у объекта (геометрия теперь в шейпкеях)
    # НЕ удаляем — оставляем, но выключаем, чтобы в Blender не было двойной деформации.
    for m in obj.modifiers:
        m.show_viewport = False
        m.show_render = False
        m.show_in_editmode = False

    # обнуляем location — базовое положение остаётся
    obj.location = (0.0, 0.0, 0.0)

    # удаляем fcurves локации — они больше не нужны
    ad = obj.animation_data
    if ad and ad.action:
        for fc in list(ad.action.fcurves):
            if fc.data_path == "location":
                ad.action.fcurves.remove(fc)

    print("Shape keys запечены: %d ключей у %s" % (len(obj.data.shape_keys.key_blocks), obj.name))


if BAKE_SHAPEKEYS:
    scene.frame_set(1)
    bake_to_shapekeys(plate_obj)
    bake_to_shapekeys(ribs_obj)

    # ---- анимация шейпкеев (чтобы в Blender тоже ехало) ----
    # ключи: [1..N], каждый активен 1 кадр, плавно
    for obj in (plate_obj, ribs_obj):
        if obj.data.shape_keys is None:
            continue
        kb = obj.data.shape_keys.key_blocks
        names = [k.name for k in kb if k.name.startswith("Frame_")]
        if not names:
            continue

        # базовое значение всех ключей = 0
        for k in kb:
            k.value = 0.0

        # каждый ключ активируем на своём кадре
        for i, name in enumerate(names):
            frame_num = int(name.split("_")[1])
            k = kb[name]
            k.value = 0.0
            k.keyframe_insert("value", frame=max(1, frame_num - SK_STEP))
            k.value = 1.0
            k.keyframe_insert("value", frame=frame_num)
            k.value = 0.0
            k.keyframe_insert("value", frame=min(last, frame_num + SK_STEP))

        # линейная интерполяция
        if obj.data.shape_keys.animation_data and obj.data.shape_keys.animation_data.action:
            for fc in obj.data.shape_keys.animation_data.action.fcurves:
                for kp in fc.keyframe_points:
                    kp.interpolation = 'LINEAR'

scene.frame_set(1)

print("=" * 60)
print("ГОТОВО!")
print("Радиус: %.3f, периметр: %.3f, зубцов: %d" % (R, path_len, count))
print()
print("BLENDER: нажми ПРОБЕЛ — конвейер поедет через шейпкеи.")
print()
print("UNITY:")
print("  1. Экспорт FBX: File → Export → FBX.")
print("  2. В FBX-экспортёре включи галочку: Geometry → Shape Keys (Blend Shapes).")
print("  3. Экспортируй коллекцию 'Conveyor'.")
print("  4. Импортируй FBX в Unity.")
print("  5. На меш-объекте (Belt_Plate / Belt_Ribs) выбери:")
print("     Import Settings → Blend Shape Normals = Calculate.")
print("     SkinnedMeshRenderer → BlendShapes — там будут все ключи Frame_001...")
print("  6. В Animation окне создай клип: анимируй blendShape.NAME с 0→1→0 по кадрам.")
print("  7. Loop Time = ON, скорость = по вкусу.")
print()
print("В Unity лента поедет перебором зубчиков, как в Blender.")