import bpy
import bmesh
import numpy as np
from pathlib import Path
from mathutils import Vector

def simple_blur(data, iterations=1):
    """
    Simple 3x3 box blur for vector fields.
    """
    result = data.copy()
    h, w, c = data.shape

    for _ in range(iterations):
        blurred = np.zeros_like(result)
        for y in range(h):
            for x in range(w):
                count = 0
                total = np.zeros(c)

                for dy in [-1, 0, 1]:
                    for dx in [-1, 0, 1]:
                        ny, nx = y + dy, x + dx
                        if 0 <= ny < h and 0 <= nx < w:
                            total += result[ny, nx]
                            count += 1

                blurred[y, x] = total / count if count > 0 else result[y, x]

        result = blurred

    return result


def generate_normal_map_from_terrain(terrain_obj, resolution=1024, blur_iterations=2):
    """
    Bake terrain surface normals to a texture for neighbor sampling.

    Args:
        terrain_obj: Blender mesh object (the terrain plane)
        resolution: Output texture resolution (square)
        blur_iterations: Number of smoothing passes
    """

    print(f"Generating normal map from terrain mesh: {terrain_obj.name}")
    print(f"Output resolution: {resolution}×{resolution}")

    mesh = terrain_obj.data
    bm = bmesh.new()
    bm.from_mesh(mesh)
    bm.verts.ensure_lookup_table()
    bm.faces.ensure_lookup_table()

    world_matrix = terrain_obj.matrix_world
    normal_matrix = world_matrix.to_3x3().inverted().transposed()

    verts_world = [world_matrix @ v.co for v in bm.verts]
    normals_world = [normal_matrix @ v.normal for v in bm.verts]

    if len(verts_world) == 0:
        raise ValueError("Terrain mesh has no vertices")

    xs = [v.x for v in verts_world]
    ys = [v.y for v in verts_world]

    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)

    print(f"Terrain bounds: X=[{min_x:.2f}, {max_x:.2f}], Y=[{min_y:.2f}, {max_y:.2f}]")

    normal_map = np.zeros((resolution, resolution, 3), dtype=np.float32)
    sample_counts = np.zeros((resolution, resolution), dtype=np.int32)

    print("Sampling vertex normals to grid...")
    for v_pos, v_normal in zip(verts_world, normals_world):
        u = (v_pos.x - min_x) / (max_x - min_x) if max_x > min_x else 0.5
        v_coord = (v_pos.y - min_y) / (max_y - min_y) if max_y > min_y else 0.5

        u = np.clip(u, 0.0, 0.9999)
        v_coord = np.clip(v_coord, 0.0, 0.9999)

        grid_x = int(u * resolution)
        grid_y = int(v_coord * resolution)

        normal_map[grid_y, grid_x, 0] += v_normal.x
        normal_map[grid_y, grid_x, 1] += v_normal.y
        normal_map[grid_y, grid_x, 2] += v_normal.z
        sample_counts[grid_y, grid_x] += 1

    mask = sample_counts > 0
    for c in range(3):
        normal_map[:, :, c][mask] /= sample_counts[mask]

    print("Interpolating missing values...")
    if not mask.all():
        filled = normal_map.copy()
        for _ in range(10):
            smoothed = simple_blur(filled, iterations=2)
            for c in range(3):
                filled[:, :, c] = np.where(mask, filled[:, :, c], smoothed[:, :, c])
        normal_map = filled

    print("Normalizing normal vectors...")
    magnitudes = np.sqrt(np.sum(normal_map**2, axis=2))
    magnitudes = np.maximum(magnitudes, 1e-8)
    for c in range(3):
        normal_map[:, :, c] /= magnitudes

    if blur_iterations > 0:
        print(f"Smoothing normal field ({blur_iterations} iterations)...")
        normal_map = simple_blur(normal_map, iterations=blur_iterations)
        magnitudes = np.sqrt(np.sum(normal_map**2, axis=2))
        magnitudes = np.maximum(magnitudes, 1e-8)
        for c in range(3):
            normal_map[:, :, c] /= magnitudes

    print("Encoding to RGB texture (world space normals)...")
    normal_encoded = normal_map * 0.5 + 0.5

    normal_rgba = np.zeros((resolution, resolution, 4), dtype=np.float32)
    normal_rgba[:, :, 0] = normal_encoded[:, :, 0]
    normal_rgba[:, :, 1] = normal_encoded[:, :, 1]
    normal_rgba[:, :, 2] = normal_encoded[:, :, 2]
    normal_rgba[:, :, 3] = 1.0

    normal_flat = normal_rgba.flatten()

    print(f"Creating Blender image...")
    output_img = bpy.data.images.new(
        name="Terrain_Normal_Map",
        width=resolution,
        height=resolution,
        alpha=True,
        float_buffer=True
    )

    output_img.pixels[:] = normal_flat

    bm.free()

    return output_img


def save_normal_map(output_path, resolution=1024, blur_iterations=2, terrain_name="TerrainPlane"):
    """
    Main function to generate and save normal map from terrain in current Blender file.

    Args:
        output_path: Path to save PNG normal map
        resolution: Output texture resolution
        blur_iterations: Smoothing iterations
        terrain_name: Name of terrain object
    """

    terrain_obj = bpy.data.objects.get(terrain_name)

    if not terrain_obj:
        print(f"Object '{terrain_name}' not found. Searching for terrain-like objects...")
        for obj in bpy.data.objects:
            if obj.type == 'MESH':
                print(f"  Found mesh: {obj.name}")
                if '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(f"Object '{terrain_name}' not found and no terrain mesh detected. Available objects: {[obj.name for obj in bpy.data.objects if obj.type == 'MESH']}")

    print(f"Using terrain object: {terrain_obj.name}")

    if terrain_obj.type != 'MESH':
        raise ValueError(f"Object '{terrain_obj.name}' is not a mesh")

    normal_img = generate_normal_map_from_terrain(
        terrain_obj,
        resolution=resolution,
        blur_iterations=blur_iterations
    )

    normal_img.filepath_raw = str(output_path)
    normal_img.file_format = 'PNG'
    normal_img.save()

    bpy.data.images.remove(normal_img)

    print(f"\n{'='*60}")
    print(f"Normal map generation complete!")
    print(f"Output: {output_path}")
    print(f"Resolution: {resolution}×{resolution}")
    print(f"{'='*60}")
    print("\nNormal encoding:")
    print("  RGB channels: World-space normal (0.5, 0.5, 0.5) = (0, 0, 0)")
    print("  Decode in shader: normal = texture.rgb * 2.0 - 1.0;")


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

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

    save_normal_map(
        output_path=output_path,
        resolution=1024,
        blur_iterations=2,
        terrain_name="TerrainPlane"
    )
