eorst 1.0.0

//! GDAL utility functions for raster processing.
//!
//! This module provides helper functions for GDAL command-line operations
//! including warping, translating, creating VRTs, and mosaic building.

use crate::metadata::Extent;
use crate::types::{BlockSize, GeoTransform, ImageResolution};
use crate::core_types::RasterType;
use anyhow::Result;
use gdal::{Dataset, DatasetOptions, DriverManager, GdalOpenFlags};
use gdal::{raster::GdalDataType, raster::RasterCreationOptions, spatial_ref::SpatialRef};
use gdal::vector::{FieldValue, LayerAccess};
use kdam::rayon::prelude::*;
use std::env;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::Command;
use uuid::Uuid;

// ─── Reusable helpers extracted from duplicated patterns ───

/// Creates a rayon ThreadPool with the given number of CPUs.
///
/// Sets the `RAYON_NUM_THREADS` environment variable and builds the pool.
/// Replaces the duplicated pattern found in 11 locations across the codebase.
pub fn create_rayon_pool(n_cpus: usize) -> rayon::ThreadPool {
    unsafe { std::env::set_var("RAYON_NUM_THREADS", n_cpus.to_string()) };
    rayon::ThreadPoolBuilder::new()
        .num_threads(n_cpus)
        .build()
        .expect("Failed to create rayon thread pool")
}

/// Extracts the file stem as a `&str` from a path.
///
/// Replaces the duplicated pattern found in 17 locations across the codebase.
pub fn file_stem_str(path: &Path) -> &str {
    path.file_stem()
        .expect("Path has no file stem")
        .to_str()
        .expect("File stem is not valid UTF-8")
}

/// Opens a GDAL dataset with update access.
///
/// Replaces the duplicated `DatasetOptions { open_flags: GDAL_OF_UPDATE, .. }` pattern
/// found in `blocks.rs::write_samples_feature()`, `blocks.rs::write3()`,
/// and `io.rs::write_window3()`.
pub fn open_for_update(path: &Path) -> Dataset {
    let opts = DatasetOptions {
        open_flags: GdalOpenFlags::GDAL_OF_UPDATE,
        ..DatasetOptions::default()
    };
    Dataset::open_ex(path, opts).expect("Failed to open dataset for update")
}

/// Writes each band of a 3D array to a GDAL dataset.
///
/// Replaces the duplicated band-write loop found in `blocks.rs::write3()`
/// and `io.rs::write_window3()`.
pub fn write_bands_to_file<T: RasterType>(
    out_ds: &Dataset,
    data: ndarray::ArrayView3<T>,
    write_offset: (isize, isize),
    write_size: (usize, usize),
) {
    use gdal::raster::Buffer;
    use ndarray::s;
    for band in 0..data.shape()[0] {
        let b = (band + 1) as isize;
        let mut out_band = out_ds.rasterband(b as usize).expect("Failed to get raster band");
        let data_vec: Vec<T> = data.slice(s![band, .., ..]).into_iter().copied().collect();
        let mut data_buffer = Buffer::new(write_size, data_vec);
        out_band
            .write(write_offset, write_size, &mut data_buffer)
            .expect("Failed to write band");
    }
}

/// Runs a GDAL command-line tool with the given arguments.
///
/// Replaces the duplicated `.spawn().expect().wait().expect()` pattern found in 3+ locations.
pub fn run_gdal_command(argv: &[&str]) {
    Command::new(argv[0])
        .args(&argv[1..])
        .spawn()
        .expect("failed to start gdal command")
        .wait()
        .expect("failed to wait for gdal command");
}

/// Reads a raster band window into a 2D array.
///
/// Replaces the duplicated `Dataset::open` → `rasterband` → `read_as` → `to_array` pattern
/// found in `io.rs::read_block()`, `io.rs::read_block_layer_idx()`, and
/// `processing.rs::read_block_hybrid()`.
pub fn read_raster_band<T: RasterType>(
    raster_path: &Path,
    band_index: usize,
    offset: (isize, isize),
    window_size: (usize, usize),
) -> ndarray::Array2<T> {
    let ds = Dataset::open(raster_path).expect(&format!("Unable to open {:?}", raster_path));
    let raster_band = ds.rasterband(band_index).expect("Failed to get raster band");
    let array_size = window_size;
    let e_resample_alg = None;
    raster_band
        .read_as::<T>(offset, window_size, array_size, e_resample_alg)
        .expect("Failed to read raster band")
        .to_array()
        .expect("Failed to convert to array")
}

/// Assembles processed block files into a final output via mosaic + translate + cleanup.
///
/// This replaces the duplicated 4-line pattern (mosaic, translate, remove VRT, remove blocks)
/// found in 9 locations across the codebase.
pub fn mosaic_translate_cleanup(
    collected: &[PathBuf],
    tmp_file: &Path,
    out_file: &Path,
    epsg_code: u32,
) {
    mosaic(collected, tmp_file, epsg_code).expect("Could not mosaic to vrt");
    translate(tmp_file, out_file).expect("Could not translate to geotiff");
    std::fs::remove_file(tmp_file).expect("Unable to remove the temporary file");
    collected
        .iter()
        .for_each(|f| std::fs::remove_file(f).expect("Unable to remove file"));
}

/// Assembles time-step block files into a final output via mosaic + translate + cleanup.
///
/// Variant for `Vec<Vec<PathBuf>>` where each inner vec corresponds to a time step.
pub fn mosaic_translate_cleanup_time_steps(
    collected: &[Vec<PathBuf>],
    out_file: &Path,
    epsg_code: u32,
    n_times: usize,
) {
    (0..n_times).into_par_iter().for_each(|time_index| {
        let mut block_fns = Vec::new();
        for block in collected.iter() {
            let block_fn = block[time_index].to_owned();
            block_fns.push(block_fn);
        }
        let tmp_vrt = PathBuf::from(create_temp_file("vrt"));
        let file_stem = file_stem_str(out_file);
        let time_fn = out_file.with_file_name(format!("{}_{}.tif", file_stem, time_index));
        mosaic(&block_fns, &tmp_vrt, epsg_code).expect("Could not mosaic to vrt");
        translate(&tmp_vrt, &time_fn).expect("Could not translate to geotiff");
        std::fs::remove_file(tmp_vrt).expect("Unable to remove the temporary file");
        block_fns
            .iter()
            .for_each(|f| std::fs::remove_file(f).expect("Unable to remove file"));
    });
}

/// Creates a temporary file with the given extension.
///
/// Uses the `TMP_DIR` environment variable if set, otherwise falls back to `/tmp`.
/// File names are prefixed with `eorst_` for easy identification.
pub fn create_temp_file(ext: &str) -> String {
    let dir = env::var("TMP_DIR").unwrap_or("/tmp".to_string());
    let dir = Path::new(&dir);
    let file_name = format!("eorst_{}.{}", Uuid::new_v4().simple(), ext);
    let file_name = dir.join(file_name);
    file_name.into_os_string().into_string().unwrap()
}

/// Warps a raster source to a different EPSG projection.
pub(crate) fn warp(
    source: PathBuf,
    target_resolution: Option<ImageResolution>,
    target_epsg: u32,
) -> PathBuf {
    let new_source = create_temp_file("vrt");
    let mut args: Vec<String> = vec!["gdalwarp".to_string(), "-q".to_string()];
    if let Some(tr) = target_resolution {
        args.extend([
            "-tr".to_string(),
            format!("{}", tr.x),
            format!("{}", tr.y),
            "-tap".to_string(),
        ]);
    }
    args.extend([
        "-t_srs".to_string(),
        format!("EPSG:{}", target_epsg),
        source.to_string_lossy().to_string(),
        new_source.clone(),
    ]);
    let argv: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
    run_gdal_command(&argv);
    PathBuf::from(new_source)
}

/// Warps a raster source to a different extent and resolution.
pub(crate) fn warp_with_te_tr(
    source: PathBuf,
    target_te: &Extent,
    target_resolution: ImageResolution,
) -> PathBuf {
    let new_source = create_temp_file("vrt");
    let xmin_s = format!("{}", (target_te.xmin * 100.).round() / 100.);
    let ymin_s = format!("{}", (target_te.ymin * 100.).round() / 100.);
    let xmax_s = format!("{}", (target_te.xmax * 100.).round() / 100.);
    let ymax_s = format!("{}", (target_te.ymax * 100.).round() / 100.);
    let trx_s = format!("{}", (target_resolution.x * 100.).round() / 100.);
    let try_s = format!("{}", (target_resolution.y * 100.).round() / 100.);
    let source_s = source.to_string_lossy();
    let argv = vec![
        "gdalwarp", "-q", "-te", &xmin_s, &ymin_s, &xmax_s, &ymax_s,
        "-tr", &trx_s, &try_s, "-r", "nearest",
        &source_s, &new_source,
    ];
    run_gdal_command(&argv);
    PathBuf::from(new_source)
}

/// Changes the resolution of a raster source.
pub(crate) fn change_res(source: PathBuf, target_resolution: ImageResolution) -> PathBuf {
    let new_source = create_temp_file("vrt");
    let trx = format!("{}", target_resolution.x);
    let try_ = format!("{}", target_resolution.y);
    let source_s = source.to_string_lossy();
    let argv = vec![
        "gdalwarp", "-q", "-tr", &trx, &try_, &source_s, &new_source,
    ];
    run_gdal_command(&argv);
    PathBuf::from(new_source)
}

/// Extracts a single band from a raster source into a VRT.
///
/// Uses GDAL's VRT driver in-process instead of spawning a `gdal_translate`
/// subprocess, which reduces overhead from ~100ms to ~0.1ms per band.
pub(crate) fn extract_band(source: &Path, band: usize) -> PathBuf {
    // Open source to extract metadata
    let src_ds = match Dataset::open(source) {
        Ok(ds) => ds,
        Err(_) => {
            // Fallback to gdal_translate if we can't open in-process
            let new_source = create_temp_file("vrt");
            let argv = &[
                "gdal_translate",
                "-b",
                &format!("{}", band),
                "-q",
                source.to_str().unwrap(),
                &new_source,
            ];
            run_gdal_command(argv);
            return PathBuf::from(new_source);
        }
    };

    let (xsize, ysize) = src_ds.raster_size();
    let gt = src_ds.geo_transform().ok();
    let srs_wkt = src_ds.spatial_ref().ok().map(|s| s.to_wkt().ok()).flatten();

    let band_idx = band;
    let band_meta = match src_ds.rasterband(band_idx) {
        Ok(b) => {
            let dtype = b.band_type();
            let no_data = b.no_data_value();
            (Some(dtype), no_data)
        }
        Err(_) => (None, None),
    };

    // Build VRT XML manually — no subprocess needed
    let wkt_part = match &srs_wkt {
        Some(wkt) => format!("  <SRS>{}</SRS>\n", wkt),
        None => String::new(),
    };
    let gt_part = match &gt {
        Some(arr) => format!(
            "  <GeoTransform> {:.15}, {:.15}, {:.15}, {:.15}, {:.15}, {:.15} </GeoTransform>\n",
            arr[0], arr[1], arr[2], arr[3], arr[4], arr[5]
        ),
        None => String::new(),
    };
    let nd_part = match band_meta.1 {
        Some(nd) => format!("<NoDataValue>{}</NoDataValue>\n", nd),
        None => String::new(),
    };

    let source_abs = std::fs::canonicalize(source)
        .unwrap_or_else(|_| source.to_path_buf())
        .to_string_lossy()
        .to_string();

    // Use Display impl for data type string (e.g., "UInt16", "Float32")
    let dtype_str = match band_meta.0 {
        Some(dt) => format!("{}", dt),
        None => "Unknown".to_string(),
    };

    let vrt_xml = format!(
        r#"<VRTDataset rasterXSize="{}" rasterYSize="{}">
{wkt_part}{gt_part}  <VRTRasterBand dataType="{dtype}" band="1">
    <SimpleSource>
      <SourceFilename relativeToVRT="0">{source}</SourceFilename>
      <SourceBand>{band}</SourceBand>
      {nd_part}    </SimpleSource>
  </VRTRasterBand>
</VRTDataset>"#,
        xsize, ysize,
        wkt_part = wkt_part,
        gt_part = gt_part,
        dtype = dtype_str,
        source = source_abs,
        band = band,
        nd_part = nd_part,
    );

    let new_source = PathBuf::from(create_temp_file("vrt"));
    std::fs::write(&new_source, vrt_xml).expect("Failed to write VRT XML");
    PathBuf::from(new_source)
}

/// Creates an empty raster file with the given size and properties.
pub fn raster_from_size<T>(
    file_name: &Path,
    geo_transform: GeoTransform,
    epsg_code: u32,
    block_size: BlockSize,
    n_bands: isize,
    na_value: T,
) where
    T: RasterType,
{
    let parent_path = file_name.parent().unwrap();
    std::fs::create_dir(parent_path).unwrap_or(());

    let size_x = block_size.cols;
    let size_y = block_size.rows;
    let driver = DriverManager::get_driver_by_name("GTIFF").unwrap();
    let options =
        RasterCreationOptions::from_iter(["COMPRESS=LZW", "BLOCKXSIZE=512", "BLOCKYSIZE=512"]);

    let mut dataset = driver
        .create_with_band_type_with_options::<T, _>(
            file_name,
            size_x,
            size_y,
            n_bands as usize,
            &options,
        )
        .unwrap();
    dataset
        .set_geo_transform(&geo_transform.to_array())
        .unwrap();
    let srs = SpatialRef::from_epsg(epsg_code).unwrap();
    dataset.set_spatial_ref(&srs).unwrap();
    for band_index in 1..n_bands + 1 {
        let mut raster_band = dataset.rasterband(band_index as usize).unwrap();
        let no_data_f64 = na_value
            .to_f64()
            .expect("Failed to convert no_data value to f64");
        raster_band.set_no_data_value(Some(no_data_f64)).unwrap();
    }
}

/// Creates a mosaic from multiple raster files.
pub fn mosaic(collected: &[PathBuf], tmp_file: &Path, epsg_code: u32) -> Result<()> {
    let collected_reproj: Vec<PathBuf> = collected
        .par_iter()
        .map(|image| {
            let new_source = create_temp_file("vrt");
            let epsg_s = format!("EPSG:{}", epsg_code);
            let image_s = image.to_string_lossy();
            let argv = vec![
                "gdalwarp", "-q", "-t_srs", &epsg_s, &image_s, &new_source,
            ];
            run_gdal_command(&argv);
            PathBuf::from(new_source)
        })
        .collect();

    let mut argv: Vec<String> = vec![
        "gdalbuildvrt".to_string(),
        "-q".to_string(),
        tmp_file.to_string_lossy().to_string(),
    ];
    argv.extend(collected_reproj.iter().map(|p| p.to_string_lossy().to_string()));
    let argv_refs: Vec<&str> = argv.iter().map(|s| s.as_str()).collect();
    run_gdal_command(&argv_refs);

    Ok(())
}

/// Translates a raster file to a different format using the specified driver.
pub fn translate_with_driver(tmp_fn: &Path, image_fn: &Path, driver_name: &str) -> Result<()> {
    let argv = vec![
        "gdal_translate",
        "-q",
        "-of",
        driver_name,
        "-co",
        "BIGTIFF=YES",
        "-co",
        "COMPRESS=DEFLATE",
        "-co",
        "NUM_THREADS=16",
        tmp_fn.to_str().unwrap(),
        image_fn.to_str().unwrap(),
    ];
    run_gdal_command(&argv);
    Ok(())
}

/// Translates a raster file to GeoTIFF format.
pub fn translate(tmp_fn: &Path, image_fn: &Path) -> Result<()> {
    translate_with_driver(tmp_fn, image_fn, "GTiff").unwrap();
    Ok(())
}

/// Returns the widest GDAL data type among all raster bands in the dataset.
#[allow(dead_code)]
pub(crate) fn get_widest_type(source: &PathBuf) -> GdalDataType {
    use log::warn;

    let dataset = Dataset::open(source).unwrap();

    let mut widest: Option<GdalDataType> = None;

    for i in 1..=dataset.raster_count() {
        let band = dataset.rasterband(i).expect("Failed to read band");
        let dtype = band.band_type();

        if let Some(existing) = widest {
            if dtype != existing {
                warn!(
                    "Band {} has different type ({:?}) than first band ({:?})",
                    i, dtype, existing
                );
            }
            widest = Some(existing.union(dtype));
        } else {
            widest = Some(dtype);
        }
    }

    widest.expect("Dataset has no bands")
}

/// Gets class IDs from a vector dataset.
#[allow(dead_code)]
pub fn get_class(
    dataset_path: &PathBuf,
    id_column: &str,
    class_column: &str,
) -> BTreeMap<i16, i16> {
    let dataset = Dataset::open(dataset_path).unwrap();
    let mut layer = dataset.layer(0).unwrap();

    let _fields_defn = layer
        .defn()
        .fields()
        .map(|field| (field.name(), field.field_type(), field.width()))
        .collect::<Vec<_>>();
    let mut class: BTreeMap<i16, i16> = BTreeMap::new();

    for feature in layer.features() {
        let id_column_idx = feature.field_index(id_column).expect("Bad column name");
        let class_column_idx = feature.field_index(class_column).expect("Bad column name");
        let id = feature
            .field(id_column_idx)
            .unwrap()
            .unwrap()
            .into_int()
            .unwrap();
        let condition = feature
            .field(class_column_idx)
            .unwrap()
            .unwrap_or(FieldValue::IntegerValue(-1))
            .into_int()
            .unwrap();
        class.insert(id as i16, condition as i16);
    }
    class
}

/// Basic metadata extracted from a single GDAL dataset open.
///
/// Consolidates the 6 separate `Dataset::open` calls in `builder.rs`
/// (`get_resolution`, `get_geotransform`, `get_extent`, `get_epsg_code`,
/// `get_na_value`, `get_image_size`) into a single dataset open.
#[derive(Debug, Clone)]
pub struct BasicRasterInfo {
    /// GeoTransform (x_ul, x_res, x_rot, y_ul, y_rot, y_res)
    pub geo_transform: [f64; 6],
    /// Image size in pixels (cols, rows)
    pub size: (usize, usize),
    /// EPSG code (0 if unknown)
    pub epsg_code: u32,
    /// No-data value (None if not set)
    pub no_data: Option<f64>,
    /// Number of raster bands
    pub n_bands: usize,
}

impl BasicRasterInfo {
    /// Resolution (x, y) derived from geo_transform.
    pub fn resolution(&self) -> crate::types::ImageResolution {
        crate::types::ImageResolution {
            x: self.geo_transform[1],
            y: self.geo_transform[5],
        }
    }

    /// GeoTransform struct derived from geo_transform array.
    pub fn geo_transform_struct(&self) -> crate::types::GeoTransform {
        crate::types::GeoTransform {
            x_ul: self.geo_transform[0],
            x_res: self.geo_transform[1],
            x_rot: self.geo_transform[2],
            y_ul: self.geo_transform[3],
            y_rot: self.geo_transform[4],
            y_res: self.geo_transform[5],
        }
    }

    /// ImageSize struct derived from size tuple.
    pub fn image_size(&self) -> crate::types::ImageSize {
        crate::types::ImageSize {
            rows: self.size.1,
            cols: self.size.0,
        }
    }

    /// No-data value cast to type T, or T::zero() if not set.
    pub fn na_value<T: crate::core_types::RasterType>(&self) -> T {
        self.no_data
            .and_then(|v| num_traits::NumCast::from(v))
            .unwrap_or(T::zero())
    }
}

/// Opens a raster dataset once and extracts all basic metadata fields.
///
/// This replaces the pattern of 6 separate `Dataset::open` calls in builder.rs
/// with a single open, reducing I/O overhead.
pub fn read_basic_raster_info(source: &Path) -> BasicRasterInfo {
    let ds = Dataset::open(source).expect(&format!("Unable to open {:?}", source));
    let geo_transform = ds.geo_transform().expect("Failed to get geo transform");
    let size = ds.raster_size();
    let spatial_ref = ds.spatial_ref().expect("Failed to get spatial ref");
    let epsg_code = spatial_ref.auth_code().unwrap_or(0);

    // Extract no-data value from first band
    let mut no_data = None;
    if size.0 > 0 && size.1 > 0 {
        if let Ok(band) = ds.rasterband(1) {
            no_data = band.no_data_value();
        }
    }

    BasicRasterInfo {
        geo_transform,
        size,
        epsg_code: epsg_code as u32,
        no_data,
        n_bands: ds.raster_count(),
    }
}