import bpy
import bmesh
import math
import os

# ============================================================
# NEON AVATAR — Realistic Female Avatar for Face Tracking
# ============================================================

# Clear scene
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()
for mesh in bpy.data.meshes:
    bpy.data.meshes.remove(mesh)

# --- HELPERS ---

def add_subsurf(obj, levels=2):
    mod = obj.modifiers.new(name="Subsurf", type='SUBSURF')
    mod.levels = levels
    mod.render_levels = levels

def shade_smooth(obj):
    for poly in obj.data.polygons:
        poly.use_smooth = True

def create_material(name, color, roughness=0.5, metallic=0.0):
    mat = bpy.data.materials.new(name=name)
    mat.use_nodes = True  # compat
    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    if bsdf:
        bsdf.inputs['Base Color'].default_value = (*color, 1.0)
        bsdf.inputs['Roughness'].default_value = roughness
        bsdf.inputs['Metallic'].default_value = metallic
    return mat

# --- HEAD ---
print("=== Creating head ===")
bpy.ops.mesh.primitive_uv_sphere_add(segments=64, ring_count=32, radius=0.5, location=(0, 0, 0))
head = bpy.context.active_object
head.name = "Head"

# Shape head
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.mesh.select_all(action='SELECT')
bpy.ops.transform.resize(value=(0.78, 1.0, 0.85))
bpy.ops.mesh.subdivide(number_cuts=3)

# Get bmesh in edit mode
bm = bmesh.from_edit_mesh(head.data)
for v in bm.verts:
    # Narrower jaw
    if v.co.z < -0.15:
        v.co.x *= 0.85
        v.co.y *= 0.9
    # Rounder forehead
    if v.co.z > 0.2:
        v.co.x *= 0.95
        v.co.y *= 1.05
    # Pointed chin
    if v.co.z < -0.35:
        v.co.x *= 0.6
        v.co.y *= 0.8
        v.co.z *= 1.1

bmesh.update_edit_mesh(head.data)
bpy.ops.object.mode_set(mode='OBJECT')

add_subsurf(head, 2)
shade_smooth(head)

# --- EYES ---
print("=== Creating eyes ===")
for side, x in [("Left", -0.15), ("Right", 0.15)]:
    bpy.ops.mesh.primitive_uv_sphere_add(segments=32, ring_count=16, radius=0.06, location=(x, -0.35, 0.05))
    eye = bpy.context.active_object
    eye.name = f"Eye_{side}"
    shade_smooth(eye)
    
    bpy.ops.mesh.primitive_uv_sphere_add(segments=16, ring_count=8, radius=0.025, location=(x, -0.40, 0.05))
    iris = bpy.context.active_object
    iris.name = f"Iris_{side}"
    iris_mat = create_material(f"Iris_{side}", (0.2, 0.4, 0.8), 0.1, 0.3)
    iris.data.materials.append(iris_mat)
    
    bpy.ops.mesh.primitive_uv_sphere_add(segments=8, ring_count=4, radius=0.012, location=(x, -0.41, 0.05))
    pupil = bpy.context.active_object
    pupil.name = f"Pupil_{side}"
    pupil_mat = create_material(f"Pupil_{side}", (0.05, 0.05, 0.05), 0.0)
    pupil.data.materials.append(pupil_mat)

# Skin
skin_mat = create_material("Skin", (0.85, 0.7, 0.6), 0.6)
head.data.materials.append(skin_mat)

# --- BODY ---
print("=== Creating body ===")
bpy.ops.mesh.primitive_cube_add(size=1, location=(0, 0, -0.7))
torso = bpy.context.active_object
torso.name = "Torso"
torso.scale = (0.22, 0.13, 0.35)
bpy.ops.object.transform_apply(scale=True)

bpy.ops.object.mode_set(mode='EDIT')
bm = bmesh.from_edit_mesh(torso.data)
for v in bm.verts:
    if abs(v.co.z) < 0.05:
        v.co.x *= 0.85
    if v.co.z < -0.15:
        v.co.x *= 1.1
bmesh.update_edit_mesh(torso.data)
bpy.ops.object.mode_set(mode='OBJECT')

add_subsurf(torso, 2)
shade_smooth(torso)
torso.data.materials.append(skin_mat)

for side, x in [("Left", -0.28), ("Right", 0.28)]:
    bpy.ops.mesh.primitive_cylinder_add(vertices=16, radius=0.05, depth=0.45, location=(x, 0, -0.75))
    arm = bpy.context.active_object
    arm.name = f"Arm_{side}"
    arm.rotation_euler = (math.radians(15), 0, math.radians(5 if side == "Left" else -5))
    shade_smooth(arm)
    arm.data.materials.append(skin_mat)

for side, x in [("Left", -0.1), ("Right", 0.1)]:
    bpy.ops.mesh.primitive_cylinder_add(vertices=16, radius=0.07, depth=0.7, location=(x, 0, -1.2))
    leg = bpy.context.active_object
    leg.name = f"Leg_{side}"
    shade_smooth(leg)
    leg.data.materials.append(skin_mat)

# --- HAIR ---
print("=== Creating hair ===")
bpy.ops.mesh.primitive_uv_sphere_add(segments=48, ring_count=24, radius=0.52, location=(0, 0.02, 0))
hair = bpy.context.active_object
hair.name = "Hair"

bpy.ops.object.mode_set(mode='EDIT')
bm = bmesh.from_edit_mesh(hair.data)
for v in bm.verts:
    if v.co.z < -0.1:
        v.co.z = -0.1
    if v.co.y > 0.3:
        v.co.y *= 1.3
        v.co.z -= 0.05
bmesh.update_edit_mesh(hair.data)
bpy.ops.object.mode_set(mode='OBJECT')

add_subsurf(hair, 2)
shade_smooth(hair)
hair_mat = create_material("Hair", (0.05, 0.05, 0.05), 0.3, 0.1)
hair.data.materials.append(hair_mat)

# --- CLOTHING ---
print("=== Creating clothing ===")
bpy.ops.mesh.primitive_cube_add(size=1, location=(0, 0.02, -0.7))
jacket = bpy.context.active_object
jacket.name = "Jacket"
jacket.scale = (0.26, 0.16, 0.38)
bpy.ops.object.transform_apply(scale=True)
add_subsurf(jacket, 1)
shade_smooth(jacket)
jacket_mat = create_material("Jacket", (0.15, 0.15, 0.2), 0.8)
jacket.data.materials.append(jacket_mat)

# --- ARKIT BLENDSHAPES ---
print("=== Creating ARKit blendshapes ===")
arkit_shapes = [
    "EyeBlinkLeft", "EyeLookInLeft", "EyeLookOutLeft", "EyeLookUpLeft", "EyeLookDownLeft",
    "EyeBlinkRight", "EyeLookInRight", "EyeLookOutRight", "EyeLookUpRight", "EyeLookDownRight",
    "EyeSquintLeft", "EyeSquintRight", "EyeWideLeft", "EyeWideRight",
    "JawOpen", "JawLeft", "JawRight", "JawForward",
    "MouthClose", "MouthFunnel", "MouthPucker", "MouthLeft", "MouthRight",
    "MouthSmileLeft", "MouthSmileRight", "MouthFrownLeft", "MouthFrownRight",
    "MouthDimpleLeft", "MouthDimpleRight", "MouthStretchLeft", "MouthStretchRight",
    "MouthRollLower", "MouthRollUpper", "MouthShrugLower", "MouthShrugUpper",
    "MouthPressLeft", "MouthPressRight", "MouthLowerOuterLeft", "MouthLowerOuterRight",
    "MouthUpperOuterLeft", "MouthUpperOuterRight",
    "CheekSquintLeft", "CheekSquintRight", "CheekPuff",
    "NoseSneerLeft", "NoseSneerRight",
    "BrowInnerUp", "BrowInnerDown",
    "BrowOuterUpLeft", "BrowOuterUpRight", "BrowOuterDownLeft", "BrowOuterDownRight",
]

head.shape_key_add(name="Basis", from_mix=False)

for shape_name in arkit_shapes:
    sk = head.shape_key_add(name=shape_name, from_mix=False)
    sk.value = 0.0
    
    verts = head.data.vertices
    n = len(verts)
    
    if "Blink" in shape_name:
        for i, v in enumerate(verts):
            if "Left" in shape_name and v.co.x < 0 and v.co.y < -0.3 and 0.0 < v.co.z < 0.15:
                sk.data[i].co.y += 0.02
            elif "Right" in shape_name and v.co.x > 0 and v.co.y < -0.3 and 0.0 < v.co.z < 0.15:
                sk.data[i].co.y += 0.02
    
    elif "Smile" in shape_name:
        for i, v in enumerate(verts):
            if "Left" in shape_name and v.co.x < -0.05 and v.co.y < -0.35 and v.co.z < -0.1:
                sk.data[i].co.y -= 0.015
                sk.data[i].co.x -= 0.01
            elif "Right" in shape_name and v.co.x > 0.05 and v.co.y < -0.35 and v.co.z < -0.1:
                sk.data[i].co.y -= 0.015
                sk.data[i].co.x += 0.01
    
    elif "JawOpen" in shape_name:
        for i, v in enumerate(verts):
            if v.co.z < -0.25 and v.co.y < -0.2:
                sk.data[i].co.z -= 0.05
    
    elif "Brow" in shape_name and "Up" in shape_name:
        for i, v in enumerate(verts):
            if v.co.y < -0.35 and v.co.z > 0.1:
                if "Left" in shape_name and v.co.x < 0:
                    sk.data[i].co.z += 0.015
                elif "Right" in shape_name and v.co.x > 0:
                    sk.data[i].co.z += 0.015
                elif "Inner" in shape_name:
                    sk.data[i].co.z += 0.012

print(f"Created {len(arkit_shapes)} blendshapes")

# --- IDLE ANIMATIONS ---
print("=== Creating idle animations ===")
scene = bpy.context.scene
scene.frame_start = 1
scene.frame_end = 120

# Breathing
torso.animation_data_create()
action = bpy.data.actions.new(name="Idle_Breathe")
torso.animation_data.action = action
for frame in range(1, 121):
    breath = math.sin(frame * 2 * math.pi / 96) * 0.015
    torso.scale[2] = 0.35 + breath
    torso.keyframe_insert(data_path="scale", index=2, frame=frame)

# Blinking
head.animation_data_create()
blink_action = bpy.data.actions.new(name="Idle_Anim")
head.animation_data.action = blink_action

blink_left = head.data.shape_keys.key_blocks.get("EyeBlinkLeft")
blink_right = head.data.shape_keys.key_blocks.get("EyeBlinkRight")

if blink_left and blink_right:
    for blink_frame in [30, 75]:
        for offset in range(-3, 5):
            f = blink_frame + offset
            if f < 1 or f > 120:
                continue
            if offset == 0:
                blink_left.value = 1.0
                blink_right.value = 1.0
            elif abs(offset) == 1:
                blink_left.value = 0.5
                blink_right.value = 0.5
            elif abs(offset) == 2:
                blink_left.value = 0.1
                blink_right.value = 0.1
            else:
                blink_left.value = 0.0
                blink_right.value = 0.0
            blink_left.keyframe_insert(data_path="value", frame=f)
            blink_right.keyframe_insert(data_path="value", frame=f)

# Weight shift
for frame in range(1, 121):
    head.location.x = math.sin(frame * 2 * math.pi / 120) * 0.01
    head.location.y = math.sin(frame * 2 * math.pi / 80) * 0.005
    head.keyframe_insert(data_path="location", frame=frame)

# Hair sway
hair.animation_data_create()
hair_action = bpy.data.actions.new(name="Idle_Hair")
hair.animation_data.action = hair_action
for frame in range(1, 121):
    hair.rotation_euler[1] = math.sin(frame * 2 * math.pi / 60) * 0.005
    hair.keyframe_insert(data_path="rotation_euler", index=1, frame=frame)

# --- CAMERA ---
print("=== Setting up camera ===")
bpy.ops.object.camera_add(location=(0, -1.2, 0.1))
cam = bpy.context.active_object
cam.name = "AvatarCamera"
cam.rotation_euler = (math.radians(88), 0, 0)
scene.camera = cam

# --- LIGHTING ---
print("=== Setting up lighting ===")
bpy.ops.object.light_add(type='AREA', location=(0.3, -0.8, 0.5))
key_light = bpy.context.active_object
key_light.data.energy = 200
key_light.data.size = 0.5
key_light.rotation_euler = (math.radians(60), math.radians(20), 0)

bpy.ops.object.light_add(type='AREA', location=(-0.3, -0.6, 0.2))
fill_light = bpy.context.active_object
fill_light.data.energy = 80
fill_light.data.size = 0.8

bpy.ops.object.light_add(type='AREA', location=(0, 0.5, 0.6))
back_light = bpy.context.active_object
back_light.data.energy = 100
back_light.data.size = 0.5

# --- RENDER ---
scene.render.engine = 'CYCLES'
scene.cycles.samples = 32
scene.cycles.use_denoising = True
scene.render.resolution_x = 1280
scene.render.resolution_y = 720
scene.render.fps = 30

world = bpy.data.worlds["World"]
world.use_nodes = True
bg = world.node_tree.nodes.get("Background")
if bg:
    bg.inputs[0].default_value = (0.05, 0.05, 0.08, 1.0)
    bg.inputs[1].default_value = 0.5

# --- SAVE ---
output = "/home/natiris/Dokumente/neon-avatar/avatar.blend"
bpy.ops.wm.save_as_mainfile(filepath=output)
print(f"=== Saved: {output} ===")

# Export FBX
fbx = "/home/natiris/Dokumente/neon-avatar/avatar.fbx"
bpy.ops.export_scene.fbx(
    filepath=fbx,
    use_mesh_modifiers=True,
    add_leaf_bones=False,
    bake_anim=True,
    bake_anim_step=1,
)
print(f"=== Exported: {fbx} ===")

# Summary
print(f"\n=== SUMMARY ===")
print(f"Head: {len(head.data.vertices)} verts, {len(head.data.shape_keys.key_blocks)} shape keys")
print(f"Blendshapes: {len(arkit_shapes)} ARKit")
print(f"Animations: Breathing, Blinking, Weight Shift, Hair Sway")
print(f"Materials: Skin, Hair(black), Eyes(blue), Jacket")
