"""Run inside Unreal (Python console / Execute Python Script).

Reads terrain_config.json (next to this script) and:
  1. builds a heightfield mesh with GeometryScript (fractal value noise),
     bakes it to a StaticMesh, and drops it into the level;
  2. scatters tree/rock StaticMeshActors across it per "scatter_layers",
     using the same noise field so foliage height always matches the
     ground beneath it (no raycasting needed).

Run import_environment_assets.py first so the mesh paths referenced in
terrain_config.json actually exist.

Set the TERRAIN_CONFIG env var to point at a different config file, e.g. to
run terrain_config_variants_demo.json without touching the default config.
"""
import os, json, math, random, unreal

CONFIG_PATH = os.environ.get(
    "TERRAIN_CONFIG",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "terrain_config.json"),
)

asset_tools = unreal.AssetToolsHelpers.get_asset_tools()
actor_subsystem = unreal.get_editor_subsystem(unreal.EditorActorSubsystem)


# ---------------------------------------------------------------------------
# Fractal value noise (pure Python, no engine dependency) -------------------
# ---------------------------------------------------------------------------
class ValueNoise2D:
    """Seeded, smooth, tileable-enough value noise for terrain-style height fields."""

    def __init__(self, seed):
        rng = random.Random(seed)
        table = list(range(256))
        rng.shuffle(table)
        self._perm = table + table  # avoid overflow when indexing i+1

    @staticmethod
    def _fade(t):
        return t * t * t * (t * (t * 6 - 15) + 10)

    @staticmethod
    def _lerp(a, b, t):
        return a + t * (b - a)

    def _gradient(self, ix, iy):
        h = self._perm[(self._perm[ix & 255] + iy) & 255]
        angle = (h / 255.0) * 2.0 * math.pi
        return math.cos(angle), math.sin(angle)

    def sample(self, x, y):
        x0, y0 = math.floor(x), math.floor(y)
        x1, y1 = x0 + 1, y0 + 1
        sx, sy = x - x0, y - y0

        def dot_grid(ix, iy, px, py):
            gx, gy = self._gradient(ix, iy)
            return gx * (px - ix) + gy * (py - iy)

        n00 = dot_grid(x0, y0, x, y)
        n10 = dot_grid(x1, y0, x, y)
        n01 = dot_grid(x0, y1, x, y)
        n11 = dot_grid(x1, y1, x, y)

        u, v = self._fade(sx), self._fade(sy)
        return self._lerp(self._lerp(n00, n10, u), self._lerp(n01, n11, u), v)


class FractalHeightField:
    def __init__(self, noise_cfg, height_scale):
        self.height_scale = height_scale
        self.frequency = noise_cfg["frequency"]
        self.octaves = noise_cfg["octaves"]
        self.lacunarity = noise_cfg["lacunarity"]
        self.gain = noise_cfg["gain"]
        self.noise = ValueNoise2D(noise_cfg["seed"])

    def height(self, x, y):
        amplitude, frequency, total, max_amp = 1.0, self.frequency, 0.0, 0.0
        for _ in range(self.octaves):
            total += amplitude * self.noise.sample(x * frequency, y * frequency)
            max_amp += amplitude
            amplitude *= self.gain
            frequency *= self.lacunarity
        return (total / max_amp) * self.height_scale

    def normal(self, x, y, step=25.0):
        h_l = self.height(x - step, y)
        h_r = self.height(x + step, y)
        h_d = self.height(x, y - step)
        h_u = self.height(x, y + step)
        # tangents along X and Y, cross product gives the surface normal.
        tangent_x = unreal.Vector(2 * step, 0.0, h_r - h_l)
        tangent_y = unreal.Vector(0.0, 2 * step, h_u - h_d)
        nx = tangent_x.y * tangent_y.z - tangent_x.z * tangent_y.y
        ny = tangent_x.z * tangent_y.x - tangent_x.x * tangent_y.z
        nz = tangent_x.x * tangent_y.y - tangent_x.y * tangent_y.x
        length = math.sqrt(nx * nx + ny * ny + nz * nz) or 1.0
        return unreal.Vector(nx / length, ny / length, nz / length)


# ---------------------------------------------------------------------------
# Terrain mesh ----------------------------------------------------------------
# ---------------------------------------------------------------------------
def build_terrain_mesh(cfg, field):
    steps = max(2, cfg["resolution"])
    mesh = unreal.DynamicMesh()
    primitive_options = unreal.GeometryScriptPrimitiveOptions()
    mesh = unreal.GeometryScript_Primitives.append_rectangle_xy(
        mesh, primitive_options, unreal.Transform(),
        dimension_x=cfg["size_x"], dimension_y=cfg["size_y"],
        steps_width=steps, steps_height=steps,
    )

    mesh, position_list, _has_gaps = unreal.GeometryScript_MeshQueries.get_all_vertex_positions(mesh, False)
    positions = unreal.GeometryScript_List.convert_vector_list_to_array(position_list)

    min_x = min_y = float("inf")
    max_x = max_y = float("-inf")
    displaced = []
    for p in positions:
        z = field.height(p.x, p.y)
        displaced.append(unreal.Vector(p.x, p.y, z))
        min_x, max_x = min(min_x, p.x), max(max_x, p.x)
        min_y, max_y = min(min_y, p.y), max(max_y, p.y)

    new_list = unreal.GeometryScript_List.convert_array_to_vector_list(displaced)
    mesh = unreal.GeometryScript_MeshEdits.set_all_mesh_vertex_positions(mesh, new_list)
    mesh = unreal.GeometryScript_Normals.set_per_vertex_normals(mesh)

    dest_path = cfg["dest_path"]
    asset_name = cfg["asset_name"]
    asset_path = f"{dest_path}/{asset_name}"
    if unreal.EditorAssetLibrary.does_asset_exist(asset_path):
        unreal.EditorAssetLibrary.delete_asset(asset_path)

    create_options = unreal.GeometryScriptCreateNewStaticMeshAssetOptions()
    static_mesh, _outcome = unreal.GeometryScript_NewAssetUtils.create_new_static_mesh_asset_from_mesh(
        mesh, asset_path, create_options,
    )
    unreal.EditorAssetLibrary.save_asset(static_mesh.get_path_name())

    bounds = (min_x, max_x, min_y, max_y)
    return static_mesh, bounds


_mesh_base_z_cache = {}


def mesh_base_z(mesh):
    """Local-space Z of the mesh's lowest point, so we can plant its base on the
    ground instead of its (often bounding-box-centered) pivot."""
    path = mesh.get_path_name()
    if path not in _mesh_base_z_cache:
        box = mesh.get_bounding_box()
        _mesh_base_z_cache[path] = box.min.z
    return _mesh_base_z_cache[path]


# ---------------------------------------------------------------------------
# Foliage scatter --------------------------------------------------------------
# ---------------------------------------------------------------------------
def scatter_layer(layer_cfg, field, bounds, origin, base_folder):
    min_x, max_x, min_y, max_y = bounds
    rng = random.Random(layer_cfg.get("seed", 0))
    meshes = [unreal.load_asset(path) for path in layer_cfg["meshes"]]
    meshes = [m for m in meshes if m is not None]
    if not meshes:
        print(f"  [{layer_cfg['name']}] no valid meshes found, skipping")
        return 0

    max_slope_cos = math.cos(math.radians(layer_cfg.get("max_slope_degrees", 90.0)))
    min_spacing = layer_cfg.get("min_spacing", 0.0)
    min_spacing_sqr = min_spacing * min_spacing

    placed_points = []
    spawned = 0
    max_attempts = layer_cfg["count"] * 25

    attempts = 0
    while spawned < layer_cfg["count"] and attempts < max_attempts:
        attempts += 1
        x = rng.uniform(min_x, max_x)
        y = rng.uniform(min_y, max_y)

        if min_spacing_sqr > 0.0 and any((x - px) ** 2 + (y - py) ** 2 < min_spacing_sqr for px, py in placed_points):
            continue

        z = field.height(x, y)
        if not (layer_cfg["min_height"] <= z <= layer_cfg["max_height"]):
            continue

        normal = field.normal(x, y)
        if normal.z < max_slope_cos:
            continue

        placed_points.append((x, y))

        mesh = rng.choice(meshes)
        scale = rng.uniform(layer_cfg["min_scale"], layer_cfg["max_scale"])
        # Shift up by -min_z*scale so the mesh's lowest vertex sits on the ground,
        # regardless of where the source asset's pivot was authored.
        ground_offset = -mesh_base_z(mesh) * scale
        location = unreal.Vector(origin.x + x, origin.y + y, origin.z + z + ground_offset)
        yaw = rng.uniform(0.0, 360.0)
        if layer_cfg.get("align_to_normal", False):
            # Approximate tilt from the surface normal; good enough for gentle slopes.
            pitch = -math.degrees(math.atan2(normal.x, normal.z))
            roll = math.degrees(math.atan2(normal.y, normal.z))
            rotation = unreal.Rotator(roll=roll, pitch=pitch, yaw=yaw)
        else:
            rotation = unreal.Rotator(roll=0.0, pitch=0.0, yaw=yaw)

        actor = actor_subsystem.spawn_actor_from_object(mesh, location, rotation)
        actor.set_actor_scale3d(unreal.Vector(scale, scale, scale))
        actor.set_folder_path(f"{base_folder}/{layer_cfg['name']}")
        actor.set_actor_label(f"{layer_cfg['name']}_{spawned:04d}")
        spawned += 1

    print(f"  [{layer_cfg['name']}] placed {spawned}/{layer_cfg['count']} ({attempts} attempts)")
    return spawned


def place_explicit(entries, field, origin, base_folder):
    """Places an exact, hand-authored (or externally-generated) list of instances.

    Unlike scatter_layer's density/slope rules, this reproduces a literal
    manifest byte-for-byte: the same JSON always yields the same named
    landmarks at the same coordinates - useful for hero props, quest markers,
    or a layout exported from an external design tool, none of which a
    rule-driven scatter (ours or a graph-based tool like PCG) is meant to pin down.
    """
    placed = 0
    for entry in entries:
        mesh = unreal.load_asset(entry["mesh"])
        if mesh is None:
            print(f"  [explicit] missing mesh asset: {entry['mesh']}")
            continue

        x, y, z_offset = entry.get("location", [0.0, 0.0, 0.0])
        scale = entry.get("scale", 1.0)
        if isinstance(scale, list):
            scale_vec = unreal.Vector(*scale)
        else:
            scale_vec = unreal.Vector(scale, scale, scale)

        z = z_offset
        if entry.get("snap_to_terrain", True):
            z += field.height(x, y) - mesh_base_z(mesh) * scale_vec.z

        roll, pitch, yaw = entry.get("rotation", [0.0, 0.0, 0.0])
        location = unreal.Vector(origin.x + x, origin.y + y, origin.z + z)
        rotation = unreal.Rotator(roll=roll, pitch=pitch, yaw=yaw)

        actor = actor_subsystem.spawn_actor_from_object(mesh, location, rotation)
        actor.set_actor_scale3d(scale_vec)
        actor.set_folder_path(f"{base_folder}/Explicit")
        actor.set_actor_label(entry.get("name", f"Explicit_{placed:04d}"))
        placed += 1

    print(f"  [explicit] placed {placed}/{len(entries)}")
    return placed


def generate_tile(config, terrain_cfg, tile_label):
    """Builds one terrain + its scatter layers + its explicit placements, all
    under a dedicated Outliner folder so multiple tiles can coexist in one level.
    """
    field = FractalHeightField(terrain_cfg["noise"], terrain_cfg["height_scale"])
    base_folder = f"ProceduralTerrain/{tile_label}"

    print(f"[{tile_label}] building terrain mesh (seed={terrain_cfg['noise']['seed']})...")
    static_mesh, bounds = build_terrain_mesh(terrain_cfg, field)

    loc = terrain_cfg["location"]
    origin = unreal.Vector(loc[0], loc[1], loc[2])
    terrain_actor = actor_subsystem.spawn_actor_from_object(static_mesh, origin, unreal.Rotator(0, 0, 0))
    terrain_actor.set_folder_path(base_folder)
    terrain_actor.set_actor_label(f"Terrain_{tile_label}")

    for layer_cfg in config.get("scatter_layers", []):
        scatter_layer(layer_cfg, field, bounds, origin, base_folder)

    explicit_entries = config.get("explicit_placements", [])
    if explicit_entries:
        place_explicit(explicit_entries, field, origin, base_folder)


def main():
    with open(CONFIG_PATH, "r") as f:
        config = json.load(f)

    base_terrain_cfg = config["terrain"]
    variants = config.get("variants")

    if not variants:
        # Single-tile mode: unchanged behavior for existing configs.
        generate_tile(config, base_terrain_cfg, "Main")
        print("Done.")
        return

    # Variant mode: one JSON list of seeds becomes N independent, deterministic
    # tiles laid out side by side in a single run - each a full reproducible
    # world driven by nothing but that seed. There's no "batch-produce N baked,
    # comparable layouts from one data file" workflow in the Foliage tool or a
    # PCG graph; you'd re-author or re-seed each one by hand, one at a time.
    gap = base_terrain_cfg["size_x"] * 1.2
    base_seed = base_terrain_cfg["noise"]["seed"]
    base_layer_seeds = {layer["name"]: layer.get("seed", 0) for layer in config.get("scatter_layers", [])}

    for i, seed_offset in enumerate(variants):
        tile_label = f"Variant_{seed_offset}"
        terrain_cfg = json.loads(json.dumps(base_terrain_cfg))  # cheap deep copy
        terrain_cfg["noise"]["seed"] = base_seed + seed_offset
        terrain_cfg["asset_name"] = f"{base_terrain_cfg['asset_name']}_{tile_label}"
        terrain_cfg["location"] = [
            base_terrain_cfg["location"][0] + i * gap,
            base_terrain_cfg["location"][1],
            base_terrain_cfg["location"][2],
        ]

        tile_config = json.loads(json.dumps(config))
        tile_config["terrain"] = terrain_cfg
        for layer_cfg in tile_config.get("scatter_layers", []):
            layer_cfg["seed"] = base_layer_seeds.get(layer_cfg["name"], 0) + seed_offset

        generate_tile(tile_config, terrain_cfg, tile_label)

    print(f"Done: {len(variants)} variant(s) generated side by side.")


main()