#!/usr/bin/env python3

# Create an ImageMosaic index for colour-corrected ice shelf imagery.
#
# Processing:
#   - Scan cwd/imagery/<ice_shelf>/colcorrect/*.tif
#   - Create polygon footprints from each GeoTIFF extent
#   - Set featname from the ice shelf folder name
#   - Extract year from the YYYYMMDD component of the filename
#   - Store image paths relative to cwd/imagery
#   - Write imagery.shp to cwd/imagery using EPSG:3031
#
# Install:
#   pip install pyshp GDAL

from __future__ import annotations

import re
from pathlib import Path

import shapefile
from osgeo import gdal


# =============================================================================
# CONFIGURATION
# =============================================================================

BASE_DIR = Path.cwd()
IMAGERY_DIR = BASE_DIR / "imagery"

INPUT_FOLDER = "colcorrect"
OUTPUT_NAME = "imagery"

YEAR_RE = re.compile(r"_(\d{4})\d{4}_")

# Exact single-line EPSG:3031 WKT required for the ImageMosaic index.
EPSG3031_WKT = 'PROJCS["WGS 84 / Antarctic Polar Stereographic", GEOGCS["WGS 84", DATUM["World Geodetic System 1984", SPHEROID["WGS 84", 6378137.0, 298.257223563, AUTHORITY["EPSG","7030"]], AUTHORITY["EPSG","6326"]], PRIMEM["Greenwich", 0.0, AUTHORITY["EPSG","8901"]], UNIT["degree", 0.017453292519943295], AXIS["Geodetic longitude", EAST], AXIS["Geodetic latitude", NORTH], AUTHORITY["EPSG","4326"]], PROJECTION["Polar Stereographic (variant B)", AUTHORITY["EPSG","9829"]], PARAMETER["central_meridian", 0.0], PARAMETER["Standard_Parallel_1", -71.0], PARAMETER["false_easting", 0.0], PARAMETER["false_northing", 0.0], UNIT["m", 1.0], AXIS["Easting", "North along 90 deg East"], AXIS["Northing", "North along 0 deg"], AUTHORITY["EPSG","3031"]]'


# =============================================================================
# MAIN
# =============================================================================

def main() -> None:
    gdal.UseExceptions()

    if not IMAGERY_DIR.is_dir():
        raise FileNotFoundError(f"Imagery directory not found:\n{IMAGERY_DIR}")

    shp_base = IMAGERY_DIR / OUTPUT_NAME

    writer = shapefile.Writer(
        str(shp_base),
        shapeType=shapefile.POLYGON,
        encoding="UTF-8",
    )

    # ImageMosaic-required fields.
    writer.field("location", "C", 254)
    writer.field("featname", "C", 64)
    writer.field("year", "N", 4, 0)

    tif_count = 0

    for shelf_dir in sorted(p for p in IMAGERY_DIR.iterdir() if p.is_dir()):
        input_dir = shelf_dir / INPUT_FOLDER

        if not input_dir.is_dir():
            continue

        featname = shelf_dir.name

        for tif in sorted(input_dir.glob("*.tif")):
            ds = gdal.Open(str(tif), gdal.GA_ReadOnly)

            if ds is None:
                print(f"[SKIP] unreadable: {tif.relative_to(IMAGERY_DIR)}")
                continue

            gt = ds.GetGeoTransform()

            def corner(px: int, py: int) -> list[float]:
                return [
                    gt[0] + px * gt[1] + py * gt[2],
                    gt[3] + px * gt[4] + py * gt[5],
                ]

            ul = corner(0, 0)
            ur = corner(ds.RasterXSize, 0)
            lr = corner(ds.RasterXSize, ds.RasterYSize)
            ll = corner(0, ds.RasterYSize)

            ds = None

            writer.poly([[
                ul,
                ur,
                lr,
                ll,
                ul,
            ]])

            # Path relative to the GeoServer mosaic directory (cwd/imagery).
            location = tif.relative_to(IMAGERY_DIR).as_posix()

            match = YEAR_RE.search(tif.name)
            year = int(match.group(1)) if match else None

            writer.record(
                location=location,
                featname=featname,
                year=year,
            )

            tif_count += 1

    writer.close()

    with open(
        f"{shp_base}.prj",
        "w",
        newline="",
        encoding="utf-8",
    ) as f:
        f.write(EPSG3031_WKT)

    print(f"Indexed {tif_count} GeoTIFF(s)")
    print(f"[OK] {shp_base.relative_to(BASE_DIR)}.shp")


if __name__ == "__main__":
    main()