eorst 1.0.1

//! RasterDatasetBuilder for creating and configuring raster datasets.
//!
//! This module provides the builder pattern for constructing [RasterDataset](crate::rasterdataset::RasterDataset)
//! instances from various data sources including files, STAC collections, or from scratch.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use log::debug;
use math::round;
use num_traits::{NumCast, ToPrimitive};
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
use stac::ItemCollection;

use crate::core_types::RasterType;
use crate::data_sources::{DataSource, DateType};
use crate::metadata::{Layer, RasterMetadata};
use crate::rasterdataset::RasterDataset;
use crate::types::{BlockSize, Dimension, GeoTransform, ImageResolution, ImageSize, Overlap, RasterDataShape, ReadWindow, Offset, Size};
use crate::gdal_utils::{BasicRasterInfo, create_temp_file, warp, warp_with_te_tr, change_res, extract_band, mosaic, read_basic_raster_info, run_gdal_command};

/// Main builder struct for constructing RasterDataset instances.
#[derive(Default)]
pub struct RasterDatasetBuilder<T>
where
    T: RasterType,
{
    /// Source file paths.
    pub sources: Vec<PathBuf>,
    /// Band indices per source.
    pub bands: HashMap<String, Vec<usize>>,
    /// Band indices as vector.
    pub bands_vec: Vec<Vec<usize>>,
    /// EPSG code for the coordinate reference system.
    pub epsg: u32,
    /// Overlap size between blocks.
    pub overlap_size: usize,
    /// Block size for reading.
    pub block_size: BlockSize,
    /// Dimension for band composition.
    pub composition_bands: Dimension,
    /// Dimension for source composition.
    pub composition_sources: Dimension,
    /// Image dimensions.
    pub image_size: ImageSize,
    /// Geotransform parameters.
    pub geo_transform: GeoTransform,
    /// Date indices for temporal data.
    pub date_indices: Vec<DateType>,
    /// Image resolution.
    pub resolution: ImageResolution,
    /// STAC feature collection.
    pub feature_collection: Option<ItemCollection>,
    /// No-data value.
    pub na_value: T,
    /// Layer names from DataSource.
    pub layer_names: Vec<String>,
}

impl<T> RasterDatasetBuilder<T>
where
    T: RasterType,
{
    /// Creates a RasterDataset from scratch (empty raster).
    ///
    /// Creates a new empty raster dataset with the specified extent, resolution, and CRS.
    /// Useful for creating output rasters for rasterization or other operations that
    /// need a pre-defined grid.
    ///
    /// # Arguments
    ///
    /// * `extent` - Geographic extent as [Extent](crate::metadata::Extent) struct
    /// * `resolution` - Pixel size (will be converted to f64)
    /// * `epsg_code` - EPSG coordinate reference system code (e.g., 32755 for UTM zone 55S)
    /// * `block_size` - Size of processing blocks as [BlockSize]
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use eorst::rasterdataset::RasterDatasetBuilder;
    /// use eorst::metadata::Extent;
    /// use eorst::types::BlockSize;
    ///
    /// let extent = Extent { xmin: 0., ymin: 0., xmax: 1000., ymax: 1000. };
    /// let block_size = BlockSize { cols: 256, rows: 256 };
    ///
    /// let rds: RasterDataset<i32> = RasterDatasetBuilder::from_scratch(
    ///     extent, 30.0, 32755, block_size
    /// );
    /// ```
    pub fn from_scratch<U>(
        extent: crate::metadata::Extent,
        resolution: U,
        epsg_code: u32,
        block_size: BlockSize,
    ) -> RasterDataset<T>
    where
        U: ToPrimitive + std::fmt::Debug,
    {
        debug!("Extent: {:?}", extent);
        let resolution = resolution.to_f64().expect("Unable to convert");
        debug!("Resolution {:?}", resolution);

        let layers = Vec::new();
        let rows = ((extent.ymax - extent.ymin) / resolution) as usize;
        let cols = ((extent.xmax - extent.xmin) / resolution) as usize;
        debug!("Creating raster with rows: {rows:?} and cols: {cols:?}");

        let shape = RasterDataShape {
            times: 1,
            layers: 1,
            rows,
            cols,
        };

        let geo_transform = GeoTransform {
            x_ul: extent.xmin,
            x_res: resolution,
            x_rot: 0.,
            y_ul: extent.ymin,
            y_rot: 0.,
            y_res: resolution,
        };

        let overlap_size = 0;
        let date_indices = Vec::new();
        let layer_indices = Vec::new();

        let metadata = RasterMetadata {
            layers: layers.clone(),
            shape,
            block_size,
            epsg_code,
            geo_transform,
            overlap_size,
            date_indices,
            layer_indices,
            na_value: T::zero(),
        };

        let tmp_layers = Vec::new();
        let image_size = ImageSize { rows, cols };

        let n_block_cols = n_block_cols(image_size, block_size);
        let n_block_rows = n_block_rows(image_size, block_size);
        let n_blocks = n_block_rows * n_block_cols;

        let mut blocks: Vec<crate::blocks::RasterBlock<T>> = Vec::new();
        (0..n_blocks).for_each(|i| {
            let block_attributes = block_from_id(
                i,
                shape,
                &layers,
                epsg_code as usize,
                image_size,
                block_size,
                overlap_size,
                geo_transform,
            );
            blocks.push(block_attributes);
        });

        let rds = RasterDataset {
            metadata,
            blocks,
            tmp_layers,
        };

        debug!("rds from scratch: {rds}");
        rds
    }

    /// Creates a builder from a single data source.
    ///
    /// Convenience method that wraps [`from_sources`](Self::from_sources).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use eorst::rasterdataset::RasterDatasetBuilder;
    /// use eorst::data_sources::DataSourceBuilder;
    ///
    /// let data_source = DataSourceBuilder::from_file(&path_to_file).build();
    /// let rds = RasterDatasetBuilder::from_source(&data_source).build();
    /// ```
    pub fn from_source(data_source: &DataSource) -> RasterDatasetBuilder<T> {
        RasterDatasetBuilder::from_sources(&vec![data_source.to_owned()])
    }

    /// Creates a builder from a STAC item collection.
    ///
    /// This is a convenience alias for [from_stac_query](Self::from_stac_query).
    pub fn from_item_collection(feature_collection: &ItemCollection) -> Self {
        RasterDatasetBuilder::from_stac_query(feature_collection)
    }

    /// Creates a RasterDataset from a STAC item collection.
    ///
    /// Downloads and mosaics imagery from a STAC ItemCollection. This is the primary way
    /// to load cloud-optimized satellite imagery from providers like DEA, Element84, or
    /// Planetary Computer.
    ///
    /// The method finds the minimum extent of all images, reprojects and resamples as
    /// needed, and creates a unified dataset. Tiles with different EPSG codes are
    /// automatically reprojected to the target CRS. Resolution can also be changed.
    ///
    /// Use builder methods to configure the output:
    /// - [`set_epsg`](Self::set_epsg) — Set the target coordinate reference system
    /// - [`set_resolution`](Self::set_resolution) — Change the output pixel resolution
    /// - [`block_size`](Self::block_size) — Set the processing block size
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use eorst::rasterdataset::RasterDatasetBuilder;
    /// use stac::ItemCollection;
    ///
    /// let items: ItemCollection = /* ... from a STAC API query ... */;
    /// let rds = RasterDatasetBuilder::from_stac_query(&items)
    ///     .set_epsg(32755)          // UTM zone 55S
    ///     .set_resolution(30.0)     // 30m resolution
    ///     .block_size(BlockSize { cols: 2048, rows: 2048 })
    ///     .build();
    /// ```
    pub fn from_stac_query(feature_collection: &ItemCollection) -> Self {
        let bands_vec = Vec::new();

        // default behaviour
        let composition_bands = Dimension::Time;
        let composition_sources = Dimension::Layer; // as many times as sources
        let geo_transform = GeoTransform::default();
        let image_size = ImageSize::default();
        let block_size = BlockSize::default();
        let overlap_size = 0;
        let first_item_assets = feature_collection.items[0].assets.values().next().unwrap();
        let href = Path::new(&first_item_assets.href);
        let epsg = Self::get_epsg_code(href);
        let resolution = Self::get_resolution(href);

        let sources: Vec<PathBuf> = Vec::new();
        let mut date_indices: Vec<DateType> = Vec::new();
        let mut bands = HashMap::new();
        for item in feature_collection.items.iter() {
            let date = item.properties.datetime.unwrap();
            date_indices.push(DateType::Date(date.into()));
            for (name, _) in item.assets.iter() {
                let title = name.to_owned();
                let band = 1; // FIX
                bands.insert(title, vec![band]);
            }
        }
        date_indices.sort();

        let date_indices = date_indices.into_iter().collect();

        RasterDatasetBuilder {
            sources,
            bands,
            bands_vec,
            epsg,
            overlap_size,
            block_size,
            composition_bands,
            composition_sources,
            image_size,
            geo_transform,
            date_indices,
            resolution,
            feature_collection: Some(feature_collection.clone()),
            na_value: NumCast::from(0).unwrap(),
            layer_names: Vec::new(),
        }
    }

    /// Creates a builder from multiple data sources.
    ///
    /// All sources must have the same shape and GeoTransform. They are stacked along
    /// the time or layer dimension depending on the composition settings.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use eorst::rasterdataset::RasterDatasetBuilder;
    /// use eorst::data_sources::DataSourceBuilder;
    ///
    /// let sources = vec![
    ///     DataSourceBuilder::from_file(&PathBuf::from("scene_2023.tif")).build(),
    ///     DataSourceBuilder::from_file(&PathBuf::from("scene_2024.tif")).build(),
    /// ];
    /// let rds = RasterDatasetBuilder::from_sources(&sources)
    ///     .block_size(BlockSize { cols: 2048, rows: 2048 })
    ///     .build();
    /// ```
    pub fn from_sources(data_sources: &Vec<DataSource>) -> RasterDatasetBuilder<T> {
        let mut sources: Vec<PathBuf> = Vec::new();
        let mut bands_vec = Vec::new();
        let bands = HashMap::new();

        let geo_transform = GeoTransform::default();
        let image_size = ImageSize::default();
        let block_size = BlockSize::default();

        let overlap_size = 0;

        // default behaviour for compositions
        let composition_bands = Dimension::Layer;
        let composition_sources = Dimension::Time; // as many times as sources

        for data_source in data_sources {
            sources.push(data_source.source.clone());
            bands_vec.push(data_source.bands.clone());
        }

        let date_indices = (0..sources.len()).map(DateType::Index).collect();

        // Single GDAL open to extract all metadata (replaces 3 separate opens)
        let info = read_basic_raster_info(&data_sources[0].source);
        let epsg = info.epsg_code;
        let resolution = info.resolution();
        let na_value = info.na_value::<T>();

        // Propagate layer_names from DataSource
        let layer_names = if !data_sources.is_empty() {
            data_sources[0].layer_names.clone()
        } else {
            Vec::new()
        };

        RasterDatasetBuilder {
            sources,
            bands,
            bands_vec,
            epsg,
            overlap_size,
            block_size,
            composition_bands,
            composition_sources,
            image_size,
            geo_transform,
            date_indices,
            resolution,
            feature_collection: None,
            na_value,
            layer_names,
        }
    }

    /// Sets the desired EPSG code for the raster dataset.
    ///
    /// If the source data has a different CRS, it will be automatically reprojected
    /// during `build()`.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// RasterDatasetBuilder::from_source(&data_source)
    ///     .set_epsg(32755)  // UTM zone 55S
    ///     .build()
    /// ```
    pub fn set_epsg(mut self, epsg_code: u32) -> Self {
        self.epsg = epsg_code;
        self
    }

    /// Sets the image size.
    pub fn set_image_size(mut self, image_size: ImageSize) -> Self {
        self.image_size = image_size;
        self
    }

    /// Set the desired geo_transform
    pub fn set_geo_transform(mut self, geo_transform: GeoTransform) -> Self {
        self.geo_transform = geo_transform;
        self
    }

    /// Sets the desired resolution for the raster dataset.
    ///
    /// If the source data has a different resolution, it will be resampled
    /// during `build()`.
    pub fn set_resolution(mut self, resolution: ImageResolution) -> Self {
        self.resolution = resolution;
        self
    }

    /// Sets the block size for chunked raster operations.
    ///
    /// Larger blocks reduce I/O overhead but use more memory. Default is 1024×1024.
    /// A good choice for Sentinel-2 data is 2048×2048.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// RasterDatasetBuilder::from_source(&data_source)
    ///     .block_size(BlockSize { cols: 2048, rows: 2048 })
    ///     .build()
    /// ```
    pub fn block_size(mut self, block_size: BlockSize) -> RasterDatasetBuilder<T> {
        self.block_size = block_size;
        self
    }

    /// Sets the overlap size for block-based reads.
    ///
    /// Overlap pixels are included on all sides of each block, useful for
    /// neighborhood operations that require data from adjacent blocks.
    pub fn overlap_size(mut self, overlap_size: usize) -> RasterDatasetBuilder<T> {
        self.overlap_size = overlap_size;
        self
    }

    /// Sets the source composition dimension.
    ///
    /// Controls how multiple data sources are arranged along the time/layer axes.
    pub fn source_composition_dimension(
        mut self,
        composition_dimension: Dimension,
    ) -> RasterDatasetBuilder<T> {
        self.composition_sources = composition_dimension;
        self
    }

    /// Sets the band composition dimension.
    ///
    /// Controls how bands from each source are arranged along the time/layer axes.
    pub fn band_composition_dimension(
        mut self,
        composition_dimension: Dimension,
    ) -> RasterDatasetBuilder<T> {
        self.composition_bands = composition_dimension;
        self
    }

    /// Sets the date indices for temporal data.
    pub fn set_date_indices(mut self, date_indices: &[DateType]) -> Self {
        self.date_indices = date_indices.to_vec();
        self
    }

    /// Sets the band mapping.
    ///
    /// Maps band names to their source band indices. Useful when working with
    /// multi-band files or semantic band names.
    pub fn bands(mut self, bands: HashMap<String, Vec<usize>>) -> Self {
        self.bands = bands;
        self
    }

    /// Sets the template for the raster dataset.
    ///
    /// Copies the image size and GeoTransform from another dataset, ensuring
    /// the output matches its geometry. The EPSG code is also copied.
    pub fn set_template<V>(mut self, other: &RasterDataset<V>) -> RasterDatasetBuilder<T>
    where
        V: RasterType,
    {
        self.image_size = ImageSize {
            rows: other.metadata.shape.rows,
            cols: other.metadata.shape.cols,
        };
        self.geo_transform = other.metadata.geo_transform;
        self.epsg = other.metadata.epsg_code;
        self
    }

    fn get_resolution(source: &Path) -> ImageResolution {
        read_basic_raster_info(source).resolution()
    }

    fn get_geotransform(source: &Path) -> GeoTransform {
        read_basic_raster_info(source).geo_transform_struct()
    }

    fn get_max_extent(extents: Vec<crate::metadata::Extent>) -> crate::metadata::Extent {
        let xmin = extents
            .iter()
            .map(|ex| ex.xmin)
            .min_by(|i, j| i.partial_cmp(j).unwrap())
            .unwrap();
        let ymin = extents
            .iter()
            .map(|ex| ex.ymin)
            .min_by(|i, j| i.partial_cmp(j).unwrap())
            .unwrap();
        let xmax = extents
            .iter()
            .map(|ex| ex.xmax)
            .max_by(|i, j| i.partial_cmp(j).unwrap())
            .unwrap();
        let ymax = extents
            .iter()
            .map(|ex| ex.ymax)
            .max_by(|i, j| i.partial_cmp(j).unwrap())
            .unwrap();
        crate::metadata::Extent {
            xmin,
            ymin,
            xmax,
            ymax,
        }
    }

    fn get_extent(source: &Path, target_epsg: u32) -> crate::metadata::Extent {
        let info = read_basic_raster_info(source);
        let gt = info.geo_transform_struct();
        let is = info.image_size();

        if info.epsg_code == target_epsg {
            let xmin = [gt.x_ul];
            let xmax = [xmin[0] + gt.x_res * is.cols as f64];
            let ymax = [gt.y_ul];
            let ymin = [ymax[0] + gt.y_res * is.rows as f64];
            crate::metadata::Extent {
                xmin: xmin[0],
                ymin: ymin[0],
                xmax: xmax[0],
                ymax: ymax[0],
            }
        } else {
            let new_source = create_temp_file("vrt");
            let epsg_s = format!("EPSG:{}", target_epsg);
            let source_s = source.to_string_lossy();
            let argv = vec![
                "gdalwarp", "-t_srs", &epsg_s, "-q", &source_s, &new_source,
            ];
            run_gdal_command(&argv);
            let new_info = read_basic_raster_info(&PathBuf::from(new_source));
            let gt = new_info.geo_transform_struct();
            let is = new_info.image_size();
            let xmin = [gt.x_ul];
            let ymax = [gt.y_ul];
            let ymin = [ymax[0] + gt.y_res * is.rows as f64];
            let xmax = [xmin[0] + gt.x_res * is.cols as f64];
            crate::metadata::Extent {
                xmin: xmin[0],
                ymin: ymin[0],
                xmax: xmax[0],
                ymax: ymax[0],
            }
        }
    }

    fn get_epsg_code(source: &Path) -> u32 {
        debug!("get_epsg_code: source {:?}", source);
        let info = read_basic_raster_info(source);
        info.epsg_code
    }

    fn get_blocks(
        &self,
        block_shape: RasterDataShape,
        layers: &[Layer],
        epsg_code: usize,
    ) -> Vec<crate::blocks::RasterBlock<T>>
    where
        T: RasterType,
    {
        let n_blocks = self.n_blocks();
        let mut blocks: Vec<crate::blocks::RasterBlock<T>> = Vec::new();
        (0..n_blocks).for_each(|i| {
            let block_attributes = self.block_from_id(i, block_shape, layers, epsg_code);
            blocks.push(block_attributes)
        });
        blocks
    }

    fn n_block_cols(&self) -> usize {
        let image_size = self.image_size;
        let block_size = self.block_size;
        round::ceil(image_size.cols as f64 / block_size.cols as f64, 0) as usize
    }

    fn n_block_rows(&self) -> usize {
        let image_size = self.image_size;
        let block_size = self.block_size;
        round::ceil(image_size.rows as f64 / block_size.rows as f64, 0) as usize
    }

    fn n_blocks(&self) -> usize {
        self.n_block_cols() * self.n_block_rows()
    }

    fn block_from_id(
        &self,
        id: usize,
        block_shape: RasterDataShape,
        _layers: &[Layer],
        epsg_code: usize,
    ) -> crate::blocks::RasterBlock<T> {
        let tile_col = self.block_col_row(id).0;
        let tile_row = self.block_col_row(id).1;
        let overlap = get_overlap(
            self.image_size,
            self.block_size,
            self.block_col_row(id),
            self.overlap_size,
        );

        let ul_x = (self.block_size.cols * tile_col) as isize;
        let ul_y = (self.block_size.rows * tile_row) as isize;

        let ul_x_overlap = ul_x - overlap.left as isize;
        let ul_y_overlap = ul_y - overlap.top as isize;
        let lr_x_overlap = std::cmp::min(
            self.image_size.cols as isize,
            ul_x + self.block_size.cols as isize + overlap.right as isize,
        );
        let lr_y_overlap = std::cmp::min(
            self.image_size.rows as isize,
            ul_y + self.block_size.rows as isize + overlap.bottom as isize,
        );

        let win_width = lr_x_overlap - ul_x_overlap;
        let win_height = lr_y_overlap - ul_y_overlap;

        let read_window = ReadWindow {
            offset: Offset {
                cols: ul_x_overlap,
                rows: ul_y_overlap,
            },
            size: Size {
                cols: win_width,
                rows: win_height,
            },
        };

        let block_geo_transform = get_block_gt(read_window, self.geo_transform);
        let empty_metadata: RasterMetadata<T> = RasterMetadata::new();
        crate::blocks::RasterBlock {
            block_index: id,
            read_window,
            overlap_size: self.overlap_size,
            geo_transform: block_geo_transform,
            overlap,
            shape: block_shape,
            epsg_code,
            raster_metadata: empty_metadata,
        }
    }

    fn block_col_row(&self, id: usize) -> (usize, usize) {
        let block_row = id / self.n_block_cols();
        let block_col = id - (block_row * self.n_block_cols());
        (block_col, block_row)
    }


    /// Builds the RasterDataset.
    pub fn build(mut self) -> RasterDataset<T> {
        debug!("Building RasterDataset");

        let raster_dataset = match self.feature_collection {
            None => {
                let mut layers = Vec::new();
                let mut blocks: Vec<crate::blocks::RasterBlock<T>> = Vec::new();
                let mut tmp_layers: Vec<PathBuf> = Vec::new();
                // Pre-read metadata for all sources (single open per source instead of 2)
                let t_meta = std::time::Instant::now();
                let source_meta = self.sources.iter()
                    .map(|s| read_basic_raster_info(s))
                    .collect::<Vec<BasicRasterInfo>>();
                let t_meta_elapsed = t_meta.elapsed().as_secs_f64() * 1000.0;

                for (i, source) in self.sources.iter_mut().enumerate() {
                    let target_epsg = self.epsg;
                    let target_resolution = self.resolution;

                    let current_epsg = source_meta[i].epsg_code;
                    let current_resolution = source_meta[i].resolution();

                    if current_epsg != target_epsg {
                        *source = warp(source.to_path_buf(), None, target_epsg);
                        tmp_layers.push(source.clone());
                    }

                    if current_resolution != target_resolution {
                        *source = change_res(source.to_path_buf(), target_resolution);
                        tmp_layers.push(source.clone());
                    }
                }

                let t_loop2 = std::time::Instant::now();
                for source in self.sources.iter_mut() {
                    let current_gt = Self::get_geotransform(source);
                    let target_gt = self.geo_transform;

                    if (current_gt != target_gt) & (target_gt != GeoTransform::default()) {
                        let current_extent =
                            Self::get_extent(source, self.epsg);
                        let x_ll = target_gt.x_ul + (self.image_size.cols as f64 * target_gt.x_res);
                        let y_ll = target_gt.y_ul + (self.image_size.rows as f64 * target_gt.y_res);

                        let target_extent = crate::metadata::Extent {
                            xmin: target_gt.x_ul,
                            ymin: y_ll,
                            xmax: x_ll,
                            ymax: target_gt.y_ul,
                        };
                        let te = if current_extent != target_extent {
                            crate::metadata::Extent {
                                xmin: target_gt.x_ul,
                                ymin: y_ll,
                                xmax: x_ll,
                                ymax: target_gt.y_ul,
                            }
                        } else {
                            current_extent
                        };

                        let tr = ImageResolution {
                            x: target_gt.x_res,
                            y: target_gt.y_res,
                        };

                        *source = warp_with_te_tr(source.to_path_buf(), &te, tr);
                        tmp_layers.push(source.clone());
                    }
                }
                let t_loop2_elapsed = t_loop2.elapsed().as_secs_f64() * 1000.0;

                let t_info = std::time::Instant::now();
                let info = read_basic_raster_info(&self.sources[0]);
                let t_info_elapsed = t_info.elapsed().as_secs_f64() * 1000.0;
                let geo_transform = info.geo_transform_struct();
                debug!("GeoTransform template from {:?}", &self.sources[0]);
                let image_size = info.image_size();

                self.geo_transform = geo_transform;
                self.image_size = image_size;

                let n_rows = image_size.rows;
                let n_cols = image_size.cols;
                // Validate that all sources match the template geometry.
                // source[0] is already confirmed by the template read above.
                for source in self.sources.iter().skip(1) {
                    debug!("checking {:?}", source);
                    let src_info = read_basic_raster_info(source);
                    let gt = src_info.geo_transform_struct();
                    let is = src_info.image_size();
                    if gt != geo_transform {
                        panic!("Sources have different geo_transforms!");
                    }

                    if is != image_size {
                        panic!("Sources have different size!")
                    }
                }
                let mut time_pos = 1;
                let mut layer_pos = 1;

                let t_bands = std::time::Instant::now();
                for (source_idx, source) in self.sources.iter().enumerate() {
                    let bands = self.bands_vec[source_idx].clone();

                    for band in bands.iter() {
                        let source = PathBuf::from(source);
                        let new_source = extract_band(&source, *band);
                        tmp_layers.push(new_source.clone());
                        let source = new_source;
                        let layer = Layer {
                            source,
                            time_pos: time_pos - 1,
                            layer_pos: layer_pos - 1,
                        };
                        layers.push(layer);

                        match self.composition_bands {
                            Dimension::Layer => layer_pos += 1,
                            Dimension::Time => time_pos += 1,
                        }
                    }

                    match self.composition_sources {
                        Dimension::Layer => {
                            if self.composition_bands != Dimension::Layer {
                                layer_pos += 1
                            }
                        }
                        Dimension::Time => {
                            if self.composition_bands != Dimension::Time {
                                time_pos += 1
                            }
                        }
                    }

                    match self.composition_sources {
                        Dimension::Layer => time_pos = 1,
                        Dimension::Time => layer_pos = 1,
                    }
                }

                let t_bands_elapsed = t_bands.elapsed().as_secs_f64() * 1000.0;

                let n_times = layers.iter().map(|l| l.time_pos).max().unwrap() + 1;
                let n_layers = layers.iter().map(|l| l.layer_pos).max().unwrap() + 1;

                // Use layer_names from DataSource if provided, otherwise generate defaults
                let layer_names = if !self.layer_names.is_empty() {
                    self.layer_names.clone()
                } else {
                    (0..n_layers)
                        .map(|i| format!("Layer_{}", i))
                        .collect()
                };

                let block_shape = RasterDataShape {
                    times: n_times,
                    layers: n_layers,
                    rows: n_rows,
                    cols: n_cols,
                };

                let t_blocks = std::time::Instant::now();
                blocks.extend(self.get_blocks(block_shape, &layers, self.epsg.try_into().unwrap()));
                let t_blocks_elapsed = t_blocks.elapsed().as_secs_f64() * 1000.0;

                debug!("build() timing: meta_read={:.1}ms loop2={:.1}ms template={:.1}ms bands={:.1}ms blocks={:.1}ms",
                    t_meta_elapsed, t_loop2_elapsed, t_info_elapsed, t_bands_elapsed, t_blocks_elapsed);

                let metadata = RasterMetadata {
                    layers,
                    shape: block_shape,
                    block_size: self.block_size,
                    epsg_code: self.epsg,
                    geo_transform,
                    overlap_size: self.overlap_size,
                    date_indices: self.date_indices,
                    layer_indices: layer_names,
                    na_value: self.na_value,
                };

                RasterDataset {
                    metadata,
                    blocks,
                    tmp_layers,
                }
            }
            Some(ref feature_collection) => {
                let mut sources = Vec::new();
                for item in &feature_collection.items {
                    for asset in item.assets.values() {
                        debug!("Asset {:?}", asset);

                        let href = &asset.href;
                        sources.push(PathBuf::from(href));
                    }
                }
                debug!("Assets added");

                let tmp_layers = Arc::new(Mutex::new(Vec::new()));
                let layers = Arc::new(Mutex::new(Vec::new()));

                let mut blocks: Vec<crate::blocks::RasterBlock<T>> = Vec::new();

                let extents: Vec<_> = sources
                    .par_iter()
                    .map(|source| Self::get_extent(source, self.epsg))
                    .collect();

                let global_extent = Self::get_max_extent(extents);

                let original_times_indices = crate::stac_helpers::get_sorted_datetimes(feature_collection);

                let target_time_indices = crate::stac_helpers::unique_datetimes_in_range(original_times_indices);

                let asset_names = crate::stac_helpers::get_asset_names(feature_collection);
                debug!("Asset names: {:?} ", asset_names);

                target_time_indices
                    .par_iter()
                    .enumerate()
                    .for_each(|(time_idx, time)| {
                        let mut layer_idx = 0;

                        let date_items = crate::stac_helpers::get_items_for_date(feature_collection, time);

                        for asset_name in asset_names.iter() {
                            let asset_bands = [1];

                            for _ in asset_bands.iter() {
                                let sources = crate::stac_helpers::get_sources_for_asset(&date_items, asset_name);
                                let date_mosaic_tmp = PathBuf::from(create_temp_file("vrt"));
                                mosaic(&sources, &date_mosaic_tmp, self.epsg, None, None).unwrap();
                                let mut single_band = date_mosaic_tmp;
                                let raster_info = read_basic_raster_info(&single_band);
                                let current_epsg = raster_info.epsg_code;

                                let target_epsg = self.epsg;
                                if current_epsg != target_epsg {
                                    single_band = warp(single_band, None, target_epsg);
                                }

                                let current_resoultion = read_basic_raster_info(&single_band).resolution();
                                let target_resolution = self.resolution;
                                let tr = if current_resoultion == target_resolution {
                                    current_resoultion
                                } else {
                                    target_resolution
                                };

                                let current_gt = read_basic_raster_info(&single_band).geo_transform_struct();
                                let target_gt = self.geo_transform;
                                let mut local_extent = global_extent.clone();
                                let mut local_tr = tr;

                                if (current_gt != target_gt)
                                    & (target_gt != GeoTransform::default())
                                {
                                    let x_ll = target_gt.x_ul
                                        + (self.image_size.cols as f64 * target_gt.x_res);
                                    let y_ll = target_gt.y_ul
                                        + (self.image_size.rows as f64 * target_gt.y_res);
                                    local_extent = crate::metadata::Extent {
                                        xmin: (target_gt.x_ul * 100.).round() / 100.,
                                        ymin: (y_ll * 100.).round() / 100.,
                                        xmax: (x_ll * 100.).round() / 100.,
                                        ymax: (target_gt.y_ul * 100.).round() / 100.,
                                    };
                                    local_tr = ImageResolution {
                                        x: (target_gt.x_res * 100.).round() / 100.,
                                        y: (target_gt.y_res * 100.).round() / 100.,
                                    };
                                }

                                single_band = warp_with_te_tr(single_band, &local_extent, local_tr);

                                let layer = Layer {
                                    source: single_band.clone(),
                                    time_pos: time_idx,
                                    layer_pos: layer_idx,
                                };

                                tmp_layers.lock().unwrap().push(single_band);
                                layers.lock().unwrap().push(layer);
                                layer_idx += 1;
                            }
                        }
                    });
                let tmp_layers = Arc::try_unwrap(tmp_layers).unwrap().into_inner().unwrap();
                let layers = Arc::try_unwrap(layers).unwrap().into_inner().unwrap();

                // update image size and geotransform fields.
                let info = read_basic_raster_info(&layers[0].source);
                self.image_size = info.image_size();
                self.geo_transform = info.geo_transform_struct();

                // Start computing the metadata
                let n_times = target_time_indices.len();

                let mut n_layers = 0;
                for (_, v) in self.bands.clone() {
                    n_layers += v.len();
                }
                let block_shape = RasterDataShape {
                    times: n_times,
                    layers: n_layers,
                    rows: self.image_size.rows,
                    cols: self.image_size.cols,
                };
                blocks.extend(self.get_blocks(block_shape, &layers, self.epsg.try_into().unwrap()));
                let target_time_indices: Vec<DateType> = target_time_indices
                    .into_iter()
                    .map(DateType::Date)
                    .collect();

                let metadata = RasterMetadata {
                    layers,
                    shape: block_shape,
                    block_size: self.block_size,
                    epsg_code: self.epsg,
                    geo_transform: self.geo_transform,
                    overlap_size: self.overlap_size,
                    date_indices: target_time_indices,
                    layer_indices: asset_names,
                    na_value: T::zero(),
                };

                RasterDataset {
                    metadata,
                    blocks,
                    tmp_layers,
                }
            }
        };
        raster_dataset
    }
}

// ─── Shared helper functions (used by both method and free-function paths) ───

/// Computes overlap for a block based on its position in the grid.
pub(crate) fn get_overlap(
    image_size: ImageSize,
    block_size: BlockSize,
    block_col_row: (usize, usize),
    overlap_size: usize,
) -> Overlap {
    let mut overlap = Overlap {
        left: overlap_size,
        top: overlap_size,
        right: overlap_size,
        bottom: overlap_size,
    };
    let tile_col = block_col_row.0;
    let tile_row = block_col_row.1;
    let n_rows_tile = image_size.rows.div_ceil(block_size.rows);
    let n_cols_tile = image_size.cols.div_ceil(block_size.cols);
    let is_top = tile_row == 0;
    let is_bottom = tile_row + 1 == n_rows_tile;
    let is_left = tile_col == 0;
    let is_right = tile_col + 1 == n_cols_tile;
    if is_top {
        overlap.top = 0
    };
    if is_bottom {
        overlap.bottom = 0
    };
    if is_left {
        overlap.left = 0
    };
    if is_right {
        overlap.right = 0
    };
    overlap
}

/// Computes the number of block columns for the given image and block sizes.
pub(crate) fn n_block_cols(image_size: ImageSize, block_size: BlockSize) -> usize {
    round::ceil(image_size.cols as f64 / block_size.cols as f64, 0) as usize
}

/// Computes the number of block rows for the given image and block sizes.
pub(crate) fn n_block_rows(image_size: ImageSize, block_size: BlockSize) -> usize {
    round::ceil(image_size.rows as f64 / block_size.rows as f64, 0) as usize
}

/// Converts a block ID to (col, row) position in the block grid.
fn block_col_row(id: usize, image_size: ImageSize, block_size: BlockSize) -> (usize, usize) {
    let ncols = n_block_cols(image_size, block_size);
    let block_row = id / ncols;
    let block_col = id - (block_row * ncols);
    (block_col, block_row)
}

/// Computes the block-level GeoTransform from a read window and image GeoTransform.
fn get_block_gt(read_window: ReadWindow, geo_transform: GeoTransform) -> GeoTransform {
    let x_ul_image = geo_transform.x_ul;
    let y_ul_image = geo_transform.y_ul;

    let x_res = geo_transform.x_res;
    let y_res = geo_transform.y_res;
    let x_pos = read_window.offset.cols;
    let y_pos = read_window.offset.rows;

    let x_ul_block = x_ul_image + x_res * x_pos as f64;
    let y_ul_block = y_ul_image + y_res * y_pos as f64;
    GeoTransform {
        x_ul: x_ul_block,
        x_res,
        x_rot: geo_transform.x_rot,
        y_ul: y_ul_block,
        y_rot: geo_transform.x_rot,
        y_res,
    }
}

/// Creates a RasterBlock from an ID and parameters.
/// Used by both `from_scratch()` and the builder's `block_from_id()` method.
fn block_from_id<U>(
    id: usize,
    block_shape: RasterDataShape,
    _layers: &[Layer],
    epsg_code: usize,
    image_size: ImageSize,
    block_size: BlockSize,
    overlap_size: usize,
    geo_transform: GeoTransform,
) -> crate::blocks::RasterBlock<U>
where
    U: RasterType,
{
    let tile_col = block_col_row(id, image_size, block_size).0;
    let tile_row = block_col_row(id, image_size, block_size).1;
    let overlap = get_overlap(
        image_size,
        block_size,
        block_col_row(id, image_size, block_size),
        overlap_size,
    );

    let ul_x = (block_size.cols * tile_col) as isize;
    let ul_y = (block_size.rows * tile_row) as isize;

    // now compute with overlap
    let ul_x_overlap = ul_x - overlap.left as isize;
    let ul_y_overlap = ul_y - overlap.top as isize;
    let lr_x_overlap = std::cmp::min(
        image_size.cols as isize,
        ul_x + block_size.cols as isize + overlap.right as isize,
    );

    let lr_y_overlap = std::cmp::min(
        image_size.rows as isize,
        ul_y + block_size.rows as isize + overlap.bottom as isize,
    );

    let win_width = lr_x_overlap - ul_x_overlap;
    let win_height = lr_y_overlap - ul_y_overlap;

    let arr_width = win_width;
    let arr_height = win_height;

    let read_window = ReadWindow {
        offset: Offset {
            cols: ul_x_overlap,
            rows: ul_y_overlap,
        },
        size: Size {
            cols: arr_width,
            rows: arr_height,
        },
    };

    let block_geo_transform = get_block_gt(read_window, geo_transform);
    let empty_metadata: RasterMetadata<U> = RasterMetadata::new();
    crate::blocks::RasterBlock {
        block_index: id,
        read_window,
        overlap_size,
        geo_transform: block_geo_transform,
        overlap,
        shape: block_shape,
        epsg_code,
        raster_metadata: empty_metadata,
    }
}