#!/usr/bin/env python3
"""Rebuild river-masked orthophotos from existing tile masks."""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import numpy as np
from osgeo import gdal

from geodata_pipeline.config import Config
from geodata_pipeline.river_erosion import _apply_water_mask_to_ortho


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Rebuild export_unity/ortho_jpg_river from work/river_masks and export_unity/ortho_jpg."
    )
    p.add_argument("--config", default="geodata_config.toml", help="Path to pipeline config TOML.")
    p.add_argument("--manifest", default="", help="Optional manifest CSV override.")
    p.add_argument("--mask-dir", default="", help="Optional mask directory override.")
    p.add_argument("--output-dir", default="", help="Optional output ortho_jpg_river directory override.")
    p.add_argument("--skip-existing", action="store_true", help="Skip tiles that already exist in output.")
    p.add_argument("--mask-res", type=int, default=0, help="Fallback mask resolution when a tile mask is missing.")
    return p.parse_args()


def load_tile_ids(manifest_path: Path) -> list[str]:
    with manifest_path.open("r", encoding="utf-8", newline="") as f:
        rows = list(csv.DictReader(f))
    return [row["tile_id"].strip() for row in rows if row.get("tile_id")]


def load_mask(mask_path: Path) -> np.ndarray | None:
    ds = gdal.Open(str(mask_path), gdal.GA_ReadOnly)
    if ds is None:
        return None
    arr = ds.ReadAsArray()
    ds = None
    if arr is None:
        return None
    if arr.ndim == 3:
        arr = arr[0]
    arr = arr.astype(np.float32)
    if arr.max() > 1.0:
        arr /= 255.0
    return np.clip(arr, 0.0, 1.0)


def main() -> int:
    args = parse_args()
    cfg = Config.load(args.config)
    gdal.UseExceptions()

    manifest_path = Path(args.manifest or cfg.export.manifest_path)
    if not manifest_path.exists():
        raise SystemExit(f"[rebuild_ortho_jpg_river] Missing manifest: {manifest_path}")

    mask_dir = Path(args.mask_dir or Path(cfg.work.work_dir) / "river_masks")
    output_dir = Path(args.output_dir or cfg.ortho.river_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    fallback_res = int(args.mask_res or cfg.heightmap.out_res)
    if fallback_res <= 0:
        fallback_res = 1025

    tile_ids = load_tile_ids(manifest_path)
    written = 0
    skipped = 0
    with_mask = 0
    without_mask = 0
    missing_source = 0

    print(f"[rebuild_ortho_jpg_river] Tiles in manifest: {len(tile_ids)}")
    print(f"[rebuild_ortho_jpg_river] Mask dir: {mask_dir}")
    print(f"[rebuild_ortho_jpg_river] Output dir: {output_dir}")

    for tile_id in tile_ids:
        out_path = output_dir / f"{tile_id}.jpg"
        if args.skip_existing and out_path.exists():
            skipped += 1
            continue

        source_path = Path(cfg.export.ortho_dir) / f"{tile_id}.jpg"
        if not source_path.exists():
            missing_source += 1
            print(f"[rebuild_ortho_jpg_river] Missing source ortho for {tile_id}: {source_path}")
            continue

        mask_path = mask_dir / f"{tile_id}_mask.png"
        mask = load_mask(mask_path)
        if mask is None:
            without_mask += 1
            mask = np.zeros((fallback_res, fallback_res), dtype=np.float32)
        else:
            with_mask += 1

        _apply_water_mask_to_ortho(tile_id, mask, cfg)
        if out_path.exists():
            written += 1

    print(
        "[rebuild_ortho_jpg_river] Summary: "
        f"written={written}, skipped={skipped}, with_mask={with_mask}, "
        f"without_mask={without_mask}, missing_source={missing_source}"
    )
    return 0 if missing_source == 0 else 2


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