from __future__ import annotations

import glob
import os
from typing import Iterable

from osgeo import gdal

from .config import Config
from .gdal_utils import build_vrt, cleanup_aux_files, ensure_dir, ensure_parent, open_dataset, safe_remove

gdal.UseExceptions()


def _cleanup_patterns(raw_dir: str) -> Iterable[str]:
    return [
        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"),
    ]


def export_heightmaps(cfg: Config, *, force_vrt: bool = False) -> int:
    ensure_dir(cfg.work.work_dir)
    ensure_dir(cfg.export.heightmap_dir)
    ensure_parent(cfg.export.manifest_path)

    tif_paths = sorted(glob.glob(os.path.join(cfg.raw.dgm1_dir, "*.tif")))
    build_vrt(cfg.work.heightmap_vrt, tif_paths, force=force_vrt)
    ds = open_dataset(cfg.work.heightmap_vrt, f"Could not open {cfg.work.heightmap_vrt} after attempting to build it.")

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

    with open(cfg.export.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 tif_paths:
            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

            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

            tmp_path = os.path.join(cfg.work.work_dir, f"{tile_id}_tmp.tif")
            out_path = os.path.join(cfg.export.heightmap_dir, f"{tile_id}.png")

            warp_opts = gdal.WarpOptions(
                outputBounds=(xmin, ymin, xmax, ymax),
                width=cfg.heightmap.out_res,
                height=cfg.heightmap.out_res,
                resampleAlg=cfg.heightmap.resample,
                srcNodata=-9999,
                dstNodata=0,  # Use 0 for nodata (safe since valid elevations scale to 1-65535)
            )
            try:
                gdal.Warp(tmp_path, ds, options=warp_opts)
            except RuntimeError as exc:
                print(f"Warp failed for {tile_id}: {exc}")
                skipped += 1
                continue

            trans_opts = gdal.TranslateOptions(
                outputType=gdal.GDT_UInt16,
                scaleParams=[(gmin, gmax, 0, 65535)],
                format="PNG",
                creationOptions=["WORLDFILE=YES"],
            )
            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},{cfg.heightmap.out_res}\n")
            print(f"Wrote {out_path}")
            written += 1

    print(f"Manifest: {cfg.export.manifest_path}")
    print(f"Summary: wrote {written} tiles; skipped {skipped}.")
    removed = cleanup_aux_files(_cleanup_patterns(cfg.raw.dgm1_dir))
    print(f"Cleanup removed {removed} temporary files/sidecars.")
    return 1 if skipped else 0
