import bpy
import sys
import os

def create_tex_node(nodes, img_path, label, color_space, location):
    img = bpy.data.images.load(img_path)
    tex = nodes.new(type='ShaderNodeTexImage')
    tex.image = img
    tex.label = label
    tex.location = location
    tex.image.colorspace_settings.name = color_space
    return tex

def setup_environment_lighting(hdri_path):
    if not bpy.data.worlds:
        bpy.data.worlds.new(name="World")
    if bpy.context.scene.world is None:
        bpy.context.scene.world = bpy.data.worlds[0]
    world = bpy.context.scene.world

    world.use_nodes = True
    nodes = world.node_tree.nodes
    links = world.node_tree.links
    nodes.clear()

    env_tex = nodes.new(type="ShaderNodeTexEnvironment")
    env_tex.image = bpy.data.images.load(hdri_path)
    env_tex.location = (-300, 0)

    bg = nodes.new(type="ShaderNodeBackground")
    bg.location = (0, 0)

    output = nodes.new(type="ShaderNodeOutputWorld")
    output.location = (300, 0)

    links.new(env_tex.outputs["Color"], bg.inputs["Color"])
    links.new(bg.outputs["Background"], output.inputs["Surface"])

def setup_gpu_rendering():
    bpy.context.scene.render.engine = 'CYCLES'
    prefs = bpy.context.preferences
    cprefs = prefs.addons['cycles'].preferences

    # Choose backend depending on GPU type: 'CUDA', 'OPTIX', 'HIP', 'METAL'
    cprefs.compute_device_type = 'CUDA'
    bpy.context.scene.cycles.device = 'GPU'

def generate_blend(obj_path, base_color_path, normal_map_path, roughness_path, metallic_path, output_blend):
    # Reset scene
    bpy.ops.wm.read_factory_settings(use_empty=True)

    # Import OBJ
    bpy.ops.import_scene.obj(filepath=obj_path)
    obj = bpy.context.selected_objects[0]

    # Create material
    mat = bpy.data.materials.new(name="BRDF_Material")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    nodes.clear()

    output = nodes.new(type='ShaderNodeOutputMaterial')
    output.location = (400, 0)

    principled = nodes.new(type='ShaderNodeBsdfPrincipled')
    principled.location = (100, 0)
    links.new(principled.outputs['BSDF'], output.inputs['Surface'])

    # Base Color
    base_color = create_tex_node(nodes, base_color_path, "Base Color", 'sRGB', (-600, 200))
    links.new(base_color.outputs['Color'], principled.inputs['Base Color'])

    # Roughness
    rough = create_tex_node(nodes, roughness_path, "Roughness", 'Non-Color', (-600, 0))
    links.new(rough.outputs['Color'], principled.inputs['Roughness'])

    # Metallic
    metal = create_tex_node(nodes, metallic_path, "Metallic", 'Non-Color', (-600, -200))
    links.new(metal.outputs['Color'], principled.inputs['Metallic'])

    # Normal Map
    normal_tex = create_tex_node(nodes, normal_map_path, "Normal Map", 'Non-Color', (-800, -400))
    normal_map = nodes.new(type='ShaderNodeNormalMap')
    normal_map.location = (-400, -400)
    links.new(normal_tex.outputs['Color'], normal_map.inputs['Color'])
    links.new(normal_map.outputs['Normal'], principled.inputs['Normal'])

    # Assign material
    if obj.data.materials:
        obj.data.materials[0] = mat
    else:
        obj.data.materials.append(mat)

    # Global Illumination using Blender's default forest HDRI
    blender_data_path = bpy.utils.resource_path('LOCAL')
    forest_hdri_path = os.path.join(blender_data_path, "datafiles", "studiolights", "world", "forest.exr")
    print(f"Using HDRI: {forest_hdri_path}")
    setup_environment_lighting(forest_hdri_path)

    # GPU rendering setup
    setup_gpu_rendering()

    # Pack textures into .blend
    bpy.ops.file.pack_all()
    
    # Set the 3D View to Rendered mode and focus on object
    for area in bpy.context.screen.areas:
        if area.type == 'VIEW_3D':
            for space in area.spaces:
                if space.type == 'VIEW_3D':
                    space.shading.type = 'RENDERED'  # Set viewport shading to Rendered
            for region in area.regions:
                if region.type == 'WINDOW':
                    override = {'area': area, 'region': region, 'scene': bpy.context.scene}
                    bpy.ops.view3d.view_all(override, center=True)

        elif area.type == 'NODE_EDITOR':
            for space in area.spaces:
                if space.type == 'NODE_EDITOR':
                    space.tree_type = 'ShaderNodeTree'  # Switch to Shader Editor
                    space.shader_type = 'OBJECT'

    # Optional: Switch active workspace to Shading (if it exists)
    for workspace in bpy.data.workspaces:
        if workspace.name == 'Shading':
            bpy.context.window.workspace = workspace
            break

    # Save the .blend file
    bpy.ops.wm.save_as_mainfile(filepath=output_blend)
    print(f"✅ Saved .blend file with BRDF, HDRI, GPU: {output_blend}")

if __name__ == "__main__":
    argv = sys.argv
    argv = argv[argv.index("--") + 1:]  # Only use args after "--"

    if len(argv) != 6:
        print("Usage:\n  blender --background --python generate_blend.py -- obj base_color normal roughness metallic output.blend")
        sys.exit(1)

    generate_blend(*argv)