#!/usr/bin/env python3
import glob
import math
import os

from osgeo import gdal

RAW_DIR = "raw_dgm1"
VRT_PATH = "work/dgm.vrt"
OUT_DIR = "export_unity/height_png16"

TILE_SIZE_M = 1000  # real-world tile size in meters
OUT_RES = 1025  # Unity Terrain-friendly resolution (2^n + 1)
RESAMPLE = "bilinear"

os.makedirs("work", exist_ok=True)
os.makedirs(OUT_DIR, exist_ok=True)

# Open VRT
ds = gdal.Open(VRT_PATH)
if ds is None:
    raise SystemExit(f"Could not open {VRT_PATH}. Did you run gdalbuildvrt?")

band = ds.GetRasterBand(1)

# Strategy B: compute global min/max for CURRENT downloaded AOI
gmin, gmax = band.ComputeRasterMinMax(True)
print(f"GLOBAL_MIN={gmin}, GLOBAL_MAX={gmax}")

# Export manifest for Unity placement later
manifest_path = os.path.join("export_unity", "tile_index.csv")
with open(manifest_path, "w", encoding="utf-8") as f:
    f.write("tile_id,xmin,ymin,xmax,ymax,global_min,global_max,out_res\n")

    for tif in sorted(glob.glob(os.path.join(RAW_DIR, "*.tif"))):
        tds = gdal.Open(tif)
        if tds is None:
            print(f"Skipping unreadable: {tif}")
            continue

        gt = tds.GetGeoTransform()
        ulx, xres, _, uly, _, yres = gt  # yres typically negative in north-up rasters

        # Snap tile bounds to the 1000 m grid to avoid drift/seams.
        # (This also neutralizes the 1001×1001 overlap nuance.)
        xmin = math.floor(ulx / TILE_SIZE_M) * TILE_SIZE_M
        ymax = math.ceil(uly / TILE_SIZE_M) * TILE_SIZE_M
        xmax = xmin + TILE_SIZE_M
        ymin = ymax - TILE_SIZE_M

        base = os.path.splitext(os.path.basename(tif))[0]
        tile_id = base  # keep stable naming = easy re-export + reimport

        tmp_path = os.path.join("work", f"{tile_id}_tmp.tif")
        out_path = os.path.join(OUT_DIR, f"{tile_id}.png")

        warp_opts = gdal.WarpOptions(
            outputBounds=(xmin, ymin, xmax, ymax),
            width=OUT_RES,
            height=OUT_RES,
            resampleAlg=RESAMPLE,
            dstNodata=-9999,
        )
        gdal.Warp(tmp_path, ds, options=warp_opts)

        # Scale to UInt16 (0..65535) using Strategy-B global min/max
        trans_opts = gdal.TranslateOptions(
            outputType=gdal.GDT_UInt16,
            scaleParams=[(gmin, gmax, 0, 65535)],
            format="PNG",
        )
        gdal.Translate(out_path, tmp_path, options=trans_opts)

        f.write(f"{tile_id},{xmin},{ymin},{xmax},{ymax},{gmin},{gmax},{OUT_RES}\n")
        print(f"Wrote {out_path}")

print(f"Manifest: {manifest_path}")