//! Tile pyramid generation for deep-zoom viewing
//!
//! Generates multi-resolution tile pyramids compatible with OpenSeadragon.

use crate::config::PyramidProcessingConfig;
use anyhow::{Result, bail};
use image::{DynamicImage, GenericImageView, ImageBuffer, RgbImage};
use rayon::prelude::*;
use slidescope_shared::{DziMetadata, TileLayer};
use std::collections::HashMap;
use std::io::Cursor;
use tracing::{debug, info};

/// Tile pyramid for deep-zoom viewing
pub struct TilePyramid {
    /// Original image dimensions
    pub width: u32,
    pub height: u32,
    /// Tile size
    pub tile_size: u32,
    /// Maximum zoom level
    pub max_level: u8,
    /// Tiles stored by (level, x, y) -> JPEG bytes
    tiles: HashMap<(u8, u32, u32), Vec<u8>>,
    /// JPEG quality (stored for metadata)
    #[allow(dead_code)]
    jpeg_quality: u8,
}

impl TilePyramid {
    /// Generate a tile pyramid from an image
    pub fn generate(image: &DynamicImage, config: &PyramidProcessingConfig) -> Result<Self> {
        let (width, height) = image.dimensions();

        if width == 0 || height == 0 {
            bail!("Image dimensions must be non-zero");
        }

        let max_level = compute_max_level(width, height, config.tile_size);

        info!(
            "Generating tile pyramid: {}x{}, tile_size={}, max_level={}",
            width, height, config.tile_size, max_level
        );

        let mut pyramid = Self {
            width,
            height,
            tile_size: config.tile_size,
            max_level,
            tiles: HashMap::new(),
            jpeg_quality: config.jpeg_quality,
        };

        // Generate tiles for each level, from highest resolution to lowest
        for level in (0..=max_level).rev() {
            let level_width = level_dimension(width, max_level, level);
            let level_height = level_dimension(height, max_level, level);

            debug!("Level {}: {}x{}", level, level_width, level_height);

            // Resize image to this level
            let scaled = image.resize_exact(
                level_width,
                level_height,
                image::imageops::FilterType::Lanczos3,
            );

            // Generate tiles for this level
            pyramid.generate_level_tiles(level, &scaled, config)?;
        }

        info!("Pyramid complete: {} total tiles", pyramid.tiles.len());

        Ok(pyramid)
    }

    /// Generate tiles for a single level
    fn generate_level_tiles(
        &mut self,
        level: u8,
        image: &DynamicImage,
        config: &PyramidProcessingConfig,
    ) -> Result<()> {
        let (width, height) = image.dimensions();
        let tile_size = config.tile_size;

        let tiles_x = width.div_ceil(tile_size);
        let tiles_y = height.div_ceil(tile_size);

        let rgb_image = image.to_rgb8();

        if config.parallel {
            // Parallel tile generation
            let tile_coords: Vec<(u32, u32)> = (0..tiles_y)
                .flat_map(|y| (0..tiles_x).map(move |x| (x, y)))
                .collect();

            let tiles: Vec<_> = tile_coords
                .par_iter()
                .map(|&(x, y)| {
                    let tile_data =
                        extract_and_encode_tile(&rgb_image, x, y, tile_size, config.jpeg_quality);
                    ((level, x, y), tile_data)
                })
                .collect();

            for ((level, x, y), data) in tiles {
                if let Ok(tile_bytes) = data {
                    self.tiles.insert((level, x, y), tile_bytes);
                }
            }
        } else {
            // Sequential tile generation
            for y in 0..tiles_y {
                for x in 0..tiles_x {
                    if let Ok(tile_bytes) =
                        extract_and_encode_tile(&rgb_image, x, y, tile_size, config.jpeg_quality)
                    {
                        self.tiles.insert((level, x, y), tile_bytes);
                    }
                }
            }
        }

        Ok(())
    }

    /// Get a tile by level and coordinates
    pub fn get_tile(&self, level: u8, x: u32, y: u32) -> Option<&Vec<u8>> {
        self.tiles.get(&(level, x, y))
    }

    /// Get DZI metadata for OpenSeadragon
    pub fn get_dzi_metadata(&self) -> DziMetadata {
        DziMetadata {
            width: self.width,
            height: self.height,
            tile_size: self.tile_size,
            overlap: 0,
            format: "jpeg".to_string(),
            max_level: self.max_level,
        }
    }

    /// Get number of tiles at a given level
    pub fn tiles_at_level(&self, level: u8) -> (u32, u32) {
        if level > self.max_level {
            return (0, 0);
        }

        let level_width = level_dimension(self.width, self.max_level, level);
        let level_height = level_dimension(self.height, self.max_level, level);

        let tiles_x = level_width.div_ceil(self.tile_size);
        let tiles_y = level_height.div_ceil(self.tile_size);

        (tiles_x, tiles_y)
    }

    /// Get total number of tiles in the pyramid
    pub fn total_tiles(&self) -> usize {
        self.tiles.len()
    }

    /// Get memory usage in bytes (approximate)
    pub fn memory_bytes(&self) -> usize {
        self.tiles.values().map(|t| t.len()).sum()
    }
}

/// Extract a tile from an image and encode as JPEG
fn extract_and_encode_tile(
    image: &RgbImage,
    tile_x: u32,
    tile_y: u32,
    tile_size: u32,
    quality: u8,
) -> Result<Vec<u8>> {
    let (img_width, img_height) = image.dimensions();

    let start_x = tile_x * tile_size;
    let start_y = tile_y * tile_size;

    // Actual tile dimensions (may be smaller at edges)
    let actual_width = (img_width - start_x).min(tile_size);
    let actual_height = (img_height - start_y).min(tile_size);

    // Extract tile pixels
    let mut tile = ImageBuffer::new(actual_width, actual_height);
    for y in 0..actual_height {
        for x in 0..actual_width {
            let pixel = image.get_pixel(start_x + x, start_y + y);
            tile.put_pixel(x, y, *pixel);
        }
    }

    // Encode as JPEG
    let mut buffer = Cursor::new(Vec::new());
    let encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buffer, quality);
    tile.write_with_encoder(encoder)?;

    Ok(buffer.into_inner())
}

/// Compute maximum zoom level for given dimensions
fn compute_max_level(width: u32, height: u32, tile_size: u32) -> u8 {
    let max_dim = width.max(height) as f64;
    let levels = (max_dim / tile_size as f64).log2().ceil();
    levels as u8
}

/// Compute dimension at a given level
fn level_dimension(full_dim: u32, max_level: u8, level: u8) -> u32 {
    if level > max_level {
        return 0;
    }

    let scale = 2_u32.pow((max_level - level) as u32);
    full_dim.div_ceil(scale)
}

/// Pyramid cache for multiple layers (original, stain1, stain2, etc.)
pub struct MultiLayerPyramid {
    pyramids: HashMap<TileLayer, TilePyramid>,
    /// Original image width
    #[allow(dead_code)]
    width: u32,
    /// Original image height
    #[allow(dead_code)]
    height: u32,
}

impl MultiLayerPyramid {
    /// Create a new multi-layer pyramid
    pub fn new(width: u32, height: u32) -> Self {
        Self {
            pyramids: HashMap::new(),
            width,
            height,
        }
    }

    /// Add a layer to the pyramid
    pub fn add_layer(&mut self, layer: TileLayer, pyramid: TilePyramid) {
        self.pyramids.insert(layer, pyramid);
    }

    /// Get a tile from a specific layer
    pub fn get_tile(&self, layer: TileLayer, level: u8, x: u32, y: u32) -> Option<&Vec<u8>> {
        self.pyramids.get(&layer)?.get_tile(level, x, y)
    }

    /// Check if a layer exists
    pub fn has_layer(&self, layer: TileLayer) -> bool {
        self.pyramids.contains_key(&layer)
    }

    /// Get DZI metadata (same for all layers)
    pub fn get_dzi_metadata(&self, layer: TileLayer) -> Option<DziMetadata> {
        self.pyramids.get(&layer).map(|p| p.get_dzi_metadata())
    }

    /// Get available layers
    pub fn layers(&self) -> Vec<TileLayer> {
        self.pyramids.keys().copied().collect()
    }
}

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

    fn test_config() -> PyramidProcessingConfig {
        PyramidProcessingConfig {
            tile_size: 64, // Small for testing
            overlap: 0,
            jpeg_quality: 80,
            parallel: false,
        }
    }

    #[test]
    fn test_compute_max_level() {
        assert_eq!(compute_max_level(256, 256, 256), 0);
        assert_eq!(compute_max_level(512, 512, 256), 1);
        assert_eq!(compute_max_level(1024, 1024, 256), 2);
    }

    #[test]
    fn test_level_dimension() {
        // 1000 pixels, max_level=4, tile_size=256
        assert_eq!(level_dimension(1000, 4, 4), 1000); // Full res
        assert_eq!(level_dimension(1000, 4, 3), 500); // Half
        assert_eq!(level_dimension(1000, 4, 2), 250); // Quarter
        assert_eq!(level_dimension(1000, 4, 1), 125);
        assert_eq!(level_dimension(1000, 4, 0), 63); // Smallest
    }

    #[test]
    fn test_pyramid_generation() {
        let config = test_config();

        // Create a small test image
        let mut img = RgbImage::new(128, 128);
        for y in 0..128 {
            for x in 0..128 {
                let r = (x * 2) as u8;
                let g = (y * 2) as u8;
                let b = 128;
                img.put_pixel(x, y, Rgb([r, g, b]));
            }
        }
        let dynamic_img = DynamicImage::ImageRgb8(img);

        let pyramid = TilePyramid::generate(&dynamic_img, &config).unwrap();

        assert_eq!(pyramid.width, 128);
        assert_eq!(pyramid.height, 128);
        assert!(pyramid.total_tiles() > 0);

        // Should have tile at level 0, position 0,0
        let tile = pyramid.get_tile(pyramid.max_level, 0, 0);
        assert!(tile.is_some());
    }

    #[test]
    fn test_dzi_metadata() {
        let config = test_config();
        let img = DynamicImage::ImageRgb8(RgbImage::new(256, 256));
        let pyramid = TilePyramid::generate(&img, &config).unwrap();

        let dzi = pyramid.get_dzi_metadata();
        assert_eq!(dzi.width, 256);
        assert_eq!(dzi.height, 256);
        assert_eq!(dzi.tile_size, 64);
        assert_eq!(dzi.format, "jpeg");
    }
}