# scripts/generate_scene.py
import bpy
import json
import sys
import os
import mathutils

def clear_scene():
    bpy.ops.wm.read_factory_settings(use_empty=True)

def make_material(name, base_color=(1,1,1,1), texture_path=None):
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    nodes.clear()

    output = nodes.new(type="ShaderNodeOutputMaterial")
    principled = nodes.new(type="ShaderNodeBsdfPrincipled")
    principled.inputs['Base Color'].default_value = base_color
    principled.inputs['Roughness'].default_value = 0.4
    principled.inputs['Metallic'].default_value = 0.0

    links.new(principled.outputs['BSDF'], output.inputs['Surface'])

    if texture_path and os.path.exists(texture_path):
        tex = nodes.new(type="ShaderNodeTexImage")
        tex.image = bpy.data.images.load(texture_path)
        links.new(tex.outputs['Color'], principled.inputs['Base Color'])

    return mat

def add_mesh_object(mesh_obj, name, location=(0,0,0), material=None):
    """Helper to configure mesh object - objects from bpy.ops are already linked"""
    mesh_obj.name = name
    mesh_obj.location = location
    if material:
        if len(mesh_obj.data.materials) == 0:
            mesh_obj.data.materials.append(material)
        else:
            mesh_obj.data.materials[0] = material
    return mesh_obj

def primitive_humanoid(name_prefix="Avatar", scale=1.0, material=None):
    # Create simple primitives for torso, head, limbs and join them
    parts = []

    # Torso - cylinder
    bpy.ops.mesh.primitive_cylinder_add(vertices=32, radius=0.25*scale, depth=0.8*scale, location=(0,0,1.0*scale))
    torso = bpy.context.object
    torso.name = f"{name_prefix}_Torso"
    parts.append(torso)

    # Head - sphere
    bpy.ops.mesh.primitive_uv_sphere_add(radius=0.18*scale, location=(0,0,1.9*scale))
    head = bpy.context.object
    head.name = f"{name_prefix}_Head"
    parts.append(head)

    # Upper arms - cubes scaled
    bpy.ops.mesh.primitive_cube_add(size=0.2*scale, location=(0.45*scale,0,1.3*scale))
    r_arm = bpy.context.object
    r_arm.scale[0] = 2.0
    r_arm.name = f"{name_prefix}_R_UpperArm"
    parts.append(r_arm)

    bpy.ops.mesh.primitive_cube_add(size=0.2*scale, location=(-0.45*scale,0,1.3*scale))
    l_arm = bpy.context.object
    l_arm.scale[0] = 2.0
    l_arm.name = f"{name_prefix}_L_UpperArm"
    parts.append(l_arm)

    # Upper legs
    bpy.ops.mesh.primitive_cube_add(size=0.25*scale, location=(0.18*scale,0,0.5*scale))
    r_leg = bpy.context.object
    r_leg.scale[2] = 1.5
    r_leg.name = f"{name_prefix}_R_UpperLeg"
    parts.append(r_leg)

    bpy.ops.mesh.primitive_cube_add(size=0.25*scale, location=(-0.18*scale,0,0.5*scale))
    l_leg = bpy.context.object
    l_leg.scale[2] = 1.5
    l_leg.name = f"{name_prefix}_L_UpperLeg"
    parts.append(l_leg)

    # Join into single mesh (optional)
    bpy.ops.object.select_all(action='DESELECT')
    for p in parts:
        p.select_set(True)
        bpy.context.view_layer.objects.active = p
    bpy.ops.object.join()  # join into active
    combined = bpy.context.object
    combined.name = f"{name_prefix}_Mesh"

    # Apply material
    if material:
        if len(combined.data.materials) == 0:
            combined.data.materials.append(material)
        else:
            combined.data.materials[0] = material

    return combined

def create_armature_for_avatar(avatar_obj, name="AvatarArmature"):
    bpy.ops.object.select_all(action='DESELECT')
    # create armature
    bpy.ops.object.armature_add(enter_editmode=True, location=(0,0,0))
    arm = bpy.context.object
    arm.name = name
    arm_data = arm.data
    arm_data.name = name + "_Data"

    # edit mode bones: root, spine, head, left_arm, right_arm, left_leg, right_leg
    eb = arm.data.edit_bones
    eb.remove(eb[0])  # remove default bone to start clean

    def add_bone(bname, head, tail, parent=None):
        b = eb.new(bname)
        b.head = head
        b.tail = tail
        if parent:
            b.parent = eb[parent]
        return b

    add_bone("root", (0,0,0), (0,0,0.5))
    add_bone("spine", (0,0,0.5), (0,0,1.2), parent="root")
    add_bone("head", (0,0,1.2), (0,0,1.9), parent="spine")
    add_bone("left_arm", (-0.45,0,1.3), (-0.9,0,1.3), parent="spine")
    add_bone("right_arm", (0.45,0,1.3), (0.9,0,1.3), parent="spine")
    add_bone("left_leg", (-0.18,0,0.5), (-0.18,0,-0.5), parent="root")
    add_bone("right_leg", (0.18,0,0.5), (0.18,0,-0.5), parent="root")

    bpy.ops.object.mode_set(mode='OBJECT')

    # Parent mesh to armature with automatic weights
    bpy.context.view_layer.objects.active = avatar_obj
    avatar_obj.select_set(True)
    arm.select_set(True)
    bpy.ops.object.parent_set(type='ARMATURE_AUTO')

    return arm

def add_ground(scene_desc):
    """Add ground plane"""
    if scene_desc.get('ground'):
        ground_data = scene_desc.get('ground')
        size = ground_data.get('size', 20)
        color = ground_data.get('color', [0.4, 0.3, 0.2, 1])
    else:
        size = 20
        color = [0.4, 0.3, 0.2, 1]
    
    bpy.ops.mesh.primitive_plane_add(size=size, location=(0,0,0))
    ground = bpy.context.object
    ground.name = "Ground"
    mat_ground = make_material("GroundMat", base_color=tuple(color), texture_path=scene_desc.get('ground_texture'))
    # Object is already linked by primitive_plane_add, just set material
    if len(ground.data.materials) == 0:
        ground.data.materials.append(mat_ground)
    else:
        ground.data.materials[0] = mat_ground

def add_objects(scene_desc):
    """Add scene objects"""
    for i, obj in enumerate(scene_desc.get('objects', [])):
        typ = obj.get('type','cube')
        pos = obj.get('location', [i*1.5, 0.0, 0.5])
        obj_name = obj.get('name', f"EnvObj_{i}")
        
        if typ == 'cube':
            bpy.ops.mesh.primitive_cube_add(size=obj.get('size', 0.8), location=pos)
            o = bpy.context.object
        elif typ == 'sphere':
            bpy.ops.mesh.primitive_uv_sphere_add(radius=obj.get('radius',0.5), location=pos)
            o = bpy.context.object
        elif typ == 'cylinder':
            bpy.ops.mesh.primitive_cylinder_add(vertices=32, radius=obj.get('radius', 0.5), depth=obj.get('size', 1), location=pos)
            o = bpy.context.object
        elif typ == 'cone':
            bpy.ops.mesh.primitive_cone_add(vertices=32, radius1=obj.get('radius', 0.5), depth=obj.get('size', 1), location=pos)
            o = bpy.context.object
        elif typ == 'plane':
            bpy.ops.mesh.primitive_plane_add(size=obj.get('size', 10), location=pos)
            o = bpy.context.object
        else:
            bpy.ops.mesh.primitive_cube_add(size=0.5, location=pos)
            o = bpy.context.object

        # Object is already linked by bpy.ops, just set name, material, and transform
        o.name = obj_name
        o.location = tuple(pos)
        if obj.get('rotation'):
            o.rotation_euler = tuple(obj.get('rotation', [0, 0, 0]))
        if obj.get('scale'):
            o.scale = tuple(obj.get('scale', [1, 1, 1]))
        
        # Handle PBR material properties
        roughness = obj.get('roughness', 0.5)
        metallic = obj.get('metallic', 0.0)
        emission = obj.get('emission')
        emission_strength = obj.get('emissionStrength', 1.0)
        
        mat = make_material(f"objmat_{i}", base_color=tuple(obj.get('color', [1,1,1,1])), texture_path=obj.get('texture'))
        # Update material properties
        nodes = mat.node_tree.nodes
        principled = nodes.get('Principled BSDF')
        if principled:
            principled.inputs['Roughness'].default_value = roughness
            principled.inputs['Metallic'].default_value = metallic
            if emission:
                # Try different emission input names for different Blender versions
                try:
                    if 'Emission Color' in principled.inputs:
                        principled.inputs['Emission Color'].default_value = (*emission[:3], 1.0)
                    elif 'Emission' in principled.inputs:
                        principled.inputs['Emission'].default_value = (*emission[:3], 1.0)
                    if 'Emission Strength' in principled.inputs:
                        principled.inputs['Emission Strength'].default_value = emission_strength
                except KeyError:
                    # Skip emission if not supported
                    pass
        
        if len(o.data.materials) == 0:
            o.data.materials.append(mat)
        else:
            o.data.materials[0] = mat

def add_lighting(scene_desc):
    """Add default lighting"""
    bpy.ops.object.light_add(type='SUN', location=(10, -10, 10))
    sun = bpy.context.object
    sun.data.energy = scene_desc.get('sun_energy', 3.0)
    sun.data.angle = 0.5

    # Optional point lights from scene_desc
    for i, L in enumerate(scene_desc.get('lights', [])):
        typ = L.get('type','POINT')
        loc = L.get('location', [0,0,2])
        bpy.ops.object.light_add(type=typ, location=loc)
        light = bpy.context.object
        light.data.energy = L.get('energy', 50.0)

def apply_camera():
    # Camera
    bpy.ops.object.camera_add(location=(4.0, -4.0, 2.2))
    cam = bpy.context.object
    cam.rotation_euler = (1.05, 0.0, 0.78)
    bpy.context.scene.camera = cam

def main():
    # CLI args: script.py -- <input_json> <output_glb>
    argv = sys.argv
    if "--" in argv:
        idx = argv.index("--")
        args = argv[idx+1:]
    else:
        args = []

    if len(args) < 2:
        print("Usage: blender --background --python generate_scene.py -- <input.json> <output.glb>")
        return

    input_json = args[0]
    output_glb = args[1]

    with open(input_json, 'r') as f:
        scene_desc = json.load(f)

    clear_scene()

    # Ground
    add_ground(scene_desc)

    # Objects
    add_objects(scene_desc)

    # Avatar creation (if present)
    if scene_desc.get('avatar') and scene_desc.get('avatar', {}).get('present'):
        tex_file = scene_desc.get('avatar_texture')  # optional
        mat_avatar = make_material("AvatarMat", base_color=(1,0.8,0.6,1), texture_path=tex_file)
        avatar = primitive_humanoid(name_prefix="Hero", scale=1.0, material=mat_avatar)
        arm = create_armature_for_avatar(avatar, name="HeroArmature")
        # Set avatar position if specified
        avatar_pos = scene_desc.get('avatar', {}).get('position', [0, 0, 0])
        avatar.location = tuple(avatar_pos)
        avatar_scale = scene_desc.get('avatar', {}).get('scale', 1.0)
        avatar.scale = (avatar_scale, avatar_scale, avatar_scale)

    # Lighting
    if scene_desc.get('lighting'):
        lighting_data = scene_desc.get('lighting')
        # Add ambient light via world
        bpy.context.scene.world.use_nodes = True
        world_nodes = bpy.context.scene.world.node_tree.nodes
        if 'Background' in world_nodes:
            world_nodes['Background'].inputs['Strength'].default_value = lighting_data.get('ambient_strength', 0.3)
            light_color = lighting_data.get('color', [1, 1, 1])
            world_nodes['Background'].inputs['Color'].default_value = (*light_color[:3], 1.0)
        
        # Add lights from lighting.lights array
        if lighting_data.get('lights'):
            for i, L in enumerate(lighting_data.get('lights', [])):
                typ = L.get('type','POINT')
                loc = L.get('location', [0,0,2])
                if typ == 'SUN':
                    bpy.ops.object.light_add(type='SUN', location=loc)
                    sun = bpy.context.object
                    sun.data.energy = L.get('energy', 3.0)
                    sun.data.angle = L.get('angle', 0.5)
                    if L.get('color'):
                        sun.data.color = tuple(L.get('color', [1,1,1])[:3])
                elif typ == 'POINT':
                    bpy.ops.object.light_add(type='POINT', location=loc)
                    light = bpy.context.object
                    light.data.energy = L.get('energy', 50.0)
                    if L.get('color'):
                        light.data.color = tuple(L.get('color', [1,1,1])[:3])
                elif typ == 'SPOT':
                    bpy.ops.object.light_add(type='SPOT', location=loc)
                    light = bpy.context.object
                    light.data.energy = L.get('energy', 50.0)
                    if L.get('color'):
                        light.data.color = tuple(L.get('color', [1,1,1])[:3])
        else:
            # Default sun if no lights specified
            bpy.ops.object.light_add(type='SUN', location=(10, -10, 10))
            sun = bpy.context.object
            sun.data.energy = lighting_data.get('sun_energy', 3.0)
            sun.data.angle = 0.5
    else:
        # Default lighting
        add_lighting(scene_desc)

    # Camera
    if scene_desc.get('camera'):
        cam_data = scene_desc.get('camera')
        bpy.ops.object.camera_add(location=tuple(cam_data.get('position', [4.0, -4.0, 2.2])))
        cam = bpy.context.object
        cam.rotation_euler = tuple(cam_data.get('rotation', [1.05, 0.0, 0.78]))
        bpy.context.scene.camera = cam
    else:
        apply_camera()

    # Scene settings (cycles or eevee)
    bpy.context.scene.render.engine = 'CYCLES'
    bpy.context.scene.cycles.device = 'CPU'

    # Export glb
    export_kwargs = {
        "filepath": output_glb,
        "export_format": "GLB",
        "export_apply": True,
        "export_texcoords": True,
        "export_normals": True,
        "export_materials": "EXPORT",
        "export_colors": True,
        "export_extras": True,
    }
    bpy.ops.export_scene.gltf(**export_kwargs)
    print("Exported:", output_glb)

if __name__ == "__main__":
    main()