import bpy
from pathlib import Path


def find_snow_modifier(terrain_obj):
    """
    Find the Geometry Nodes modifier that contains snow_depth attribute.
    Returns the modifier or None if not found.
    """
    for mod in terrain_obj.modifiers:
        if mod.type == 'NODES' and mod.node_group:
            # Check if this modifier's node tree has Store Named Attribute with "snow_depth"
            for node in mod.node_group.nodes:
                if node.type == 'STORE_NAMED_ATTRIBUTE':
                    if hasattr(node, 'data_type') and node.name and 'snow' in node.name.lower():
                        return mod
                    # Check inputs for the name
                    for input in node.inputs:
                        if input.name == 'Name' and hasattr(input, 'default_value'):
                            if input.default_value == 'snow_depth':
                                return mod
            # Fallback: check modifier name
            if 'snow' in mod.name.lower():
                return mod
    return None


def bake_snow_depth(terrain_obj, resolution=512, output_path=None, modifier_name=None):
    """
    Bake snow depth attribute to texture using shader-based Cycles baking.
    Uses the same approach as generate_heightmap.py.

    Requires:
    - Terrain object with Geometry Nodes modifier that stores 'snow_depth' attribute
    - UV map on terrain mesh

    Args:
        terrain_obj: Terrain mesh object with snow_depth attribute
        resolution: Texture resolution (square)
        output_path: Path to save EXR file
        modifier_name: Optional specific modifier name to use (e.g., "Snow")
    """

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

    # Find the snow geometry nodes modifier
    if modifier_name:
        geo_nodes_modifier = terrain_obj.modifiers.get(modifier_name)
        if not geo_nodes_modifier:
            raise ValueError(f"Modifier '{modifier_name}' not found on {terrain_obj.name}")
        print(f"Using specified modifier: {modifier_name}")
    else:
        geo_nodes_modifier = find_snow_modifier(terrain_obj)
        if not geo_nodes_modifier:
            print("\nAvailable Geometry Nodes modifiers:")
            for mod in terrain_obj.modifiers:
                if mod.type == 'NODES':
                    print(f"  - {mod.name}")
            raise ValueError(
                f"No Geometry Nodes modifier with 'snow_depth' attribute found on {terrain_obj.name}!\n"
                f"Either add snow accumulation modifier, or specify modifier_name parameter."
            )
        print(f"Found snow modifier: {geo_nodes_modifier.name}")

    modifier_states = {}
    print(f"\nDisabling modifiers after '{geo_nodes_modifier.name}' for baking...")
    target_mod_index = list(terrain_obj.modifiers).index(geo_nodes_modifier)

    for i, mod in enumerate(terrain_obj.modifiers):
        modifier_states[mod.name] = {
            'show_viewport': mod.show_viewport,
            'show_render': mod.show_render
        }

        if i > target_mod_index:
            print(f"  Temporarily disabling: {mod.name}")
            mod.show_viewport = False
            mod.show_render = False

    bpy.context.view_layer.update()
    depsgraph = bpy.context.evaluated_depsgraph_get()
    evaluated_obj = terrain_obj.evaluated_get(depsgraph)
    eval_mesh = evaluated_obj.to_mesh()

    if 'snow_depth' not in eval_mesh.attributes:
        evaluated_obj.to_mesh_clear()
        for mod_name, state in modifier_states.items():
            mod = terrain_obj.modifiers.get(mod_name)
            if mod:
                mod.show_viewport = state['show_viewport']
                mod.show_render = state['show_render']
        raise ValueError("snow_depth attribute missing from evaluated geometry")

    print(f"✓ Verified 'snow_depth' attribute exists")
    evaluated_obj.to_mesh_clear()

    # 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="SnowDepth_Bake",
        width=resolution,
        height=resolution,
        alpha=False,
        float_buffer=True,
        is_data=True
    )

    print(f"Created bake image: {bake_image.name}")

    original_materials = list(terrain_obj.data.materials)
    print(f"Object has {len(original_materials)} material slot(s): {[mat.name if mat else 'None' for mat in original_materials]}")

    mat = bpy.data.materials.new(name="SnowDepth_BakeMaterial")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    nodes.clear()

    attr_node = nodes.new(type='ShaderNodeAttribute')
    attr_node.attribute_name = 'snow_depth'

    emission_node = nodes.new(type='ShaderNodeEmission')
    mat.node_tree.links.new(attr_node.outputs['Fac'], emission_node.inputs['Color'])

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

    image_node = nodes.new(type='ShaderNodeTexImage')
    image_node.image = bake_image
    image_node.select = True
    nodes.active = image_node

    terrain_obj.data.materials.clear()
    terrain_obj.data.materials.append(mat)
    print(f"Temporarily replaced all materials with bake material")

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

    # Ensure we're in object mode
    if bpy.context.object and bpy.context.object.mode != 'OBJECT':
        bpy.ops.object.mode_set(mode='OBJECT')

    # 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 with Cycles (EMIT)...")
    bpy.ops.object.bake(type='EMIT', use_clear=True)

    print("Bake complete!")

    # Verify bake has data (not all black/zero)
    pixels = list(bake_image.pixels)
    max_value = max(pixels) if pixels else 0.0
    avg_value = sum(pixels) / len(pixels) if pixels else 0.0
    non_zero_count = sum(1 for p in pixels if p > 0.0001)
    print(f"Baked image stats: max={max_value:.4f}, avg={avg_value:.4f}")
    print(f"Non-zero pixels: {non_zero_count} ({non_zero_count / len(pixels) * 100:.1f}%)")

    if max_value < 0.0001:
        print("\n⚠️  WARNING: Baked image appears to be all black!")
        print("   Possible causes:")
        print("   - 'snow_depth' attribute doesn't exist in the geometry")
        print("   - Geometry Nodes modifier is disabled")
        print("   - Store Named Attribute node is not connected")
        print("   - Wrong modifier selected (try specifying modifier_name)")
        print("\n   Continuing anyway, but check your setup...")
    else:
        print(f"✓ Bake contains data (values up to {max_value:.4f}m)")

    # 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
        original_exr_codec = scene.render.image_settings.exr_codec

        # Use BW mode for single channel (same as heightmap)
        scene.render.image_settings.color_mode = 'BW'
        scene.render.image_settings.color_depth = '32'
        scene.render.image_settings.exr_codec = 'ZIP'

        print(f"Saving EXR with settings: color_mode=BW, depth=32, codec=ZIP")
        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
        scene.render.image_settings.exr_codec = original_exr_codec

        print(f"Saved to: {output_path}")
        print(f"Format: OpenEXR, 32-bit float, ZIP compression")
        print(f"File size: {output_path.stat().st_size / 1024:.1f} KB")

    bpy.data.images.remove(bake_image)
    bpy.data.materials.remove(mat)

    terrain_obj.data.materials.clear()
    for original_mat in original_materials:
        terrain_obj.data.materials.append(original_mat)
    print(f"Restored {len(original_materials)} original material(s)")

    print("\nRestoring modifier states...")
    for mod_name, state in modifier_states.items():
        mod = terrain_obj.modifiers.get(mod_name)
        if mod:
            mod.show_viewport = state['show_viewport']
            mod.show_render = state['show_render']
    print("✓ Modifiers restored")

    return True


if __name__ == "__main__":
    project_root = Path(bpy.data.filepath).parent.parent
    output_path = project_root / "textures" / "snow_depth.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!")

    # CONFIGURATION: Specify modifier name if you have multiple Geometry Nodes modifiers
    # Leave as None to auto-detect the snow modifier
    # Example: modifier_name = "Snow" or "Snow Accumulation"
    modifier_name = "Snow Accumulation"  # Auto-detect by looking for 'snow_depth' attribute

    bake_snow_depth(
        terrain_obj=terrain_obj,
        resolution=1000,
        output_path=output_path,
        modifier_name=modifier_name  # Specify "Snow" if auto-detect fails
    )

    print("\n" + "="*60)
    print("Snow depth baking complete!")
    print(f"Output: {output_path}")
    print("="*60)
    print("\nNext steps:")
    print("1. Verify snow_depth.exr in textures/ directory")
    print("2. Open in image viewer to check it's not black")
    print("3. Load in game with SnowLayer::load()")
    print("4. Test deformation with player movement")
    print("\nIf bake is black:")
    print("- Check that 'snow_depth' attribute exists in Spreadsheet Editor")
    print("- Verify Geometry Nodes modifier has Store Named Attribute node")
    print("- Try specifying modifier_name='Snow' explicitly in script")
