#!/usr/bin/env python3
import glob
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)

gdal.UseExceptions()


def build_dgm_vrt_if_needed():
    """Build the DGM VRT automatically when missing."""
    if os.path.exists(VRT_PATH):
        return
    tif_paths = sorted(glob.glob(os.path.join(RAW_DIR, "*.tif")))
    if not tif_paths:
        raise SystemExit(f"No TIFFs found in {RAW_DIR}; cannot build {VRT_PATH}.")
    print(f"Building {VRT_PATH} from {len(tif_paths)} GeoTIFFs...")
    try:
        gdal.BuildVRT(VRT_PATH, tif_paths)
    except RuntimeError as exc:
        raise SystemExit(f"Could not build {VRT_PATH}: {exc}") from exc


def open_dataset(path, purpose):
    """Open a dataset and fail fast with context."""
    try:
        ds = gdal.Open(path)
    except RuntimeError as exc:
        raise SystemExit(f"{purpose}: {exc}") from exc
    if ds is None:
        raise SystemExit(f"{purpose}: GDAL returned None for {path}")
    return ds


def safe_remove(path):
    """Remove a file if present; return True when deleted."""
    try:
        os.remove(path)
        return True
    except FileNotFoundError:
        return False
    except OSError as exc:
        print(f"Warning: could not remove {path}: {exc}")
        return False


def cleanup_aux_files():
    """Clear GDAL sidecars and leftover temp files to keep the repo tidy."""
    patterns = [
        os.path.join("work", "*_tmp.tif"),
        os.path.join("work", "*_tmp.tif.aux.xml"),
        os.path.join("work", "*.aux.xml"),
        os.path.join(RAW_DIR, "*.aux.xml"),
    ]
    removed = 0
    for pattern in patterns:
        for path in glob.glob(pattern):
            if safe_remove(path):
                removed += 1
    print(f"Cleanup removed {removed} temporary files/sidecars.")


build_dgm_vrt_if_needed()
ds = open_dataset(VRT_PATH, f"Could not open {VRT_PATH} after attempting to build it.")

band = ds.GetRasterBand(1)

gmin, gmax = band.ComputeRasterMinMax(False)
print(f"GLOBAL_MIN={gmin}, GLOBAL_MAX={gmax}")

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")

    skipped = 0
    written = 0

    for tif in sorted(glob.glob(os.path.join(RAW_DIR, "*.tif"))):
        try:
            tds = open_dataset(tif, f"Skipping unreadable {tif}")
        except SystemExit as exc:
            print(exc)
            skipped += 1
            continue

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

        # Use the source tile footprint directly to avoid shifting during export.
        xmax = ulx + xres * tds.RasterXSize
        ymin = uly + yres * tds.RasterYSize
        xmin = ulx
        ymax = uly

        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,
            srcNodata=-9999,
            dstNodata=gmin,  # fill nodata with global min to avoid deep pits
        )
        try:
            gdal.Warp(tmp_path, ds, options=warp_opts)
        except RuntimeError as exc:
            print(f"Warp failed for {tile_id}: {exc}")
            skipped += 1
            continue

        # 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",
            creationOptions=["WORLDFILE=YES"],  # emit .wld so GIS tools place tiles correctly
        )
        try:
            gdal.Translate(out_path, tmp_path, options=trans_opts)
        except RuntimeError as exc:
            print(f"Translate failed for {tile_id}: {exc}")
            skipped += 1
            continue
        safe_remove(tmp_path)
        safe_remove(f"{tmp_path}.aux.xml")

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

print(f"Manifest: {manifest_path}")
print(f"Summary: wrote {written} tiles; skipped {skipped}.")
cleanup_aux_files()

if skipped:
    raise SystemExit(1)