#!/usr/bin/env python3
"""Generate a first real-elevation heightmap for the Ground Zero terrain tile.

This prototype intentionally uses only the Python standard library so it can run
on Ubuntu-Codex without extra GIS packages. It samples USGS EPQS elevations over
the selected 1 km UTM tile, writes CSV samples, writes a little-endian R16
heightmap, and records generation metadata.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import struct
import time
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[1]
REGISTRY_PATH = PROJECT_ROOT / "Data" / "Tiles" / "ground_zero_tiles.json"
OUTPUT_ROOT = PROJECT_ROOT / "Data" / "Terrain" / "Generated"
USGS_EPQS_URL = "https://epqs.nationalmap.gov/v1/json"
TARGET_TILE_ID = "gz_us_ca_pacifica_utm10n_e544_n4160"


def utm_to_lat_lon(easting: float, northing: float, zone_number: int, northern: bool = True) -> tuple[float, float]:
    """Convert UTM WGS84 coordinates to latitude/longitude in degrees."""
    a = 6378137.0
    ecc_squared = 0.0066943799901413165
    k0 = 0.9996

    x = easting - 500000.0
    y = northing
    if not northern:
        y -= 10000000.0

    lon_origin = (zone_number - 1) * 6 - 180 + 3
    ecc_prime_squared = ecc_squared / (1 - ecc_squared)
    m = y / k0
    mu = m / (a * (1 - ecc_squared / 4 - 3 * ecc_squared**2 / 64 - 5 * ecc_squared**3 / 256))

    e1 = (1 - math.sqrt(1 - ecc_squared)) / (1 + math.sqrt(1 - ecc_squared))
    j1 = 3 * e1 / 2 - 27 * e1**3 / 32
    j2 = 21 * e1**2 / 16 - 55 * e1**4 / 32
    j3 = 151 * e1**3 / 96
    j4 = 1097 * e1**4 / 512
    fp = mu + j1 * math.sin(2 * mu) + j2 * math.sin(4 * mu) + j3 * math.sin(6 * mu) + j4 * math.sin(8 * mu)

    sin_fp = math.sin(fp)
    cos_fp = math.cos(fp)
    tan_fp = math.tan(fp)

    c1 = ecc_prime_squared * cos_fp**2
    t1 = tan_fp**2
    n1 = a / math.sqrt(1 - ecc_squared * sin_fp**2)
    r1 = a * (1 - ecc_squared) / ((1 - ecc_squared * sin_fp**2) ** 1.5)
    d = x / (n1 * k0)

    lat = fp - (n1 * tan_fp / r1) * (
        d**2 / 2
        - (5 + 3 * t1 + 10 * c1 - 4 * c1**2 - 9 * ecc_prime_squared) * d**4 / 24
        + (61 + 90 * t1 + 298 * c1 + 45 * t1**2 - 252 * ecc_prime_squared - 3 * c1**2) * d**6 / 720
    )
    lon = math.radians(lon_origin) + (
        d
        - (1 + 2 * t1 + c1) * d**3 / 6
        + (5 - 2 * c1 + 28 * t1 - 3 * c1**2 + 8 * ecc_prime_squared + 24 * t1**2) * d**5 / 120
    ) / cos_fp

    return math.degrees(lat), math.degrees(lon)


def query_elevation(latitude: float, longitude: float, retries: int = 3) -> dict:
    params = urllib.parse.urlencode({"x": f"{longitude:.8f}", "y": f"{latitude:.8f}", "units": "Meters"})
    url = f"{USGS_EPQS_URL}?{params}"
    last_error = None

    for attempt in range(retries):
        try:
            with urllib.request.urlopen(url, timeout=20) as response:
                payload = json.loads(response.read().decode("utf-8"))
                return {
                    "elevation_m": float(payload["value"]),
                    "raster_id": payload.get("rasterId"),
                    "resolution_m": payload.get("resolution"),
                    "query_url": url,
                }
        except (urllib.error.URLError, TimeoutError, KeyError, ValueError) as exc:
            last_error = exc
            time.sleep(0.5 * (attempt + 1))

    raise RuntimeError(f"USGS elevation query failed for {latitude}, {longitude}: {last_error}")


def load_ground_zero_tile() -> dict:
    registry = json.loads(REGISTRY_PATH.read_text())
    for tile in registry["tiles"]:
        if tile["tile_id"] == TARGET_TILE_ID:
            return tile
    raise RuntimeError(f"Could not find {TARGET_TILE_ID} in {REGISTRY_PATH}")


def build_sample_points(tile: dict, grid_size: int) -> list[dict]:
    grid = tile["grid"]
    e_min = grid["easting_min_m"]
    n_min = grid["northing_min_m"]
    spacing = grid["tile_size_m"] / (grid_size - 1)
    points = []

    for row in range(grid_size):
        northing = n_min + row * spacing
        for col in range(grid_size):
            easting = e_min + col * spacing
            latitude, longitude = utm_to_lat_lon(easting, northing, 10, True)
            points.append(
                {
                    "row": row,
                    "col": col,
                    "easting_m": easting,
                    "northing_m": northing,
                    "latitude": latitude,
                    "longitude": longitude,
                }
            )

    return points


def sample_elevations(points: list[dict], workers: int) -> list[dict]:
    sampled = []
    with ThreadPoolExecutor(max_workers=workers) as executor:
        future_to_point = {
            executor.submit(query_elevation, point["latitude"], point["longitude"]): point for point in points
        }
        for future in as_completed(future_to_point):
            point = future_to_point[future]
            elevation = future.result()
            point.update(elevation)
            sampled.append(point)

    return sorted(sampled, key=lambda item: (item["row"], item["col"]))


def write_csv(samples: list[dict], output_path: Path) -> None:
    fieldnames = [
        "row",
        "col",
        "easting_m",
        "northing_m",
        "latitude",
        "longitude",
        "elevation_m",
        "raster_id",
        "resolution_m",
    ]
    with output_path.open("w", newline="") as output_file:
        writer = csv.DictWriter(output_file, fieldnames=fieldnames)
        writer.writeheader()
        for sample in samples:
            writer.writerow({field: sample.get(field, "") for field in fieldnames})


def write_r16(samples: list[dict], output_path: Path) -> tuple[float, float]:
    elevations = [sample["elevation_m"] for sample in samples]
    min_elevation = min(elevations)
    max_elevation = max(elevations)
    elevation_range = max(max_elevation - min_elevation, 0.001)

    with output_path.open("wb") as output_file:
        for sample in samples:
            normalized = (sample["elevation_m"] - min_elevation) / elevation_range
            output_file.write(struct.pack("<H", round(normalized * 65535)))

    return min_elevation, max_elevation


def write_metadata(tile: dict, samples: list[dict], grid_size: int, min_elevation: float, max_elevation: float, output_path: Path) -> None:
    raster_ids = sorted({sample["raster_id"] for sample in samples if sample.get("raster_id") is not None})
    resolutions = sorted({sample["resolution_m"] for sample in samples if sample.get("resolution_m") is not None})
    metadata = {
        "tile_id": tile["tile_id"],
        "generated_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "source": {
            "name": "USGS Elevation Point Query Service",
            "url": USGS_EPQS_URL,
            "units": "Meters",
            "raster_ids": raster_ids,
            "reported_resolutions_m": resolutions,
        },
        "grid": tile["grid"],
        "sample_grid": {
            "width": grid_size,
            "height": grid_size,
            "spacing_m": tile["grid"]["tile_size_m"] / (grid_size - 1),
            "sample_count": len(samples),
        },
        "heightmap": {
            "format": "r16_little_endian_unsigned",
            "min_elevation_m": min_elevation,
            "max_elevation_m": max_elevation,
            "vertical_range_m": max_elevation - min_elevation,
        },
        "unreal_import_notes": [
            "Prototype only; final landscape import should use a higher-resolution DEM raster or lidar-derived terrain.",
            "R16 values are normalized from min_elevation_m to max_elevation_m.",
            "Use the metadata min/max values to restore vertical scale during Unreal import.",
            "Horizontal tile size is 1000 m x 1000 m."
        ],
    }
    output_path.write_text(json.dumps(metadata, indent=2) + "\n")


def update_registry(tile_id: str) -> None:
    registry = json.loads(REGISTRY_PATH.read_text())
    for tile in registry["tiles"]:
        if tile["tile_id"] == tile_id:
            tile["status"] = "generated"
            tile["generation_version"] = max(tile.get("generation_version", 0), 1)
            tile["package_version"] = max(tile.get("package_version", 0), 0)
            tile["notes"] = (
                "Prototype USGS EPQS elevation heightmap generated. "
                "Use Data/Terrain/Generated/gz_us_ca_pacifica_utm10n_e544_n4160 for samples, R16 heightmap, and metadata."
            )
            break
    else:
        raise RuntimeError(f"Could not update missing tile {tile_id}")

    REGISTRY_PATH.write_text(json.dumps(registry, indent=2) + "\n")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--grid-size", type=int, default=33, help="Sample grid width/height. Default: 33")
    parser.add_argument("--workers", type=int, default=8, help="Concurrent USGS requests. Default: 8")
    args = parser.parse_args()

    if args.grid_size < 2:
        raise SystemExit("--grid-size must be at least 2")

    tile = load_ground_zero_tile()
    points = build_sample_points(tile, args.grid_size)
    output_dir = OUTPUT_ROOT / tile["tile_id"]
    output_dir.mkdir(parents=True, exist_ok=True)

    samples = sample_elevations(points, args.workers)
    csv_path = output_dir / f"{tile['tile_id']}_elevation_samples_{args.grid_size}.csv"
    r16_path = output_dir / f"{tile['tile_id']}_heightmap_{args.grid_size}.r16"
    metadata_path = output_dir / f"{tile['tile_id']}_terrain_metadata.json"

    write_csv(samples, csv_path)
    min_elevation, max_elevation = write_r16(samples, r16_path)
    write_metadata(tile, samples, args.grid_size, min_elevation, max_elevation, metadata_path)
    update_registry(tile["tile_id"])

    print(f"Generated {len(samples)} samples for {tile['tile_id']}")
    print(f"Elevation range: {min_elevation:.3f}m to {max_elevation:.3f}m")
    print(f"CSV: {csv_path}")
    print(f"R16: {r16_path}")
    print(f"Metadata: {metadata_path}")


if __name__ == "__main__":
    main()
