#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List


def load_manifest(path: Path) -> Dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def resolve_dem_source(arg_value: str | None, manifest: Dict[str, Any]) -> Path:
    candidates: List[str] = []
    if arg_value:
        candidates.append(arg_value)
    resolved = manifest.get("resolved_dependencies", {})
    for key in ("dem_path_wsl", "dem_path_windows"):
        value = resolved.get(key)
        if value:
            candidates.append(value)

    for candidate in candidates:
        path = Path(candidate)
        if path.exists() and Path(str(path) + ".xml").exists():
            return path
    raise FileNotFoundError("Unable to resolve source DEM from arguments or stack manifest")


def compute_bbox(manifest: Dict[str, Any], margin_deg: float) -> List[float]:
    lons: List[float] = []
    lats: List[float] = []
    for scene in manifest["scenes"]:
        lon = scene.get("scene_center_lon")
        lat = scene.get("scene_center_lat")
        if lon is None or lat is None:
            source_scene_json = scene.get("source_scene_json_wsl") or scene.get("source_scene_json_windows")
            if source_scene_json and Path(source_scene_json).exists():
                source_payload = json.loads(Path(source_scene_json).read_text(encoding="utf-8"))
                lon = source_payload.get("scene_center_lon")
                lat = source_payload.get("scene_center_lat")
        if lon is not None and lat is not None:
            lons.append(float(lon))
            lats.append(float(lat))
    if not lons or not lats:
        raise ValueError("Stack manifest does not include usable scene center coordinates")
    west = min(lons) - margin_deg
    east = max(lons) + margin_deg
    south = min(lats) - margin_deg
    north = max(lats) + margin_deg
    return [south, north, west, east]


def prepare_dem(source_dem: Path, output_dem: Path, bbox: List[float]) -> None:
    from osgeo import gdal
    from isce.applications.gdal2isce_xml import gdal2isce_xml

    south, north, west, east = bbox
    src_open_path = Path(str(source_dem) + ".vrt")
    if not src_open_path.exists():
        src_open_path = source_dem

    output_dem.parent.mkdir(parents=True, exist_ok=True)
    output_vrt = Path(str(output_dem) + ".vrt")
    output_xml = Path(str(output_dem) + ".xml")
    output_hdr = Path(str(output_dem) + ".hdr")
    fallback_hdr = output_dem.with_suffix(".hdr")

    src_ds = gdal.Open(str(src_open_path), gdal.GA_ReadOnly)
    if src_ds is None:
        raise RuntimeError(f"Unable to open DEM source: {src_open_path}")

    translate_options = gdal.TranslateOptions(
        format="ENVI",
        projWin=[west, north, east, south],
    )
    out_ds = gdal.Translate(str(output_dem), src_ds, options=translate_options)
    if out_ds is None:
        raise RuntimeError("gdal.Translate failed while clipping the DEM")
    out_ds = None
    src_ds = None

    vrt_ds = gdal.Open(str(output_dem), gdal.GA_ReadOnly)
    if vrt_ds is None:
        raise RuntimeError(f"Unable to reopen clipped DEM: {output_dem}")
    gdal.Translate(str(output_vrt), vrt_ds, options=gdal.TranslateOptions(format="VRT"))
    vrt_ds = None

    gdal2isce_xml(str(output_vrt))
    if not output_xml.exists():
        raise RuntimeError(f"Expected ISCE XML was not created: {output_xml}")
    if not output_hdr.exists() and not fallback_hdr.exists():
        raise RuntimeError(f"Expected ENVI header was not created: {output_hdr} or {fallback_hdr}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Clip a local DEM window for one LT-1 stack workspace."
    )
    parser.add_argument(
        "--stack-manifest",
        required=True,
        help="Path to stack_input_manifest.json generated by build_lt1_stack_prep.py",
    )
    parser.add_argument(
        "--source-dem",
        default=None,
        help="Override source DEM base path.",
    )
    parser.add_argument(
        "--margin-deg",
        type=float,
        default=1.0,
        help="Margin around stack scene-center extents in degrees.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    manifest_path = Path(args.stack_manifest)
    if not manifest_path.exists():
        raise FileNotFoundError(f"Stack manifest not found: {manifest_path}")

    manifest = load_manifest(manifest_path)
    source_dem = resolve_dem_source(args.source_dem, manifest)
    bbox = compute_bbox(manifest, margin_deg=args.margin_deg)

    workspace = manifest["workspace"]
    dem_dir = Path(workspace["inputs_dir_wsl"]) / "dem"
    output_dem = dem_dir / "stack_dem_window.wgs84"
    prepare_dem(source_dem=source_dem, output_dem=output_dem, bbox=bbox)

    report = {
        "stack_manifest": str(manifest_path),
        "source_dem": str(source_dem),
        "output_dem": str(output_dem),
        "bbox_south_north_west_east": bbox,
    }
    report_path = dem_dir / "stack_dem_window_report.json"
    report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")

    print(f"Source DEM: {source_dem}")
    print(f"Output DEM: {output_dem}")
    print(f"BBox:       {bbox}")
    print(f"Report:     {report_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())