//! # IDX3 File Reader
//!
//! This module provides a reader for IDX3 file format, which is commonly used
//! for storing multidimensional arrays, such as images in machine learning
//! datasets.
//!
//! ## Usage
//!
//! To use the `Idx3Reader`, you need to define a type that implements the
//! `Idx3` as well as the `Idx3Read` trait. This trait specifies the magic word
//! and expected dimensions of the IDX3 file format.
//!
//! Here's an example of how to use the `Idx3Reader`:
//!
//! ```rust
//! use smyl::idx::{Idx, Read, Reader};
//!
//! #[derive(Debug, PartialEq, Eq)]
//! struct MnistImage([u8; 784]);
//!
//! impl Idx for MnistImage {
//!     const EXPECTED_DIMENSIONS: &'static [u32] = &[28, 28];
//!     const MAGIC_WORD: u32 = 0x0803;
//! }
//!
//! impl Read for MnistImage {
//!     fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> {
//!         let mut pixels = [0u8; 784];
//!         reader.read_buf(&mut pixels)?;
//!         Ok(MnistImage(pixels))
//!     }
//! }
//!
//! fn main() -> std::io::Result<()> {
//!     let mut image_reader = Reader::<MnistImage>::try_new(include_bytes!(
//!         "../datasets/mnist/t10k-images-idx3-ubyte.gz"
//!     ))?;
//!
//!     for _ in 0..image_reader.item_count {
//!         let image: MnistImage = image_reader.read()?;
//!         // Do something with the image
//!     }
//!
//!     Ok(())
//! }
//! ```
//!
//! ## Traits
//!
//! The module defines the following traits:
//!
//! - `Idx3Read`: Defines the `read` method for reading values of a specific
//!   type from the IDX3 reader.
//! - `Idx3`: Defines the magic word and expected dimensions of the IDX3 file
//!   format.
//!
//! ## `Idx3Reader`
//!
//! The `Idx3Reader` struct is the main entry point for reading IDX3 files. It
//! takes a byte slice and creates a reader that can be used to read the data
//! from the file.
//!
//! The `try_new` method initializes the reader and performs some basic
//! validation on the file format. The `read` method can then be used to read
//! values of the type specified by the `Idx3` trait implementation.
use std::io::{Error, ErrorKind, Read as IoRead};
use std::marker::PhantomData;

use flate2::read::GzDecoder;

/// Defines the `read` method for reading values of a specific type from the
/// IDX3 reader.
pub trait Read: Sized {
    /// Reads a value of the specified type from the IDX3 reader.
    /// # Errors
    /// Will return `Err` if the data is invalid.
    fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self>;
}

/// Defines the magic word and expected dimensions of the IDX3 file format.
pub trait Idx: Read {
    /// The magic word that identifies the IDX3 file format.
    const MAGIC_WORD: u32;

    /// The expected dimensions of the data in the IDX3 file.
    const EXPECTED_DIMENSIONS: &'static [u32];
}

/// A reader for IDX3 file format.
pub struct Reader<'a, I> {
    decoder:        GzDecoder<&'a [u8]>,
    pub item_count: usize,
    marker:         PhantomData<I>,
}

impl<'a, I> Reader<'a, I> {
    /// Creates a new `Idx3Reader` from the provided byte slice.
    ///
    /// # Errors
    ///
    /// Will return Err if the file cannot be opened, if it's not correctly
    /// compressed, if the magic word is different from the one specified or
    /// if the item dimensions does not match the specified.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::idx::{Idx, Read, Reader};
    ///
    /// #[derive(Debug, PartialEq, Eq)]
    /// struct MnistImage([u8; 784]);
    ///
    /// impl Idx for MnistImage {
    ///     const EXPECTED_DIMENSIONS: &'static [u32] = &[28, 28];
    ///     const MAGIC_WORD: u32 = 0x0803;
    /// }
    ///
    /// impl Read for MnistImage {
    ///     fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> {
    ///         let mut pixels = [0u8; 784];
    ///         reader.read_buf(&mut pixels)?;
    ///         Ok(MnistImage(pixels))
    ///     }
    /// }
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let bytes = include_bytes!("../datasets/mnist/t10k-images-idx3-ubyte.gz");
    ///     let mut reader = Reader::<MnistImage>::try_new(bytes)?;
    ///     // Use the reader to read data
    ///     Ok(())
    /// }
    /// ```
    pub fn try_new(bytes: &'a [u8]) -> std::io::Result<Self>
    where
        I: Idx,
    {
        let decoder = GzDecoder::new(bytes);
        let mut reader = Self {
            decoder,
            item_count: 0,
            marker: PhantomData,
        };

        if I::MAGIC_WORD != reader.read_other()? {
            return Err(Error::new(ErrorKind::InvalidData, "Invalid Magic Word"));
        }

        reader.item_count = reader.read_other::<u32>()? as usize;

        for dim in I::EXPECTED_DIMENSIONS {
            if dim != &reader.read_other()? {
                return Err(Error::new(ErrorKind::InvalidData, "Invalid Dimensions"));
            }
        }

        Ok(reader)
    }

    /// Reads a value of the type specified by the `Idx3` trait implementation.
    ///
    /// # Errors
    /// Will return `Err` if the data is invalid.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::idx::{Idx, Read, Reader};
    ///
    /// #[derive(Debug, PartialEq, Eq)]
    /// struct MnistImage([u8; 784]);
    ///
    /// impl Idx for MnistImage {
    ///     const EXPECTED_DIMENSIONS: &'static [u32] = &[28, 28];
    ///     const MAGIC_WORD: u32 = 0x0803;
    /// }
    ///
    /// impl Read for MnistImage {
    ///     fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> {
    ///         let mut pixels = [0u8; 784];
    ///         reader.read_buf(&mut pixels)?;
    ///         Ok(MnistImage(pixels))
    ///     }
    /// }
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let bytes = include_bytes!("../datasets/mnist/t10k-images-idx3-ubyte.gz");
    ///     let mut reader = Reader::<MnistImage>::try_new(bytes)?;
    ///
    ///     for _ in 0..reader.item_count {
    ///         let image: MnistImage = reader.read()?;
    ///         // Do something with the image
    ///     }
    ///
    ///     Ok(())
    /// }
    /// ```
    #[inline]
    pub fn read(&mut self) -> std::io::Result<I>
    where
        I: Read,
    {
        self.read_other::<I>()
    }

    /// Reads a value of the specified type from the IDX3 reader.
    ///
    /// # Errors
    /// Will return `Err` if the data is invalid.
    #[inline]
    pub fn read_other<T>(&mut self) -> std::io::Result<T>
    where
        T: Read,
    {
        T::read(self)
    }

    /// Reads a buffer of bytes from the IDX3 reader.
    ///
    /// # Errors
    /// Will return `Err` if not enough bytes are available to fill the buffer.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::idx::{Idx, Read, Reader};
    ///
    /// #[derive(Debug, PartialEq, Eq)]
    /// struct Dummy;
    ///
    /// impl Read for Dummy {
    ///     fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> { Ok(Dummy) }
    /// }
    ///
    /// impl Idx for Dummy {
    ///     const EXPECTED_DIMENSIONS: &'static [u32] = &[];
    ///     const MAGIC_WORD: u32 = 0x0803;
    /// }
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let bytes = include_bytes!("../datasets/mnist/t10k-images-idx3-ubyte.gz");
    ///     let mut reader = Reader::<Dummy>::try_new(bytes)?;
    ///
    ///     let mut buffer = [0u8; 100];
    ///     reader.read_buf(&mut buffer)?;
    ///     // Use the buffer
    ///     Ok(())
    /// }
    /// ```
    #[inline]
    pub fn read_buf(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
        self.decoder.read_exact(buf)
    }
}

impl Read for u8 {
    fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> {
        let mut buf = [0u8];
        reader.read_buf(&mut buf)?;
        Ok(buf[0])
    }
}

impl Read for u32 {
    fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> {
        let mut buf = [0u8; 4];
        reader.read_buf(&mut buf)?;
        Ok(Self::from_be_bytes(buf))
    }
}
