#!/usr/bin/env python3
"""Convert the extracted Ground Zero 1m DEM subset into Unreal heightmaps."""

from __future__ import annotations

import json
import struct
import time
from pathlib import Path

import numpy as np
import rasterio
from rasterio.enums import Resampling
from rasterio.warp import reproject

from prototype_ground_zero_terrain import TARGET_TILE_ID


PROJECT_ROOT = Path(__file__).resolve().parents[1]
EXTRACT_ROOT = PROJECT_ROOT / "Data" / "Terrain" / "Extracted" / TARGET_TILE_ID
UNREAL_ROOT = PROJECT_ROOT / "Data" / "Terrain" / "Unreal" / TARGET_TILE_ID
SOURCE_TIFF = EXTRACT_ROOT / f"{TARGET_TILE_ID}_1m_dem_subset.tif"
UNREAL_SIZE = 1009
TILE_SIZE_M = 1000.0
UNREAL_Z_SCALE_CM = 100.0


def encode_to_unreal_uint16(data: np.ma.MaskedArray) -> tuple[np.ndarray, float, float]:
    min_elevation = float(data.min())
    max_elevation = float(data.max())
    # Unreal Landscape encodes zero height around 32768. At Z scale 100, each
    # height unit is 100 / 128 cm and the full range is roughly +/- 256 m.
    height_values = 32768.0 + (data.filled(0.0) * 100.0 * 128.0 / UNREAL_Z_SCALE_CM)
    height_uint16 = np.rint(np.clip(height_values, 0.0, 65535.0)).astype("<u2")
    return height_uint16, min_elevation, max_elevation


def resample_to_unreal_size(source: rasterio.io.DatasetReader) -> np.ma.MaskedArray:
    source_data = source.read(1, masked=True)
    destination = np.empty((UNREAL_SIZE, UNREAL_SIZE), dtype=np.float32)

    dst_transform = rasterio.transform.from_bounds(
        source.bounds.left,
        source.bounds.bottom,
        source.bounds.right,
        source.bounds.top,
        UNREAL_SIZE,
        UNREAL_SIZE,
    )

    reproject(
        source=source_data.filled(source.nodata),
        destination=destination,
        src_transform=source.transform,
        src_crs=source.crs,
        src_nodata=source.nodata,
        dst_transform=dst_transform,
        dst_crs=source.crs,
        dst_nodata=source.nodata,
        resampling=Resampling.bilinear,
    )

    return np.ma.masked_equal(destination, source.nodata)


def write_r16(height_data: np.ndarray, output_path: Path) -> None:
    output_path.write_bytes(height_data.tobytes(order="C"))


def write_png16(height_data: np.ndarray, output_path: Path) -> bool:
    try:
        from PIL import Image
    except ModuleNotFoundError:
        return False

    image = Image.fromarray(height_data, mode="I;16")
    image.save(output_path)
    return True


def write_preview_pgm(height_data: np.ndarray, output_path: Path) -> None:
    preview = (height_data.astype(np.float32) / 65535.0 * 255.0).astype(np.uint8)
    header = f"P5\n{UNREAL_SIZE} {UNREAL_SIZE}\n255\n".encode("ascii")
    output_path.write_bytes(header + preview.tobytes(order="C"))


def write_import_metadata(
    output_path: Path,
    source: rasterio.io.DatasetReader,
    min_elevation: float,
    max_elevation: float,
    r16_path: Path,
    png_path: Path,
    png_written: bool,
    preview_path: Path,
) -> None:
    vertical_range = max_elevation - min_elevation
    metadata = {
        "tile_id": TARGET_TILE_ID,
        "generated_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "source_dem": str(SOURCE_TIFF.relative_to(PROJECT_ROOT)),
        "source_crs": str(source.crs),
        "source_bounds": list(source.bounds),
        "source_size_pixels": [source.width, source.height],
        "source_pixel_size_m": [abs(source.transform.a), abs(source.transform.e)],
        "heightmap": {
            "format": "r16_little_endian_unsigned",
            "width": UNREAL_SIZE,
            "height": UNREAL_SIZE,
            "r16_path": str(r16_path.relative_to(PROJECT_ROOT)),
            "png16_path": str(png_path.relative_to(PROJECT_ROOT)) if png_written else "",
            "preview_pgm_path": str(preview_path.relative_to(PROJECT_ROOT)),
            "min_elevation_m": min_elevation,
            "max_elevation_m": max_elevation,
            "vertical_range_m": vertical_range,
            "encoding": "unreal_landscape_midpoint_32768_sea_level",
        },
        "unreal_landscape_import": {
            "landscape_resolution": "1009 x 1009",
            "tile_world_size_m": TILE_SIZE_M,
            "x_scale_cm": 100.0 * TILE_SIZE_M / (UNREAL_SIZE - 1),
            "y_scale_cm": 100.0 * TILE_SIZE_M / (UNREAL_SIZE - 1),
            "z_scale_cm": UNREAL_Z_SCALE_CM,
            "z_offset_m": 0.0,
            "notes": [
                "Use the R16 file for import.",
                "Set X/Y scale to x_scale_cm/y_scale_cm so the 1009 samples span 1000 real meters.",
                "Set Z scale to z_scale_cm.",
                "Height values are encoded so Unreal landscape zero height corresponds to approximately sea level.",
                "1009 x 1009 is a valid Unreal Landscape import size close to the 1000m source tile."
            ],
        },
    }
    output_path.write_text(json.dumps(metadata, indent=2) + "\n")


def main() -> None:
    UNREAL_ROOT.mkdir(parents=True, exist_ok=True)
    r16_path = UNREAL_ROOT / f"{TARGET_TILE_ID}_unreal_1009.r16"
    png_path = UNREAL_ROOT / f"{TARGET_TILE_ID}_unreal_1009.png"
    preview_path = UNREAL_ROOT / f"{TARGET_TILE_ID}_unreal_1009_preview.pgm"
    metadata_path = UNREAL_ROOT / f"{TARGET_TILE_ID}_unreal_heightmap_metadata.json"

    with rasterio.open(SOURCE_TIFF) as source:
        resampled = resample_to_unreal_size(source)
        height_data, min_elevation, max_elevation = encode_to_unreal_uint16(resampled)
        write_r16(height_data, r16_path)
        png_written = write_png16(height_data, png_path)
        write_preview_pgm(height_data, preview_path)
        write_import_metadata(metadata_path, source, min_elevation, max_elevation, r16_path, png_path, png_written, preview_path)

    print(f"R16: {r16_path}")
    print(f"Preview: {preview_path}")
    if png_written:
        print(f"PNG16: {png_path}")
    else:
        print("PNG16 skipped: Pillow is not installed")
    print(f"Metadata: {metadata_path}")
    print(f"Elevation range: {min_elevation:.3f}m to {max_elevation:.3f}m")


if __name__ == "__main__":
    main()