cogrs 0.0.4

use crate::raster::RasterSource;
use crate::tiff_utils::{
    AnyResult, TAG_GDAL_METADATA, parse_gdal_metadata_stats, read_ifd, read_tag_string_from_ifd,
    read_tiff_header,
};
use crate::tile_cache;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tiff::decoder::{Decoder, DecodingResult, Limits};
use tiff::tags::Tag;

pub struct TiffChunkedRasterSource {
    path: PathBuf,
    image_width: usize,
    image_length: usize,
    chunk_width: usize,
    chunk_height: usize,
    chunks_across: usize,
    chunks_down: usize,
    samples_per_pixel: usize,
    pixel_scale: Option<[f64; 3]>,
    tiepoint: Option<[f64; 6]>,
    min_max: Mutex<Option<(f32, f32)>>,
}

impl TiffChunkedRasterSource {
    /// # Errors
    /// Returns an error if the file cannot be opened or is not a valid tiled TIFF.
    pub fn open(path: &PathBuf) -> AnyResult<Self> {
        let mut decoder = Decoder::new(File::open(path)?)?;
        decoder = decoder.with_limits(Limits::unlimited());

        let (width_u32, height_u32) = decoder.dimensions()?;
        let width = width_u32 as usize;
        let height = height_u32 as usize;

        let color_type = decoder.colortype()?;
        let samples_per_pixel = match color_type {
            tiff::ColorType::Gray(_) => 1,
            _ => {
                return Err(
                    "Unsupported color type for chunked reader (only grayscale supported)".into(),
                );
            }
        };

        let (chunk_width_u32, chunk_height_u32) = decoder.chunk_dimensions();
        let chunk_width = chunk_width_u32.max(1) as usize;
        let chunk_height = chunk_height_u32.max(1) as usize;

        let chunks_across = width.div_ceil(chunk_width);
        let chunks_down = height.div_ceil(chunk_height);

        let pixel_scale = decoder
            .get_tag_f64_vec(Tag::ModelPixelScaleTag)
            .ok()
            .and_then(|v| {
                if v.len() >= 3 {
                    Some([v[0], v[1], v[2]])
                } else {
                    None
                }
            });

        let tiepoint = decoder
            .get_tag_f64_vec(Tag::ModelTiepointTag)
            .ok()
            .and_then(|v| {
                if v.len() >= 6 {
                    Some([v[0], v[1], v[2], v[3], v[4], v[5]])
                } else {
                    None
                }
            });

        let gdal_metadata = decoder
            .get_tag_ascii_string(Tag::from_u16_exhaustive(TAG_GDAL_METADATA))
            .ok();

        drop(decoder);

        let mut min_max_hint = gdal_metadata
            .as_ref()
            .and_then(|s| parse_gdal_metadata_stats(s));

        if min_max_hint.is_none() {
            // fallback to reading via low-level utils to attempt metadata parse
            if let Ok(mut file) = File::open(path)
                && let Ok(header) = read_tiff_header(&mut file)
                    && let Ok(ifd_entries) = read_ifd(&mut file, header.little_endian)
                        && let Ok(meta_string) = read_tag_string_from_ifd(
                            &mut file,
                            &ifd_entries,
                            header.little_endian,
                            TAG_GDAL_METADATA,
                        ) {
                            min_max_hint = parse_gdal_metadata_stats(&meta_string);
                        }
        }

        Ok(Self {
            path: path.clone(),
            image_width: width,
            image_length: height,
            chunk_width,
            chunk_height,
            chunks_across,
            chunks_down,
            samples_per_pixel,
            pixel_scale,
            tiepoint,
            min_max: Mutex::new(min_max_hint),
        })
    }

    fn chunk_dimensions_for_index(&self, chunk_index: usize) -> (usize, usize) {
        let chunk_col = chunk_index % self.chunks_across;
        let chunk_row = chunk_index / self.chunks_across;

        let width = if chunk_col == self.chunks_across - 1 {
            self.image_width - chunk_col * self.chunk_width
        } else {
            self.chunk_width
        };

        let height = if chunk_row == self.chunks_down - 1 {
            self.image_length - chunk_row * self.chunk_height
        } else {
            self.chunk_height
        };

        (width, height)
    }

    fn neighbor_indices(&self, chunk_index: usize) -> Vec<usize> {
        let chunk_col = chunk_index % self.chunks_across;
        let chunk_row = chunk_index / self.chunks_across;
        let mut neighbors = Vec::with_capacity(4);

        if chunk_col > 0 {
            neighbors.push(chunk_index - 1);
        }
        if chunk_col + 1 < self.chunks_across {
            neighbors.push(chunk_index + 1);
        }
        if chunk_row > 0 {
            neighbors.push(chunk_index - self.chunks_across);
        }
        if chunk_row + 1 < self.chunks_down {
            neighbors.push(chunk_index + self.chunks_across);
        }

        neighbors
    }

    fn prefetch_neighbors(&self, chunk_index: usize) {
        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            let neighbors = self.neighbor_indices(chunk_index);
            let config = ChunkPrefetchConfig {
                path: self.path.clone(),
                chunk_width: self.chunk_width,
                chunk_height: self.chunk_height,
                samples_per_pixel: self.samples_per_pixel,
            };

            for neighbor in neighbors {
                if tile_cache::contains_by_path(&config.path, neighbor) {
                    continue;
                }

                let cfg = config.clone();
                handle.spawn_blocking(move || {
                    if let Ok(tile) = load_chunk_with_config(
                        cfg.path.as_path(),
                        neighbor,
                        cfg.chunk_width,
                        cfg.chunk_height,
                        cfg.samples_per_pixel,
                    ) {
                        tile_cache::insert_by_path(&cfg.path, neighbor, tile);
                    }
                });
            }
        }
    }

    fn load_chunk(&self, chunk_index: usize) -> AnyResult<Arc<Vec<f32>>> {
        load_chunk_with_config(
            self.path.as_path(),
            chunk_index,
            self.chunk_width,
            self.chunk_height,
            self.samples_per_pixel,
        )
    }

    fn fetch_chunk(&self, chunk_index: usize) -> AnyResult<Arc<Vec<f32>>> {
        if let Some(entry) = tile_cache::get_by_path(&self.path, chunk_index) {
            return Ok(entry);
        }

        let entry = self.load_chunk(chunk_index)?;
        tile_cache::insert_by_path(&self.path, chunk_index, Arc::clone(&entry));
        self.prefetch_neighbors(chunk_index);
        Ok(entry)
    }

    pub fn pixel_scale(&self) -> Option<[f64; 3]> {
        self.pixel_scale
    }

    pub fn tiepoint(&self) -> Option<[f64; 6]> {
        self.tiepoint
    }

    /// # Errors
    /// Returns an error if chunk reading fails.
    ///
    /// # Panics
    /// Panics if the mutex lock is poisoned.
    pub fn compute_min_max(&self) -> AnyResult<(f32, f32)> {
        if let Some(result) = *self.min_max.lock().unwrap() {
            return Ok(result);
        }

        let mut min_value = f32::INFINITY;
        let mut max_value = f32::NEG_INFINITY;

        for chunk_index in 0..(self.chunks_across * self.chunks_down) {
            let chunk = self.fetch_chunk(chunk_index)?;
            for &value in chunk.iter() {
                if value.is_nan() {
                    continue;
                }
                if value < min_value {
                    min_value = value;
                }
                if value > max_value {
                    max_value = value;
                }
            }
        }

        let result = if min_value.is_infinite() || max_value.is_infinite() {
            (0.0, 0.0)
        } else {
            (min_value, max_value)
        };

        *self.min_max.lock().unwrap() = Some(result);
        Ok(result)
    }
}

impl RasterSource for TiffChunkedRasterSource {
    fn bands(&self) -> usize {
        self.samples_per_pixel
    }

    fn width(&self) -> usize {
        self.image_width
    }

    fn height(&self) -> usize {
        self.image_length
    }

    fn sample(&self, band: usize, x: usize, y: usize) -> Option<f32> {
        if band >= self.samples_per_pixel || x >= self.image_width || y >= self.image_length {
            return None;
        }

        let chunk_col = x / self.chunk_width;
        let chunk_row = y / self.chunk_height;
        let chunk_index = chunk_row * self.chunks_across + chunk_col;

        let within_x = x % self.chunk_width;
        let within_y = y % self.chunk_height;

        let (actual_width, actual_height) = self.chunk_dimensions_for_index(chunk_index);
        if within_x >= actual_width || within_y >= actual_height {
            return None;
        }

        let Ok(chunk) = self.fetch_chunk(chunk_index) else {
            return None;
        };

        let offset = (within_y * self.chunk_width + within_x) * self.samples_per_pixel + band;
        chunk.get(offset).copied()
    }
}

fn convert_decoding_result(result: DecodingResult) -> Vec<f32> {
    match result {
        DecodingResult::U8(data) => data.into_iter().map(f32::from).collect(),
        DecodingResult::U16(data) => data.into_iter().map(f32::from).collect(),
        // Allow precision loss: converting large integers to f32 for visualization purposes
        #[allow(clippy::cast_precision_loss)]
        DecodingResult::U32(data) => data.into_iter().map(|v| v as f32).collect(),
        DecodingResult::I8(data) => data.into_iter().map(f32::from).collect(),
        DecodingResult::I16(data) => data.into_iter().map(f32::from).collect(),
        #[allow(clippy::cast_precision_loss)]
        DecodingResult::I32(data) => data.into_iter().map(|v| v as f32).collect(),
        DecodingResult::F32(data) => data,
        // Allow truncation: converting f64 to f32 for performance reasons
        #[allow(clippy::cast_possible_truncation)]
        DecodingResult::F64(data) => data.into_iter().map(|v| v as f32).collect(),
        DecodingResult::F16(data) => data.into_iter().map(f32::from).collect(),
        #[allow(clippy::cast_precision_loss)]
        DecodingResult::U64(data) => data.into_iter().map(|v| v as f32).collect(),
        #[allow(clippy::cast_precision_loss)]
        DecodingResult::I64(data) => data.into_iter().map(|v| v as f32).collect(),
    }
}

#[derive(Clone)]
struct ChunkPrefetchConfig {
    path: PathBuf,
    chunk_width: usize,
    chunk_height: usize,
    samples_per_pixel: usize,
}

fn load_chunk_with_config(
    path: &Path,
    chunk_index: usize,
    chunk_width: usize,
    chunk_height: usize,
    samples_per_pixel: usize,
) -> AnyResult<Arc<Vec<f32>>> {
    let file_decoder = Decoder::new(File::open(path)?)?;
    let mut decoder = file_decoder.with_limits(Limits::unlimited());
    // Cast is safe: chunk_index is derived from tile grid dimensions
    #[allow(clippy::cast_possible_truncation)]
    let data_dims = decoder.chunk_data_dimensions(chunk_index as u32);
    let actual_width = data_dims.0 as usize;
    let actual_height = data_dims.1 as usize;

    #[allow(clippy::cast_possible_truncation)]
    let chunk_data = decoder.read_chunk(chunk_index as u32)?;
    let values = convert_decoding_result(chunk_data);
    let expected_len = actual_width * actual_height * samples_per_pixel;
    if values.len() != expected_len {
        return Err(format!(
            "Decoded chunk has unexpected length {} (expected {})",
            values.len(),
            expected_len
        )
        .into());
    }

    let mut padded = vec![f32::NAN; chunk_width * chunk_height * samples_per_pixel];

    for row in 0..actual_height {
        for col in 0..actual_width {
            let base_src = (row * actual_width + col) * samples_per_pixel;
            let base_dst = (row * chunk_width + col) * samples_per_pixel;
            padded[base_dst..base_dst + samples_per_pixel]
                .copy_from_slice(&values[base_src..base_src + samples_per_pixel]);
        }
    }

    Ok(Arc::new(padded))
}