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

Same pipeline as generate_procedural_terrain.py, but the height field comes
from an actual image (e.g. DemoAssets/poland_shaded.png) instead of procedural
noise. Reads heightmap_config.json (next to this script).

By default darker pixels = higher elevation and brighter pixels = lower, to
match shaded-relief style maps (set "invert": false in the config for the
opposite, standard-heightmap convention where white = high).

If the image has an alpha channel, pixels with alpha below "alpha_threshold"
(0-1, default 0.5) are cut out of the terrain mesh entirely (real holes, not
just low ground) and are excluded from foliage scatter/placement - useful for
a heightmap shaped like a country/island silhouette with a transparent
background.

Run import_environment_assets.py first so the mesh paths referenced in the
config actually exist. Set the HEIGHTMAP_CONFIG env var to point at a
different config file.
"""
import os, json, unreal
import terrain_common as tc
from png_reader import decode_png

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


class HeightmapField:
    def __init__(self, cfg):
        image_path = cfg["image_path"]
        print(f"Decoding heightmap image: {image_path}")
        width, height, channels, pixels = decode_png(image_path)
        print(f"  {width}x{height}, {channels} channel(s)")

        self.width = width
        self.height_px = height
        self.invert = cfg.get("invert", True)  # darker = higher, by default (shaded-relief convention)
        self.height_scale = cfg["height_scale"]
        self.size_x = cfg["size_x"]
        self.size_y = cfg["size_y"]
        self.alpha_threshold = cfg.get("alpha_threshold", 0.5)
        self.exclude_white = cfg.get("exclude_white", True)
        self.white_threshold = cfg.get("white_threshold", 0.98)
        self._lum = self._build_luminance(pixels, width, height, channels)
        self._alpha = self._build_alpha(pixels, width, height, channels)
        self._min_channel = self._build_min_channel(pixels, width, height, channels)

    @staticmethod
    def _build_luminance(pixels, width, height, channels):
        count = width * height
        lum = [0.0] * count
        if channels <= 2:
            for i in range(count):
                lum[i] = pixels[i * channels] / 255.0
        else:
            for i in range(count):
                o = i * channels
                lum[i] = (0.299 * pixels[o] + 0.587 * pixels[o + 1] + 0.114 * pixels[o + 2]) / 255.0
        return lum

    @staticmethod
    def _build_min_channel(pixels, width, height, channels):
        """Per-pixel min(R,G,B) (or the sole channel, if grayscale) in [0,1] -
        used to detect "fully white" fill pixels regardless of luminance
        formula quirks (e.g. pure yellow has high luminance but isn't white)."""
        count = width * height
        arr = [0.0] * count
        if channels <= 2:
            for i in range(count):
                arr[i] = pixels[i * channels] / 255.0
        else:
            for i in range(count):
                o = i * channels
                arr[i] = min(pixels[o], pixels[o + 1], pixels[o + 2]) / 255.0
        return arr

    @staticmethod
    def _build_alpha(pixels, width, height, channels):
        # Only grayscale+alpha (2) and RGBA (4) carry an alpha channel; anything
        # else is fully opaque, so there is nothing to mask out.
        if channels not in (2, 4):
            return None
        count = width * height
        alpha = [0.0] * count
        last = channels - 1
        for i in range(count):
            alpha[i] = pixels[i * channels + last] / 255.0
        return alpha

    def _world_to_image(self, x, y):
        # Map world (x, y) in [-size/2, size/2] to continuous image space, with
        # +Y world mapped toward row 0 (image top), matching a north-up map.
        u = (x + self.size_x / 2.0) / self.size_x
        v = (y + self.size_y / 2.0) / self.size_y
        u = min(max(u, 0.0), 1.0)
        v = min(max(v, 0.0), 1.0)
        return u * (self.width - 1), (1.0 - v) * (self.height_px - 1)

    def _sample_bilinear(self, channel_array, x, y):
        px, py = self._world_to_image(x, y)
        x0, y0 = int(px), int(py)
        x1, y1 = min(x0 + 1, self.width - 1), min(y0 + 1, self.height_px - 1)
        fx, fy = px - x0, py - y0

        def S(xx, yy):
            return channel_array[yy * self.width + xx]

        s00, s10, s01, s11 = S(x0, y0), S(x1, y0), S(x0, y1), S(x1, y1)
        s0 = s00 + (s10 - s00) * fx
        s1 = s01 + (s11 - s01) * fx
        return s0 + (s1 - s0) * fy

    def height(self, x, y):
        lum = self._sample_bilinear(self._lum, x, y)
        if self.invert:
            lum = 1.0 - lum
        return lum * self.height_scale

    def normal(self, x, y, step=200.0):
        return tc.finite_difference_normal(self.height, x, y, step)

    def is_masked(self, x, y):
        if self._alpha is not None and self._sample_bilinear(self._alpha, x, y) < self.alpha_threshold:
            return True
        if self.exclude_white and self._sample_bilinear(self._min_channel, x, y) >= self.white_threshold:
            return True
        return False


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

    terrain_cfg = config["terrain"]
    field = HeightmapField(terrain_cfg)
    base_folder = "HeightmapTerrain"

    print("Building terrain mesh from heightmap...")
    static_mesh, bounds = tc.build_terrain_mesh(terrain_cfg, field)

    loc = terrain_cfg["location"]
    origin = unreal.Vector(loc[0], loc[1], loc[2])
    tc.spawn_terrain_actor(static_mesh, origin, base_folder, "HeightmapTerrain")

    print("Scattering foliage layers...")
    for layer_cfg in config.get("scatter_layers", []):
        tc.scatter_layer(layer_cfg, field, bounds, origin, base_folder)

    explicit_entries = config.get("explicit_placements", [])
    if explicit_entries:
        print("Placing explicit manifest entries...")
        tc.place_explicit(explicit_entries, field, origin, base_folder)

    print("Done.")


main()