#!/usr/bin/env python3
"""Export orthophoto tiles aligned to the terrain grid.

Inputs
- raw_dop/jp2/*.jp2 : orthophoto source tiles (DOP20 RGB)
- export_unity/tile_index.csv : manifest produced by export_heightmaps.py

Outputs
- work/dop.vrt : auto-built VRT mosaic of all JP2 tiles
- export_unity/ortho_jpg/<tile_id>.jpg : cropped JPEG tiles + .jgw worldfiles
"""
from __future__ import annotations

import csv
import glob
import os
from typing import Iterable

from osgeo import gdal

RAW_ORTHO_DIR = "raw_dop/jp2"
VRT_PATH = "work/dop.vrt"
TILE_INDEX = "export_unity/tile_index.csv"
OUT_DIR = "export_unity/ortho_jpg"

# 1000 m tiles at ~0.5 m/pixel give good visual quality in Unity while staying light.
OUT_RES = 2048
JPEG_QUALITY = 90

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


def build_vrt(jp2_paths: Iterable[str]) -> None:
    """Build the orthophoto VRT if missing."""
    print(f"Building VRT at {VRT_PATH} from {len(list(jp2_paths))} JP2 files...")
    gdal.BuildVRT(VRT_PATH, list(jp2_paths))


def open_dataset(path: str, purpose: str):
    """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 main() -> None:
    jp2_paths = sorted(glob.glob(os.path.join(RAW_ORTHO_DIR, "*.jp2")))
    if not jp2_paths:
        raise SystemExit(f"No JP2 files found in {RAW_ORTHO_DIR}. Run raw_dop/dlscript.sh first.")

    if not os.path.exists(VRT_PATH):
        build_vrt(jp2_paths)

    vrt_ds = open_dataset(VRT_PATH, f"Could not open VRT at {VRT_PATH}")

    if not os.path.exists(TILE_INDEX):
        raise SystemExit(f"Tile index missing: {TILE_INDEX}. Run export_heightmaps.py first.")

    with open(TILE_INDEX, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        written = 0
        skipped = 0

        for row in reader:
            try:
                tile_id = row["tile_id"]
                xmin = float(row["xmin"])
                ymin = float(row["ymin"])
                xmax = float(row["xmax"])
                ymax = float(row["ymax"])
            except (KeyError, ValueError) as exc:
                print(f"Skipping malformed row {row}: {exc}")
                skipped += 1
                continue

            out_path = os.path.join(OUT_DIR, f"{tile_id}.jpg")
            opts = gdal.TranslateOptions(
                format="JPEG",
                width=OUT_RES,
                height=OUT_RES,
                projWin=(xmin, ymax, xmax, ymin),  # xmin,xmax,ymax,ymin (upper-left origin)
                creationOptions=[f"QUALITY={JPEG_QUALITY}", "WORLDFILE=YES"],
            )
            try:
                gdal.Translate(out_path, vrt_ds, options=opts)
            except RuntimeError as exc:
                print(f"Translate failed for {tile_id}: {exc}")
                skipped += 1
                continue

            written += 1
            print(f"Wrote {out_path}")

    print(f"Summary: wrote {written} orthophoto tiles; skipped {skipped}.")
    if skipped:
        raise SystemExit(1)


if __name__ == "__main__":
    main()