//! Image loading and saving utilities
//!
//! Handles loading various image formats and converting between
//! image crate types and ndarray for processing.

use anyhow::{Context, Result, bail};
use image::{DynamicImage, GenericImageView, ImageBuffer, Rgb, RgbImage};
use ndarray::Array3;
use slidescope_shared::ImageFormat;
use std::path::Path;
use tracing::info;

/// Load an image from a file path
pub fn load_image(path: &Path) -> Result<DynamicImage> {
    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");

    let format = ImageFormat::from_extension(ext)
        .with_context(|| format!("Unsupported image format: {}", ext))?;

    if format.requires_openslide() {
        bail!(
            "Format {} requires OpenSlide support (enable 'openslide' feature)",
            ext
        );
    }

    let img =
        image::open(path).with_context(|| format!("Failed to open image: {}", path.display()))?;

    let (width, height) = img.dimensions();
    info!("Loaded image: {}x{} from {}", width, height, path.display());

    Ok(img)
}

/// Convert DynamicImage to ndarray Array3<u8>
pub fn image_to_array(img: &DynamicImage) -> Array3<u8> {
    let rgb = img.to_rgb8();
    let (width, height) = rgb.dimensions();

    let mut array = Array3::zeros((height as usize, width as usize, 3));

    for y in 0..height {
        for x in 0..width {
            let pixel = rgb.get_pixel(x, y);
            array[[y as usize, x as usize, 0]] = pixel[0];
            array[[y as usize, x as usize, 1]] = pixel[1];
            array[[y as usize, x as usize, 2]] = pixel[2];
        }
    }

    array
}

/// Convert ndarray Array3<u8> to RgbImage
pub fn array_to_image(array: &Array3<u8>) -> Result<RgbImage> {
    let shape = array.raw_dim();
    if shape[2] != 3 {
        bail!("Expected 3 channels, got {}", shape[2]);
    }

    let height = shape[0] as u32;
    let width = shape[1] as u32;

    let mut img = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let r = array[[y as usize, x as usize, 0]];
            let g = array[[y as usize, x as usize, 1]];
            let b = array[[y as usize, x as usize, 2]];
            img.put_pixel(x, y, Rgb([r, g, b]));
        }
    }

    Ok(img)
}

/// Convert f32 array (0.0-1.0) to u8 image
pub fn array_f32_to_image(array: &Array3<f32>) -> Result<RgbImage> {
    let shape = array.raw_dim();
    if shape[2] != 3 {
        bail!("Expected 3 channels, got {}", shape[2]);
    }

    let height = shape[0] as u32;
    let width = shape[1] as u32;

    let mut img = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let r = (array[[y as usize, x as usize, 0]].clamp(0.0, 1.0) * 255.0) as u8;
            let g = (array[[y as usize, x as usize, 1]].clamp(0.0, 1.0) * 255.0) as u8;
            let b = (array[[y as usize, x as usize, 2]].clamp(0.0, 1.0) * 255.0) as u8;
            img.put_pixel(x, y, Rgb([r, g, b]));
        }
    }

    Ok(img)
}

/// Save an image to a file
pub fn save_image(img: &RgbImage, path: &Path) -> Result<()> {
    img.save(path)
        .with_context(|| format!("Failed to save image to {}", path.display()))?;

    info!("Saved image to {}", path.display());
    Ok(())
}

/// Generate a thumbnail from an image
pub fn generate_thumbnail(img: &DynamicImage, max_size: u32) -> DynamicImage {
    img.thumbnail(max_size, max_size)
}

/// Extract image metadata
#[derive(Debug, Clone)]
pub struct ImageMetadata {
    pub width: u32,
    pub height: u32,
    pub format: ImageFormat,
    pub color_type: String,
    pub file_size_bytes: u64,
}

/// Get metadata from an image file
pub fn get_metadata(path: &Path) -> Result<ImageMetadata> {
    let file_size = std::fs::metadata(path)
        .with_context(|| format!("Cannot read file metadata: {}", path.display()))?
        .len();

    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");

    let format = ImageFormat::from_extension(ext).unwrap_or(ImageFormat::Jpeg);

    // Load just enough to get dimensions
    let img = image::open(path)?;
    let (width, height) = img.dimensions();

    let color_type = format!("{:?}", img.color());

    Ok(ImageMetadata {
        width,
        height,
        format,
        color_type,
        file_size_bytes: file_size,
    })
}

/// Create a grayscale concentration image from a single channel
///
/// Useful for visualizing individual stain concentrations.
pub fn concentration_to_grayscale(concentration: &ndarray::Array2<f32>) -> Result<RgbImage> {
    let shape = concentration.raw_dim();
    let height = shape[0] as u32;
    let width = shape[1] as u32;

    // Find min/max for normalization
    let min_val = concentration.iter().fold(f32::INFINITY, |a, &b| a.min(b));
    let max_val = concentration
        .iter()
        .fold(f32::NEG_INFINITY, |a, &b| a.max(b));
    let range = (max_val - min_val).max(1e-6);

    let mut img = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let val = concentration[[y as usize, x as usize]];
            let normalized = ((val - min_val) / range).clamp(0.0, 1.0);
            let gray = (normalized * 255.0) as u8;
            img.put_pixel(x, y, Rgb([gray, gray, gray]));
        }
    }

    Ok(img)
}

/// Create a colored concentration image using a stain color
pub fn concentration_to_colored(
    concentration: &ndarray::Array2<f32>,
    stain_color: [f32; 3],
) -> Result<RgbImage> {
    let shape = concentration.raw_dim();
    let height = shape[0] as u32;
    let width = shape[1] as u32;

    // Find max for normalization
    let max_val = concentration.iter().fold(0.0f32, |a, &b| a.max(b));
    let scale = if max_val > 0.0 { 1.0 / max_val } else { 1.0 };

    let mut img = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let val = concentration[[y as usize, x as usize]] * scale;
            let r = (stain_color[0] * val * 255.0).clamp(0.0, 255.0) as u8;
            let g = (stain_color[1] * val * 255.0).clamp(0.0, 255.0) as u8;
            let b = (stain_color[2] * val * 255.0).clamp(0.0, 255.0) as u8;
            img.put_pixel(x, y, Rgb([r, g, b]));
        }
    }

    Ok(img)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_array_roundtrip() {
        // Create a test image
        let mut img = RgbImage::new(10, 10);
        for y in 0..10 {
            for x in 0..10 {
                img.put_pixel(x, y, Rgb([(x * 25) as u8, (y * 25) as u8, 128]));
            }
        }

        let dynamic = DynamicImage::ImageRgb8(img.clone());
        let array = image_to_array(&dynamic);

        assert_eq!(array.shape(), &[10, 10, 3]);

        let back = array_to_image(&array).unwrap();
        assert_eq!(back.dimensions(), (10, 10));

        // Check values preserved
        for y in 0..10 {
            for x in 0..10 {
                assert_eq!(img.get_pixel(x, y), back.get_pixel(x, y));
            }
        }
    }

    #[test]
    fn test_thumbnail() {
        let img = DynamicImage::ImageRgb8(RgbImage::new(1000, 800));
        let thumb = generate_thumbnail(&img, 256);

        let (w, h) = thumb.dimensions();
        assert!(w <= 256);
        assert!(h <= 256);
    }

    #[test]
    fn test_concentration_to_grayscale() {
        let concentration =
            ndarray::Array2::from_shape_fn((10, 10), |(y, x)| (x + y) as f32 / 20.0);

        let img = concentration_to_grayscale(&concentration).unwrap();
        assert_eq!(img.dimensions(), (10, 10));
    }
}