import bpy
from pathlib import Path


def bake_heightmap(terrain_obj, resolution=1024, output_path=None):
    """
    Bake terrain heightmap using Blender's render/bake system.

    Args:
        terrain_obj: Terrain mesh object
        resolution: Texture resolution (square)
        output_path: Path to save EXR file
    """

    print(f"Baking heightmap for: {terrain_obj.name}")
    print(f"Resolution: {resolution}×{resolution}")

    # Ensure object has UV map
    if not terrain_obj.data.uv_layers:
        print("Adding UV map...")
        terrain_obj.data.uv_layers.new(name="UVMap")

    # Create new image for baking
    bake_image = bpy.data.images.new(
        name="Heightmap_Bake",
        width=resolution,
        height=resolution,
        alpha=False,
        float_buffer=True,
        is_data=True
    )

    # Setup material for baking
    if not terrain_obj.data.materials:
        mat = bpy.data.materials.new(name="Heightmap_Material")
        terrain_obj.data.materials.append(mat)
    else:
        mat = terrain_obj.data.materials[0]

    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    nodes.clear()

    # Create nodes for height baking
    # Geometry node to get position
    geo_node = nodes.new(type='ShaderNodeNewGeometry')

    # Separate XYZ to get Z (height)
    separate_node = nodes.new(type='ShaderNodeSeparateXYZ')
    mat.node_tree.links.new(geo_node.outputs['Position'], separate_node.inputs['Vector'])

    # Emission shader to output height value
    emission_node = nodes.new(type='ShaderNodeEmission')
    mat.node_tree.links.new(separate_node.outputs['Z'], emission_node.inputs['Color'])

    # Material output
    output_node = nodes.new(type='ShaderNodeOutputMaterial')
    mat.node_tree.links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])

    # Add image texture node (required for baking target)
    image_node = nodes.new(type='ShaderNodeTexImage')
    image_node.image = bake_image
    image_node.select = True
    nodes.active = image_node

    # Select object and set mode
    bpy.context.view_layer.objects.active = terrain_obj
    terrain_obj.select_set(True)

    # Setup render settings for baking
    bpy.context.scene.render.engine = 'CYCLES'
    bpy.context.scene.cycles.samples = 1
    bpy.context.scene.cycles.bake_type = 'EMIT'

    print("Baking...")
    bpy.ops.object.bake(type='EMIT', use_clear=True)

    print("Bake complete!")

    # Save as EXR
    if output_path:
        bake_image.filepath_raw = str(output_path)
        bake_image.file_format = 'OPEN_EXR'
        bake_image.use_half_precision = False

        scene = bpy.context.scene
        original_color_mode = scene.render.image_settings.color_mode
        original_color_depth = scene.render.image_settings.color_depth

        scene.render.image_settings.color_mode = 'BW'
        scene.render.image_settings.color_depth = '32'

        bake_image.save_render(str(output_path), scene=scene)

        scene.render.image_settings.color_mode = original_color_mode
        scene.render.image_settings.color_depth = original_color_depth

        print(f"Saved to: {output_path}")

    # Cleanup
    bpy.data.images.remove(bake_image)

    return True


if __name__ == "__main__":
    project_root = Path(bpy.data.filepath).parent.parent
    output_path = project_root / "textures" / "terrain_heightmap.exr"

    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Find terrain object
    terrain_obj = bpy.data.objects.get("TerrainPlane")

    if not terrain_obj:
        print("'TerrainPlane' not found. Searching for terrain mesh...")
        for obj in bpy.data.objects:
            if obj.type == 'MESH' and ('terrain' in obj.name.lower() or 'plane' in obj.name.lower()):
                terrain_obj = obj
                print(f"Using: {obj.name}")
                break

    if not terrain_obj:
        raise ValueError("No terrain object found!")

    bake_heightmap(
        terrain_obj=terrain_obj,
        resolution=1000,
        output_path=output_path
    )

    print("\n" + "="*60)
    print("Heightmap baking complete!")
    print(f"Output: {output_path}")
    print("="*60)