tileyolo 0.2.3

use crate::models::geometry::GeometryExtent;
use crate::models::layer::Layer;
use crate::{Config, utils::style::get_builtin_gradient};
use gdal::spatial_ref::SpatialRef;
use gdal::{Dataset, DriverManager, errors::GdalError};
use gdal_sys::{GDALReprojectImage, GDALResampleAlg};
use image::{ColorType, ImageEncoder, Rgba, RgbaImage, codecs::png::PngEncoder};
use proj::Proj;
use std::{io::Cursor, path::PathBuf};
use tokio::task;

// Returns true if the value should be treated as nodata (currently, if it is NaN)
fn is_nodata(val: f32) -> bool {
    val.is_nan()
}

pub async fn process_cog(
    input_path: PathBuf,
    extent_3857: GeometryExtent,
    layer_obj: Layer,
    tile_size: (usize, usize),
) -> gdal::errors::Result<Vec<u8>> {
    task::spawn_blocking(move || {
        let (tile_size_x, tile_size_y) = tile_size;
        let source_crs = format!("{}:{}", "EPSG", layer_obj.source_geometry.crs_code);
        let to_merc = Proj::new_known_crs(&source_crs, "EPSG:3857", None)
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;

        // Reproject both corners into 3857
        let (x0, y0) = to_merc
            .convert((
                layer_obj.source_geometry.extent.minx,
                layer_obj.source_geometry.extent.miny,
            ))
            .map_err(|e| GdalError::BadArgument(format!("failed to reproj min corner: {}", e)))?;
        let (x1, y1) = to_merc
            .convert((
                layer_obj.source_geometry.extent.maxx,
                layer_obj.source_geometry.extent.maxy,
            ))
            .map_err(|e| GdalError::BadArgument(format!("failed to reproj max corner: {}", e)))?;
        let orig_minx_3857 = x0.min(x1);
        let orig_maxx_3857 = x0.max(x1);
        let orig_miny_3857 = y0.min(y1);
        let orig_maxy_3857 = y0.max(y1);

        // Open source dataset, S3 is /vsis3/{bucket}/{key}, otherwise file.
        let src_ds = Dataset::open(&input_path)?;

        // Retrieve nodata value from the source dataset (band 1)
        let src_band = src_ds
            .rasterband(Config::default().default_raster_band)
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;
        let src_nodata_opt: Option<f32> = src_band.no_data_value().map(|v| v as f32);

        // Prepare an in‐memory 256×256 target in Web mercator 3857
        // let (minx, miny, maxx, maxy) = bbox_3857;
        let res_x = (extent_3857.maxx - extent_3857.minx) / (tile_size_x as f64);
        let res_y = (extent_3857.maxy - extent_3857.miny) / (tile_size_y as f64);

        let mem_drv = DriverManager::get_driver_by_name("MEM")
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;
        let mut dst_ds = mem_drv
            .create_with_band_type::<f32, _>(
                "memory_dataset",
                tile_size_x,
                tile_size_y,
                Config::default().default_raster_band,
            )
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;

        let merc_sref =
            SpatialRef::from_epsg(3857).map_err(|e| GdalError::BadArgument(e.to_string()))?;
        dst_ds
            .set_projection(
                &merc_sref
                    .to_wkt()
                    .map_err(|e| GdalError::BadArgument(e.to_string()))?,
            )
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;
        dst_ds
            .set_geo_transform(&[extent_3857.minx, res_x, 0.0, extent_3857.maxy, 0.0, -res_y])
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;

        // Set the nodata value for the destination raster band BEFORE reprojection
        if let Some(src_nodata) = src_nodata_opt {
            let mut dst_band = dst_ds
                .rasterband(Config::default().default_raster_band)
                .map_err(|e| GdalError::BadArgument(e.to_string()))?;
            dst_band
                .set_no_data_value(Some(src_nodata as f64))
                .map_err(|e| GdalError::BadArgument(e.to_string()))?;
        }

        // Setup reprojection of tile. Potential memory issues with unsafe code
        // however gdalwarp is not available in gdal crate as yet.
        unsafe {
            GDALReprojectImage(
                src_ds.c_dataset(),
                std::ptr::null(),
                dst_ds.c_dataset(),
                std::ptr::null(),
                GDALResampleAlg::GRA_NearestNeighbour,
                f64::NAN, // treat outside pixels as nodata
                f64::NAN,
                None,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
            );
        }

        let dst_band = dst_ds
            .rasterband(Config::default().default_raster_band)
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;

        // Read the warped 256×256 band into a buffer
        let mut buffer = dst_band
            .read_as::<f32>((0, 0), tile_size, tile_size, None)?
            .data()
            .to_vec();

        // Map nodata values (including 0.0) to NaN in the buffer used for rendering
        if let Some(src_nodata) = src_nodata_opt {
            for value in buffer.iter_mut() {
                if *value == src_nodata {
                    *value = f32::NAN;
                }
            }
        }
        // Also treat 0.0 as nodata
        for value in buffer.iter_mut() {
            if *value == 0.0 {
                *value = f32::NAN;
            }
        }

        // Any pixel whose geographic coordinate falls outside the original extent
        // should be treated as nodata (NaN), not 0.0.

        for y in 0..tile_size_y {
            for x in 0..tile_size_x {
                let gx = extent_3857.minx + (x as f64) * res_x;
                let gy = extent_3857.maxy - (y as f64) * res_y;
                if gx < orig_minx_3857
                    || gx > orig_maxx_3857
                    || gy < orig_miny_3857
                    || gy > orig_maxy_3857
                {
                    buffer[y * tile_size_x + x] = f32::NAN;
                }
            }
        }

        // Colourise into a 256×256 RGBA image
        let mut img = RgbaImage::new(tile_size_x as u32, tile_size_y as u32);

        if let Some(grad) = get_builtin_gradient(&layer_obj.style) {
            // Use the gradient to colourise the image
            for (i, &raw) in buffer.iter().enumerate() {
                let px = if is_nodata(raw) {
                    Rgba([0, 0, 0, 0])
                } else {
                    let t = ((raw - layer_obj.min_value)
                        / (layer_obj.max_value - layer_obj.min_value))
                        .clamp(0.0, 1.0);
                    let [r, g, b, a] = grad.at(t).to_rgba8();
                    Rgba([r, g, b, a])
                };
                let x = (i % tile_size_x) as u32;
                let y = (i / tile_size_y) as u32;
                img.put_pixel(x, y, px);
            }
        } else if layer_obj.colour_stops.is_empty() {
            // Fallback to grayscale
            for (i, &raw) in buffer.iter().enumerate() {
                let px = if is_nodata(raw) {
                    Rgba([0, 0, 0, 0])
                } else {
                    let norm =
                        (raw - layer_obj.min_value) / (layer_obj.max_value - layer_obj.min_value);
                    let lum = (norm.clamp(0.0, 1.0) * 255.0) as u8;
                    Rgba([lum, lum, lum, 255])
                };
                let x = (i % tile_size_x) as u32;
                let y = (i / tile_size_y) as u32;
                img.put_pixel(x, y, px);
            }
        } else {
            // Use the colour stops to colourise the image
            let cs = &layer_obj.colour_stops;
            let style_min = cs.first().unwrap().value;
            let style_max = cs.last().unwrap().value;
            for (i, &raw) in buffer.iter().enumerate() {
                let px = if is_nodata(raw) {
                    Rgba([0, 0, 0, 0])
                } else {
                    let norm =
                        (raw - layer_obj.min_value) / (layer_obj.max_value - layer_obj.min_value);
                    let scaled = style_min + norm.clamp(0.0, 1.0) * (style_max - style_min);
                    let mut colour = Rgba([0, 0, 0, 0]);
                    for w in cs.windows(2) {
                        let a = &w[0];
                        let b = &w[1];
                        if scaled >= a.value && scaled <= b.value {
                            let t = (scaled - a.value) / (b.value - a.value);
                            let r = ((1.0 - t) * a.red as f32 + t * b.red as f32) as u8;
                            let g = ((1.0 - t) * a.green as f32 + t * b.green as f32) as u8;
                            let b_ = ((1.0 - t) * a.blue as f32 + t * b.blue as f32) as u8;
                            let a_ = ((1.0 - t) * a.alpha as f32 + t * b.alpha as f32) as u8;
                            colour = Rgba([r, g, b_, a_]);
                            break;
                        }
                    }
                    colour
                };
                let x = (i % tile_size_x) as u32;
                let y = (i / tile_size_y) as u32;
                img.put_pixel(x, y, px);
            }
        }

        let mut png_data = Vec::new();
        PngEncoder::new(Cursor::new(&mut png_data))
            .write_image(
                img.as_raw(),
                tile_size_x as u32,
                tile_size_y as u32,
                ColorType::Rgba8.into(),
            )
            .map_err(|e| GdalError::BadArgument(e.to_string()))?;

        Ok(png_data)
    })
    .await
    .map_err(|e| GdalError::BadArgument(e.to_string()))?
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::{
        geometry::GeometryExtent,
        layer::{Layer, LayerGeometry},
        style::ColourStop,
    };
    use crate::reader::cog::process_cog;
    use gdal::spatial_ref::SpatialRef;
    use gdal::{Dataset, DriverManager};
    use image::{ColorType, ImageDecoder, codecs::png::PngDecoder};
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};
    use std::{fs, io::Cursor, path::PathBuf};
    use tempfile::TempDir;

    async fn make_layer(min_value: f32, max_value: f32) -> Layer {
        let path = PathBuf::new(); // will be set per-test
        let colour_stops = vec![
            ColourStop {
                value: 0.0,
                red: 215,
                green: 25,
                blue: 28,
                alpha: 255,
            },
            ColourStop {
                value: 100.0,
                red: 253,
                green: 174,
                blue: 97,
                alpha: 255,
            },
            ColourStop {
                value: 200.0,
                red: 255,
                green: 255,
                blue: 191,
                alpha: 255,
            },
            ColourStop {
                value: 300.0,
                red: 171,
                green: 221,
                blue: 164,
                alpha: 255,
            },
            ColourStop {
                value: 400.0,
                red: 43,
                green: 131,
                blue: 186,
                alpha: 255,
            },
        ];

        let source_geometry = LayerGeometry {
            crs_code: 3857,
            extent: GeometryExtent {
                minx: 0.0,
                miny: 0.0,
                maxx: 256.0,
                maxy: 256.0,
            },
        };
        let cached_geometry = source_geometry.generate_cached_geometry_sync().unwrap();
        Layer {
            layer: "test".to_string(),
            style: "default".to_string(),
            path,
            size_bytes: 0,
            source_geometry,
            cached_geometry,
            colour_stops,
            min_value,
            max_value,
            is_cog: true,
            last_modified: std::time::SystemTime::UNIX_EPOCH,
        }
    }

    /// Generates a temporary GeoTIFF in EPSG:3857 with reproducible random data,
    /// injecting ~10% NaN as no-data.
    fn generate_random_cog(tile_size: (usize, usize)) -> (TempDir, PathBuf) {
        let tmp = TempDir::new().expect("failed to create temp dir");
        let file_path = tmp.path().join("test.tif");

        let (tile_size_x, tile_size_y) = tile_size;

        let driver = DriverManager::get_driver_by_name("GTIFF").unwrap();
        let mut ds = driver
            .create_with_band_type::<f32, _>(
                file_path.to_str().unwrap(),
                tile_size_x,
                tile_size_y,
                1,
            )
            .unwrap();

        // Use Web Mercator so reprojection is identity
        let sref = SpatialRef::from_epsg(3857).unwrap();
        ds.set_projection(&sref.to_wkt().unwrap()).unwrap();
        ds.set_geo_transform(&[0.0, 1.0, 0.0, 0.0, 0.0, -1.0])
            .unwrap();

        // Fill with reproducible data and inject NaNs
        let mut rng = StdRng::seed_from_u64(42);
        let data: Vec<f32> = (0..tile_size_x * tile_size_y)
            .map(|_| {
                if rng.random_bool(0.1) {
                    f32::NAN
                } else {
                    rng.random_range(0.0..100.0)
                }
            })
            .collect();

        let mut band = ds
            .rasterband(Config::default().default_raster_band)
            .unwrap();

        let mut buffer = gdal::raster::Buffer::<f32>::new((tile_size_x, tile_size_y), data);
        band.write((0, 0), (tile_size_x, tile_size_y), &mut buffer)
            .unwrap();
        ds.flush_cache().unwrap();

        (tmp, file_path)
    }

    #[tokio::test]
    async fn test_process_cog_data_length() {
        let tile_size = (256, 256);
        let (tmp, path) = generate_random_cog(tile_size);
        let mut layer = make_layer(1.0, 100.0).await;
        layer.path = path.clone();
        layer.size_bytes = fs::metadata(&path).unwrap().len();

        let buffer = process_cog(
            path.clone(),
            (0.0, 256.0, 0.0, 256.0).into(),
            layer,
            tile_size,
        )
        .await
        .expect("process_cog should succeed");

        assert!(!buffer.is_empty(), "Output buffer must not be empty");
        let decoder = PngDecoder::new(Cursor::new(&buffer)).unwrap();
        assert_eq!(decoder.color_type(), ColorType::Rgba8, "Expected RGBA8");

        drop(tmp);
    }

    #[test]
    fn test_nodata_mask_generation() {
        let (tmp, path) = generate_random_cog((256, 256));
        let ds = Dataset::open(&path).unwrap();
        let band = ds
            .rasterband(Config::default().default_raster_band)
            .unwrap();

        let data: Vec<f32> = band
            .read_as::<f32>((0, 0), (10, 10), (10, 10), None)
            .unwrap()
            .data()
            .to_vec();

        let nodata_opt = band.no_data_value().map(|v| v as f32);
        let mask: Vec<bool> = data
            .iter()
            .map(|&v| v.is_nan() || nodata_opt.map(|nd| v == nd).unwrap_or(false))
            .collect();

        assert_eq!(mask.len(), data.len(), "Mask length must match data length");
        drop(tmp);
    }
}