imzml 0.1.3

use flate2::bufread::ZlibDecoder;
use quick_xml::events::{BytesStart, Event};
use std::collections::VecDeque;
use std::io::{BufRead, Cursor, Read, Write};

use super::cvparam::{search_all_params, CVParam, HasCVParams, HasParamGroupRefs};
use super::writer::Writer;
use super::{MzMLReader, MzMLTag, Tag};
use crate::error::{FatalParseError, ParseError};
use crate::mzml::attributes::{AttributeValue, LIST_ATTRIBUTES};
use crate::mzml::cvparam::{CVParamValue, RawTerm, UserParam};

use super::referenceableparamgroup::ReferenceableParamGroupRef;

/// Accession in OBO for m/z array
pub const ACCESSION_MZ_ARRAY: &str = "MS:1000514";
/// Accession in OBO for intensity array
pub const ACCESSION_INTENSITY_ARRAY: &str = "MS:1000515";

/// Accession in OBO for binary data type
pub const ACCESSION_BINARY_DATA_TYPE: &str = "MS:1000518";
/// Accession in OBO for 64-bit float
pub const ACCESSION_64_BIT_FLOAT: &str = "MS:1000523";
/// Accession in OBO for 32-bit float
pub const ACCESSION_32_BIT_FLOAT: &str = "MS:1000521";

/// Accession in OBO for compression (parent term)
pub const ACCESSION_COMPRESSION: &str = "MS:1000572";
/// Accession in OBO for no compression
pub const ACCESSION_NO_COMPRESSION: &str = "MS:1000576";
/// Accession in OBO for Zlib compression
pub const ACCESSION_ZLIB_COMPRESSION: &str = "MS:1000574";

/// Accession in OBO for external array length
pub const ACCESSION_EXTERNAL_ARRAY_LENGTH: &str = "IMS:1000103";
const B_ACCESSION_EXTERNAL_ARRAY_LENGTH: &[u8] = b"IMS:1000103";

/// Accession in OBO for external encoded length
pub const ACCESSION_EXTERNAL_ENCODED_LENGTH: &str = "IMS:1000104";
const B_ACCESSION_EXTERNAL_ENCODED_LENGTH: &[u8] = b"IMS:1000104";

/// Accession in OBO for external offset
pub const ACCESSION_EXTERNAL_OFFSET: &str = "IMS:1000102";
const B_ACCESSION_EXTERNAL_OFFSET: &[u8] = b"IMS:1000102";

/// Accession in OBO for external data
pub const ACCESSION_EXTERNAL_DATA: &str = "IMS:1000101";
const B_ACCESSION_EXTERNAL_DATA: &[u8] = b"IMS:1000101";

/// The type of data being stored (e.g. m/z values, intensities)
#[derive(Debug, PartialEq, Eq)]
pub enum BinaryDataArrayType {
    /// m/z array
    MZArray,
    /// Intensity array
    IntenistyArray,
}

/// The type used to store binary data (e.g. m/zs or intensities)
#[derive(Debug, Clone, Copy)]
pub enum BinaryDataType {
    /// Another format not specified here
    Undefined,
    /// f64
    Float64,
    /// f32
    Float32,
}

/// <binaryDataArrayList> tag representation. Contains a list of BinaryDataArray
pub struct BinaryDataArrayList {
    // TODO: Currently need this to avoid breaking code with the old bda_list: Vec<BinaryDataArray> parameter on MzML struct
    pub(crate) list: Vec<BinaryDataArray>,
}

impl BinaryDataArrayList {
    fn new(count: usize) -> Self {
        BinaryDataArrayList {
            list: Vec::with_capacity(count),
        }
    }
}

impl MzMLTag for BinaryDataArrayList {
    fn parse_start_tag<B: BufRead>(
        parser: &mut MzMLReader<B>,
        start_event: &BytesStart,
    ) -> Result<Option<Self>, FatalParseError>
    where
        Self: std::marker::Sized,
    {
        if start_event.name().local_name().as_ref() != b"binaryDataArrayList" {
            Err(FatalParseError::UnexpectedTag(format!(
                "Unexpected event {:?} when processing BinaryDataArrayList",
                start_event,
            )))
        } else {
            let attributes = parser.process_attributes(
                Tag::BinaryDataArrayList,
                &LIST_ATTRIBUTES,
                start_event,
            )?;

            let count = match attributes.get("count") {
                Some(&AttributeValue::Integer(count)) => count as usize,
                _ => 0,
            };

            parser
                .breadcrumbs
                .push_back((Tag::BinaryDataArrayList, None));

            Ok(Some(BinaryDataArrayList::new(count)))
        }
    }

    fn parse_xml<B: BufRead>(
        &mut self,
        parser: &mut MzMLReader<B>,
        buffer: &mut Vec<u8>,
    ) -> Result<(), FatalParseError> {
        let mut last_num_params = 0;

        // Check what comes next
        loop {
            // Clear the buffer ready for the next tag
            buffer.clear();

            let next_event = parser.next(buffer)?;

            match next_event {
                Event::Start(start_event) | Event::Empty(start_event) => {
                    match start_event.name().as_ref() {
                        b"binaryDataArray" => {
                            if let Some(mut bda) =
                                BinaryDataArray::parse_start_tag(parser, &start_event)?
                            {
                                bda.cv_params.reserve(last_num_params);
                                bda.parse_xml(parser, buffer)?;

                                last_num_params = bda.cv_params.len();

                                self.list.push(bda);
                            }
                        }
                        _ => parser.errors.push_back(ParseError::UnexpectedTag(format!(
                            "{:?} unexpected when processing {:?}",
                            std::str::from_utf8(start_event.name().as_ref()),
                            Tag::BinaryDataArrayList
                        ))),
                    }
                }
                Event::End(end_event) => {
                    if let b"binaryDataArrayList" = end_event.name().as_ref() {
                        parser.breadcrumbs.pop_back();

                        break;
                    }
                }
                Event::Eof => {
                    return Err(FatalParseError::MissingClosingTag(
                        "binaryDataArrayList".to_string(),
                    ));
                }
                _ => {}
            }
        }

        Ok(())
    }

    fn tag() -> Tag {
        Tag::BinaryDataArrayList
    }

    fn write_xml<W: Write>(&self, writer: &mut Writer<W>) -> Result<(), quick_xml::Error> {
        writer.write_list("binaryDataArrayList", &self.list)
    }
}

/// Possible means of compressing binary data stored in (i)mzML
#[derive(Debug)]
pub enum Compression {
    /// No compression
    None,
    /// Compression has not been defined in the metadata
    Undefined,
    /// Z-lib compression
    Zlib,
}

/// Provides a single interface to compressed data, decompressing the data on `read()` if necessary
pub struct CompressionReader {
    reader: Box<dyn Read>,
}

impl Read for CompressionReader {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.reader.read(buf)
    }
}

impl Compression {
    /// Create a `CompressionReader` for the supplied data
    pub fn to_reader(&self, data: Vec<u8>) -> CompressionReader {
        match self {
            Compression::None | Compression::Undefined => CompressionReader {
                reader: Box::new(Cursor::new(data)),
            },
            Compression::Zlib => CompressionReader {
                reader: Box::new(ZlibDecoder::new(Cursor::new(data))),
            },
        }
    }
}

/// BinaryDataArray represents the <binaryDataArray> tag and captures information required to read and parse the stored binary data.
#[derive(Debug)]
pub struct BinaryDataArray {
    param_group_refs: Vec<ReferenceableParamGroupRef>,
    cv_params: Vec<CVParam>,
    user_params: Vec<UserParam>,

    is_external: bool,

    array_length: Option<u64>,
    encoded_length: Option<u64>,
    offset: Option<u64>,
}

impl Clone for BinaryDataArray {
    fn clone(&self) -> Self {
        Self {
            param_group_refs: self.param_group_refs.clone(),
            cv_params: self.cv_params.clone(),
            user_params: self.user_params.clone(),
            is_external: self.is_external,
            array_length: self.array_length,
            encoded_length: self.encoded_length,
            offset: self.offset,
        }
    }
}

impl Default for BinaryDataArray {
    fn default() -> Self {
        Self::new()
    }
}

impl BinaryDataArray {
    /// Create a new `BinaryDataArray` with no CV params
    pub fn new() -> Self {
        BinaryDataArray {
            param_group_refs: Vec::new(),
            cv_params: Vec::new(),
            user_params: Vec::new(),

            is_external: true,

            array_length: None,
            encoded_length: None,
            offset: None,
        }
    }

    pub(crate) fn set_external_params(
        &mut self,
        offset: u64,
        array_length: u64,
        encoded_length: u64,
    ) {
        self.is_external = true;
        self.offset = Some(offset);
        self.array_length = Some(array_length);
        self.encoded_length = Some(encoded_length);
    }

    /// Return the `BinaryDataArrayType` (e.g. m/z array, intensity array) based on the CV params stored within.
    pub fn array_type(&self) -> Option<BinaryDataArrayType> {
        if let Some(_cv_param) = search_all_params(self, ACCESSION_MZ_ARRAY) {
            return Some(BinaryDataArrayType::MZArray);
        }
        if let Some(_cv_param) = search_all_params(self, ACCESSION_INTENSITY_ARRAY) {
            return Some(BinaryDataArrayType::IntenistyArray);
        }
        None
    }

    /// Return the `BinaryDataType` (e.g. f64, f32, ...) based on the CV params stored within.
    pub fn binary_type(&self) -> BinaryDataType {
        if let Some(_cv_param) = search_all_params(self, ACCESSION_64_BIT_FLOAT) {
            return BinaryDataType::Float64;
        }
        if let Some(_cv_param) = search_all_params(self, ACCESSION_32_BIT_FLOAT) {
            return BinaryDataType::Float32;
        }

        BinaryDataType::Undefined
    }

    /// Return the `Compression` (e.g. None, Z-lib, ...) used to compress the binary data, based on the CV params stored within.
    pub fn compression(&self) -> Compression {
        if let Some(_cv_param) = search_all_params(self, ACCESSION_NO_COMPRESSION) {
            return Compression::None;
        }
        if let Some(_cv_param) = search_all_params(self, ACCESSION_ZLIB_COMPRESSION) {
            return Compression::Zlib;
        }

        Compression::Undefined
    }

    /// Return the number of elements in the array.
    pub fn array_length(&self) -> Option<u64> {
        self.array_length
    }

    /// Return the number of bytes needed to store the array (in the specified data type, and after compression).
    pub fn encoded_length(&self) -> Option<u64> {
        self.encoded_length
    }

    /// Return the file offset depicting where the array is stored.
    pub fn offset(&self) -> Option<u64> {
        self.offset
    }

    /// Return true if the array is stored in an external file (e.g. ibd). Always true for imzML files.
    pub fn is_data_external(&self) -> bool {
        self.is_external
    }

    pub(crate) fn add_raw_term<'a>(
        &mut self,
        raw_term: &RawTerm<'a>,
        breadcrumbs: &VecDeque<(Tag, Option<String>)>,
        errors: &mut VecDeque<ParseError>,
        ignore_uncommon_tags: bool,
    ) {
        match raw_term.raw_accession() {
            B_ACCESSION_EXTERNAL_ARRAY_LENGTH => {
                self.array_length = Some(raw_term.value_as_u64());
            }
            B_ACCESSION_EXTERNAL_ENCODED_LENGTH => {
                self.encoded_length = Some(raw_term.value_as_u64());
            }
            B_ACCESSION_EXTERNAL_OFFSET => {
                self.offset = Some(raw_term.value_as_u64());
            }
            B_ACCESSION_EXTERNAL_DATA => {}
            _ => {
                //let cv_param: Result<CVParam<'a>, crate::ParseError> = raw_term.into();

                if !ignore_uncommon_tags {
                    self.cv_params
                        .push(raw_term.to_cv_param(breadcrumbs, errors))
                }
            }
        }
    }
}

impl MzMLTag for BinaryDataArray {
    fn parse_start_tag<B: BufRead>(
        parser: &mut MzMLReader<B>,
        start_event: &BytesStart,
    ) -> Result<Option<Self>, FatalParseError>
    where
        Self: std::marker::Sized,
    {
        if start_event.name().local_name().as_ref() != b"binaryDataArray" {
            Err(FatalParseError::UnexpectedTag(format!(
                "Unexpected event {:?} when processing BinaryDataArray",
                start_event,
            )))
        } else {
            // TODO: encodedLength

            parser.breadcrumbs.push_back((Tag::BinaryDataArray, None));

            Ok(Some(BinaryDataArray::new()))
        }
    }

    fn parse_xml<B: BufRead>(
        &mut self,
        parser: &mut MzMLReader<B>,
        buffer: &mut Vec<u8>,
    ) -> Result<(), FatalParseError> {
        // Check what comes next
        loop {
            // Clear the buffer ready for the next tag
            buffer.clear();

            let next_event = parser.next(buffer)?;
            let is_empty = matches!(next_event, Event::Empty(_));

            match next_event {
                Event::Start(start_event) | Event::Empty(start_event) => {
                    match start_event.name().as_ref() {
                        b"cvParam" => {
                            let raw_term = RawTerm::parse_start_tag(parser, &start_event)?;
                            self.add_raw_term(
                                &raw_term,
                                &parser.breadcrumbs,
                                &mut parser.errors,
                                parser.ignore_uncommon_tags,
                            );
                        }
                        b"referenceableParamGroupRef" => {
                            let param_group_ref =
                                ReferenceableParamGroupRef::parse_start_tag(parser, &start_event)?;
                            self.param_group_refs.push(param_group_ref);
                        }
                        b"binary" => {
                            if !is_empty && self.offset.is_none() {
                                self.is_external = false;

                                self.offset = Some(parser.reader.buffer_position() as u64);
                            }
                        }
                        _ => parser.errors.push_back(ParseError::UnexpectedTag(format!(
                            "{:?} unexpected when processing {:?}",
                            std::str::from_utf8(start_event.name().as_ref()),
                            Tag::BinaryDataArray
                        ))),
                    }
                }
                Event::Text(_) => {
                    if !self.is_external {
                        self.encoded_length =
                            Some(parser.reader.buffer_position() as u64 - self.offset.unwrap());
                    }
                }
                Event::End(end_event) => {
                    if let b"binaryDataArray" = end_event.name().as_ref() {
                        // Finished with gathering information, so now we can generate the DataAccess
                        // self.data_access = Some(DataAccess::OnDisk(OnDiskAccess {
                        //     file: parser.data_location.clone(),
                        //     encoded_length: self.encoded_length.unwrap(),
                        //     offset: self.offset.unwrap(),
                        // }));

                        parser.breadcrumbs.pop_back();

                        break;
                    }
                }
                Event::Eof => {
                    return Err(FatalParseError::MissingClosingTag(
                        "binaryDataArray".to_string(),
                    ));
                }
                _ => {}
            }
        }

        Ok(())
    }

    fn tag() -> Tag {
        Tag::BinaryDataArray
    }

    fn write_xml<W: Write>(&self, writer: &mut Writer<W>) -> Result<(), quick_xml::Error> {
        writer.start_tag_with_attr(
            "binaryDataArray",
            "encodedLength",
            self.encoded_length.unwrap_or(0),
        )?;

        self.write_ref_param_groups_xml(writer)?;

        if let Some(array_length) = self.array_length {
            writer.write_param(
                ACCESSION_EXTERNAL_ARRAY_LENGTH,
                CVParamValue::NonNegativeInteger(array_length),
            )?;
        }

        if let Some(encoded_length) = self.encoded_length {
            writer.write_param(
                ACCESSION_EXTERNAL_ENCODED_LENGTH,
                CVParamValue::NonNegativeInteger(encoded_length),
            )?;
        }

        if let Some(offset) = self.offset {
            writer.write_param(
                ACCESSION_EXTERNAL_OFFSET,
                CVParamValue::NonNegativeInteger(offset),
            )?;
        }

        // Include the parameter stating that the data is external
        // TODO: Check it doesn't exist already?
        writer.write_param(ACCESSION_EXTERNAL_DATA, CVParamValue::Empty)?;

        self.write_params_xml(writer)?;

        writer.empty_tag("binary")?;
        writer.end_tag("binaryDataArray")
    }
}

// struct BinaryDataArrayCVParamIterator<'obo, 'a> {
//     bda: &'a BinaryDataArray<'obo>,
//     extra: Vec<CVParam<'obo>>,
// }

// impl<'obo, 'a> BinaryDataArrayCVParamIterator<'obo, 'a> {
//     fn new(bda: &'obo BinaryDataArray) -> Self {
//         Self {
//             bda,
//             extra: Vec::new(),
//         }
//     }
// }

// impl<'obo, 'a> Iterator for BinaryDataArrayCVParamIterator<'obo, 'a> {
//     type Item = CVParam<'a>;

//     fn next(&mut self) -> Option<Self::Item> {
//         Some(self.extra.get(0).unwrap().clone())
//     }
// }

//impl<'a> WriteCVParams<'a> for BinaryDataArray<'a> {}
impl HasCVParams for BinaryDataArray {
    fn add_cv_param(&mut self, param: CVParam) {
        match param.accession() {
            ACCESSION_EXTERNAL_ARRAY_LENGTH => {
                self.array_length = Some(param.value_as_u64().unwrap());
            }
            ACCESSION_EXTERNAL_ENCODED_LENGTH => {
                self.encoded_length = Some(param.value_as_u64().unwrap());
            }
            ACCESSION_EXTERNAL_OFFSET => {
                self.offset = Some(param.value_as_u64().unwrap());
            }
            ACCESSION_EXTERNAL_DATA => {}
            /*ACCESSION_EXTERNAL_OFFSET => match param.get_value().as_ref().unwrap().parse() {
                Ok(value) => self.offset = Some(value),
                Err(_err) => self.cv_params.push(param),
            },*/
            _ => {
                self.cv_params.push(param);
            }
        }
    }

    fn cv_params(&self) -> &Vec<CVParam> {
        &self.cv_params
    }

    fn cv_params_mut(&mut self) -> &mut Vec<CVParam> {
        self.cv_params.as_mut()
    }

    fn add_user_param(&mut self, param: UserParam) {
        self.user_params.push(param);
    }

    fn user_params(&self) -> &Vec<UserParam> {
        &self.user_params
    }

    // fn cv_param_iter(&self) -> CVParamIterator {
    //     let params = self.clone_params_with_groups();

    //     // Add in

    //     CVParamIterator::from_cv_params(params)
    // }
}

impl HasParamGroupRefs for BinaryDataArray {
    fn add_param_group_ref(&mut self, param_group_ref: ReferenceableParamGroupRef) {
        self.param_group_refs.push(param_group_ref);
    }

    fn param_group_refs(&self) -> &Vec<ReferenceableParamGroupRef> {
        &self.param_group_refs
    }
}

// impl<'a> ImzMLParser<'a> {
//     pub(super) fn start_binary_data_array_list(
//         &mut self,
//         e: &BytesStart,
//     ) -> Result<Tag, FatalParseError> {
//         let attributes = self.process_attributes(Tag::BinaryDataArrayList, &LIST_ATTRIBUTES, e);

//         if let Some(&AttributeValue::Integer(count)) = attributes.get("count") {
//             if let Some(current_spectrum) = self.current_spectrum.as_mut() {
//                 current_spectrum.binary_data_array_list = Vec::with_capacity(count as usize);
//             } else if let Some(current_chromatogram) = self.current_chromatogram.as_mut() {
//                 current_chromatogram.binary_data_array_list = Vec::with_capacity(count as usize);
//             }
//         }

//         /*let mut has_error: bool = false;

//         for att in e.attributes() {
//             match att {
//                 Ok(attribute) => match attribute.key {
//                     b"count" => match parse_usize(attribute) {
//                         Ok(value) => {
//                             self.current_spectrum
//                                 .as_mut()
//                                 .unwrap()
//                                 .binary_data_array_list = Vec::with_capacity(value);
//                         }
//                         Err(error) => {
//                             self.errors
//                                 .push_back(ParseError::IntError((Tag::BinaryDataArrayList, error)));
//                             has_error = true;
//                             break;
//                         }
//                     },
//                     _ => {
//                         self.errors.push_back(ParseError::UnexpectedAttribute((
//                             Tag::BinaryDataArrayList,
//                             std::str::from_utf8(attribute.key).unwrap().to_string(),
//                         )));
//                         has_error = true;
//                         break;
//                     }
//                 },
//                 Err(error) => {
//                     self.errors
//                         .push_back(ParseError::XMLError((Tag::BinaryDataArrayList, error)));
//                     has_error = true;
//                     break;
//                 }
//             };
//         }*/
//         /*if has_error {
//             ParserState::Error
//         } else {
//             ParserState::Processing
//         }*/
//         Ok(Tag::BinaryDataArrayList)
//     }

//     pub(super) fn end_binary_data_array(&mut self) -> ParserState {
//         match self.current_binary_data_array.take() {
//             Some(binary_data_array) => {
//                 if let Some(current_spectrum) = self.current_spectrum.as_mut() {
//                     current_spectrum.add_binary_data_array(binary_data_array);
//                 } else if let Some(current_chromatogram) = self.current_chromatogram.as_mut() {
//                     current_chromatogram.add_binary_data_array(binary_data_array);
//                 } else {
//                     // TODO: ERROR neither current_spectrum nor current_chromatogram set
//                 }
//             }
//             None => {
//                 // TODO: Error here?
//                 self.errors.push_back(ParseError::UnexpectedTag(
//                     "Found </binaryDataArray> but no <binaryDataArray>".to_string(),
//                 ))
//             }
//         };

//         ParserState::Processing
//     }
// }