imzml 0.1.3

use std::io::{BufRead, Write};
use std::sync::Arc;

use quick_xml::events::{BytesStart, Event};

use crate::mzml::cvlist::{CVList, CV};
use crate::mzml::data_processing::{DataProcessing, DataProcessingList, DataProcessingRef};
use crate::mzml::filedescription::{FileDescription, SourceFileRef};
use crate::mzml::instrument::{
    InstrumentConfiguration, InstrumentConfigurationList, InstrumentConfigurationRef,
};
use crate::mzml::referenceableparamgroup::{
    ReferenceableParamGroup, ReferenceableParamGroupList, ReferenceableParamGroupRef,
};
use crate::mzml::run::Run;
use crate::mzml::sample::{Sample, SampleList, SampleRef};
use crate::mzml::scan_settings::{ScanSettings, ScanSettingsList};
use crate::mzml::software::SoftwareList;
use crate::mzml::software::{Software, SoftwareRef};
use crate::mzml::spectrum::Spectrum;
use crate::mzml::writer::Writer;
use crate::obo::Ontology;
use crate::{FatalParseError, ParseError, SpectrumList, SpectrumRef, Tag};

use super::{MzMLReader, MzMLTag};

const XMLNS: &str = "http://psi.hupo.org/ms/mzml";
const XMLNS_XSI: &str = "http://www.w3.org/2001/XMLSchema-instance";
const XSI_SCHEMA_LOCATION: &str =
    "http://psi.hupo.org/ms/mzml http://psidev.info/files/ms/mzML/xsd/mzML1.1.0_idx.xsd";
/// Current mzML version
pub const MZML_VERSION: &str = "1.1.0";

// pub trait Element {
//     fn follow(path: &str) {}
// }

/// IndexedMzML represents an mzML file with an <indexedMzML> tag present
pub struct IndexedMzML {
    mzml: Option<MzML>,

    pub(crate) index_list: Vec<Index>,
}

impl IndexedMzML {
    pub(crate) fn new() -> Self {
        IndexedMzML {
            mzml: None,
            index_list: Vec::new(),
        }
    }
}

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

impl From<IndexedMzML> for MzML {
    fn from(mut indexed_mzml: IndexedMzML) -> Self {
        let mut mzml = indexed_mzml.mzml.take().unwrap();
        mzml.index_list = Some(indexed_mzml.index_list);

        mzml
    }
}

impl MzMLTag for IndexedMzML {
    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.local_name().as_ref() != b"indexedmzML" {
            return Err(FatalParseError::UnexpectedTag(format!(
                "Unexpected event {:?} when processing IndexedMzML",
                start_event,
            )));
        }

        for att in start_event
            .attributes()
            .with_checks(parser.with_attribute_checks)
        {
            match att {
                Ok(att) => {
                    let key = att.key.as_ref();

                    match key {
                        b"xmlns" => {}
                        b"xmlns:xsi" => {}
                        b"xsi:schemaLocation" => {}
                        _ => {
                            parser.errors.push_back(ParseError::UnexpectedAttribute((
                                Tag::IndexedMzML,
                                std::str::from_utf8(key).unwrap().to_string(),
                            )));
                        }
                    }
                }
                Err(error) => {
                    parser
                        .errors
                        .push_back(ParseError::XMLError((Tag::IndexedMzML, error.into())));
                }
            };
        }

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

        Ok(Some(IndexedMzML::new()))
    }

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

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

            let current_event = parser.next(buffer)?;

            match current_event {
                Event::Start(start_event) | Event::Empty(start_event) => {
                    match start_event.name().as_ref() {
                        b"mzML" => {
                            if let Some(mut mzml) = MzML::parse_start_tag(parser, &start_event)? {
                                mzml.parse_xml(parser, buffer)?;
                                self.mzml = Some(mzml);
                            }
                        }
                        _ => parser.errors.push_back(ParseError::UnexpectedTag(format!(
                            "{:?} unexpected when processing {:?}",
                            std::str::from_utf8(start_event.name().as_ref()),
                            Tag::IndexedMzML
                        ))),
                    }
                }
                Event::End(end_event) => {
                    if let b"indexedmzML" = end_event.name().as_ref() {
                        parser.breadcrumbs.pop_back();

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

        Ok(())
    }

    fn write_xml<W: Write>(&self, _writer: &mut Writer<W>) -> Result<(), quick_xml::Error> {
        todo!()
    }
}

/// Index represents a <spectrum> tag offset in the mzML file
pub struct Index {
    name: String,
    offsets: Vec<Offset>,
}

impl Clone for Index {
    fn clone(&self) -> Self {
        Self {
            name: self.name.clone(),
            offsets: self.offsets.clone(),
        }
    }
}

/// <spectrum> tag offset in the mzML file
pub struct Offset {
    //<offset idRef="controllerType=0 controllerNumber=1 scan=1">4588</offset>
    id_ref: String,
    offset: u64,
}

impl Clone for Offset {
    fn clone(&self) -> Self {
        Self {
            id_ref: self.id_ref.clone(),
            offset: self.offset,
        }
    }
}

/// MzML represents a <mzML> tag and its children parsed from an mzML file
pub struct MzML {
    accession: Option<String>,
    id: Option<String>,
    version: String,

    ontology: Ontology,

    pub(super) cv_list: Vec<CV>,
    pub(super) file_description: Option<FileDescription>,
    pub(crate) referenceable_param_group_list: Vec<Arc<ReferenceableParamGroup>>,
    pub(crate) sample_list: Vec<Arc<Sample>>,
    pub(crate) software_list: Vec<Arc<Software>>,
    pub(crate) scan_settings_list: ScanSettingsList,
    pub(crate) instrument_configuration_list: Vec<Arc<InstrumentConfiguration>>,
    pub(crate) data_processing_list: Vec<Arc<DataProcessing>>,

    pub(super) run: Option<Run>,

    index_list: Option<Vec<Index>>,
}

impl Clone for MzML {
    fn clone(&self) -> Self {
        let referenceable_param_group_list = self
            .referenceable_param_group_list
            .iter()
            .map(|param_group| Arc::new((**param_group).clone()))
            .collect();

        let sample_list = self
            .sample_list
            .iter()
            .map(|sample| Arc::new((**sample).clone()))
            .collect();

        let software_list = self
            .software_list
            .iter()
            .map(|software| Arc::new((**software).clone()))
            .collect();

        let scan_settings_list = self.scan_settings_list.clone();

        let instrument_configuration_list = self
            .instrument_configuration_list
            .iter()
            .map(|instrument| Arc::new((**instrument).clone()))
            .collect();

        let data_processing_list = self
            .data_processing_list
            .iter()
            .map(|data_processing| Arc::new((**data_processing).clone()))
            .collect();

        // TODO: Refs are not properly cloned - ideally we should replace them with Arc<> references to the cloned values

        Self {
            accession: self.accession.clone(),
            id: self.id.clone(),
            version: self.version.clone(),
            ontology: self.ontology.clone(),
            cv_list: self.cv_list.clone(),
            file_description: self.file_description.clone(),
            referenceable_param_group_list,
            sample_list,
            software_list,
            scan_settings_list,
            instrument_configuration_list,
            data_processing_list,
            run: self.run.clone(),

            index_list: self.index_list.clone(),
        }
    }
}

// pub enum ParseImzMLError {
//     InvalidUTF8(std::str::Utf8Error),
// }

impl MzML {
    /// Create a new `MzML` instance with the specified version and ontology
    pub fn new(version: &str, ontology: Ontology) -> Self {
        MzML {
            accession: None,
            id: None,
            version: version.to_string(),
            //obo_dict: std::collections::HashMap::with_capacity(100),
            ontology,

            cv_list: Vec::new(),
            file_description: None,
            referenceable_param_group_list: Vec::new(),
            sample_list: Vec::new(),
            software_list: Vec::new(),
            scan_settings_list: ScanSettingsList::new(),
            instrument_configuration_list: Vec::new(),
            data_processing_list: Vec::new(),
            run: None,
            index_list: None,
        }
    }

    /// Returns the ontology used to link all accessions (cvParam) found in the mzML file
    pub fn ontology(&self) -> &Ontology {
        &self.ontology
    }

    /// Returns the `FileDescription` associated with the mzML tag.
    pub fn file_description(&self) -> Option<&FileDescription> {
        self.file_description.as_ref()
    }

    /// Returns mutable access to the `FileDescription` associated with the mzML tag.
    pub fn file_description_mut(&mut self) -> Option<&mut FileDescription> {
        self.file_description.as_mut()
    }

    /// Sets the `FileDescription` associated with the mzML tag.
    pub fn set_file_description(&mut self, file_description: FileDescription) {
        self.file_description = Some(file_description);
    }

    /// Sets the `Run` associated with the mzML tag.
    pub fn set_run(&mut self, run: Run) {
        self.run = Some(run);
    }

    /// Adds a `ScanSettings` tag to the `ScanSettingsList`
    pub fn add_scan_settings(&mut self, scan_settings: ScanSettings) {
        self.scan_settings_list.push(scan_settings);
    }

    /// Add an `InstrumentConfiguration` to the mzML tag
    pub fn add_instrument_configuration(
        &mut self,
        instrument_configuration: InstrumentConfiguration,
    ) {
        self.instrument_configuration_list
            .push(Arc::new(instrument_configuration));
    }

    /// Add a new data processing element to the MzML
    pub fn add_data_processing(&mut self, data_processing: DataProcessing) {
        self.data_processing_list.push(Arc::new(data_processing));
    }

    /// Adds a `ReferenceableParamGroup` to the mzML tag.
    pub fn add_referenceable_param_group(
        &mut self,
        referenceable_param_group: ReferenceableParamGroup,
    ) {
        self.referenceable_param_group_list
            .push(Arc::new(referenceable_param_group));
    }

    /// Returns a `ReferenceableParamGroupRef` with the specified `id` if a `ReferenceableParamGroup` with the
    /// corresponding id is present
    pub fn referenceable_param_group_ref(&self, id: &str) -> Option<ReferenceableParamGroupRef> {
        for referenceable_param_group in &self.referenceable_param_group_list {
            if referenceable_param_group.id() == id {
                return Some(ReferenceableParamGroupRef::Ref(
                    referenceable_param_group.clone(),
                ));
            }
        }

        None
    }

    /// Returns a `InstrumentConfigurationRef` with the specified `id` if a `InstrumentConfiguration` with the
    /// corresponding id is present
    pub fn instrument_configuration_ref(&self, id: &str) -> Option<InstrumentConfigurationRef> {
        for instrument_configuration in &self.instrument_configuration_list {
            if instrument_configuration.id() == id {
                return Some(InstrumentConfigurationRef::Ref(
                    instrument_configuration.clone(),
                ));
            }
        }

        None
    }

    /// Returns a `DataProcessingRef` with the specified `id` if a `DataProcessing` with the
    /// corresponding id is present
    pub fn data_processing_ref(&self, id: &str) -> Option<DataProcessingRef> {
        for data_processing in &self.data_processing_list {
            if data_processing.id() == id {
                return Some(DataProcessingRef::Ref(data_processing.clone()));
            }
        }

        None
    }

    /// Returns a `SourceFileRef` with the specified `id` if a `SourceFile` with the
    /// corresponding id is present
    pub fn source_file_ref(&self, id: &str) -> Option<SourceFileRef> {
        for source_file in &self.file_description.as_ref().unwrap().source_file_list {
            if source_file.id() == id {
                return Some(SourceFileRef::Ref(source_file.clone()));
            }
        }

        None
    }

    /// Returns a `SpectrumRef` with the specified `id` if a `Spectrum` with the
    /// corresponding id is present
    pub fn spectrum_ref(&self, id: &str) -> Option<SpectrumRef> {
        match &self.run.as_ref().unwrap().spectrum_list {
            Some(spectrum_list) => {
                for spectrum in spectrum_list {
                    if spectrum.id() == id {
                        return Some(SpectrumRef::Ref(spectrum.clone()));
                    }
                }

                None
            }
            None => None,
        }
    }

    /// Returns a `SampleRef` with the specified `id` if a `Sample` with the
    /// corresponding id is present
    pub fn sample_ref(&self, id: &str) -> Option<SampleRef> {
        for sample in &self.sample_list {
            if sample.id() == id {
                return Some(SampleRef::Ref(sample.clone()));
            }
        }

        None
    }

    /// Returns a `SoftwareRef` with the specified `id` if a `Software` with the
    /// corresponding id is present
    pub fn software_ref(&self, id: &str) -> Option<SoftwareRef> {
        for software in &self.software_list {
            if software.id() == id {
                return Some(SoftwareRef::Ref(software.clone()));
            }
        }

        None
    }

    /// Returns the number of spectra in the file
    pub fn num_spectra(&self) -> usize {
        match self.run.as_ref() {
            Some(run) => run.num_spectra(),
            None => 0,
        }
    }

    /// Returns the spectrum at the specified index
    pub fn spectrum(&self, index: usize) -> Option<&Arc<Spectrum>> {
        match self.run.as_ref() {
            Some(run) => run.spectrum(index),
            None => None,
        }
    }

    /// Returns the `Run` associated with the mzML tag
    pub fn run(&self) -> Option<&Run> {
        self.run.as_ref()
    }

    /// Returns a mutable reference to `Run` tag
    pub fn run_mut(&mut self) -> Option<&mut Run> {
        self.run.as_mut()
    }

    /// Returns the `SpectrumList` in the `Run`
    pub fn spectrum_list(&self) -> Option<&SpectrumList> {
        self.run().unwrap().spectrum_list()
    }

    /// Returns mutable access the `SpectrumList` in the `Run`
    pub fn spectrum_list_mut(&mut self) -> Option<&mut SpectrumList> {
        self.run.as_mut().unwrap().spectrum_list_mut()
    }
}

impl MzMLTag for MzML {
    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().as_ref() != b"mzML" {
            return Err(FatalParseError::UnexpectedTag(format!(
                "Unexpected event {:?} when processing MzML",
                start_event,
            )));
        }

        let mut _accession = None;
        let mut _id = None;
        let mut version = None;

        for att in start_event
            .attributes()
            .with_checks(parser.with_attribute_checks)
        {
            match att {
                Ok(att) => match att.key.as_ref() {
                    b"accession" => _accession = Some(att.value),
                    b"id" => _id = Some(att.value),
                    b"version" => version = Some(att.value),
                    b"xmlns" => {}
                    b"xmlns:xsi" => {}
                    b"xsi:schemaLocation" => {}
                    _ => {
                        parser.errors.push_back(ParseError::UnexpectedAttribute((
                            Tag::MzML,
                            std::str::from_utf8(att.key.as_ref())?.to_string(),
                        )));
                    }
                },
                Err(error) => {
                    parser
                        .errors
                        .push_back(ParseError::XMLError((Tag::MzML, error.into())));
                }
            };
        }

        // TODO: mzmlBuilder pattern, to ensure that all necessary features are present?
        let mzml = match version {
            Some(version) => MzML::new(std::str::from_utf8(&version)?, parser.ontology.clone()),
            None => {
                parser.errors.push_back(ParseError::MissingAttribute((
                    Tag::MzML,
                    "version".to_string(),
                )));
                MzML::new("", parser.ontology.clone())
            }
        };

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

        Ok(Some(mzml))
    }

    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 current_event = parser.next(buffer)?;

            match current_event {
                Event::Start(start_event) | Event::Empty(start_event) => {
                    match start_event.name().as_ref() {
                        b"cvList" => {
                            if let Some(mut cv_list) =
                                CVList::parse_start_tag(parser, &start_event)?
                            {
                                cv_list.parse_xml(parser, buffer)?;
                                self.cv_list = cv_list.cvs;
                            }
                        }
                        b"fileDescription" => {
                            if let Some(mut file_description) =
                                FileDescription::parse_start_tag(parser, &start_event)?
                            {
                                file_description.parse_xml(parser, buffer)?;
                                self.file_description = Some(file_description);
                            }
                        }
                        b"referenceableParamGroupList" => {
                            if let Some(mut param_group_list) =
                                ReferenceableParamGroupList::parse_start_tag(parser, &start_event)?
                            {
                                param_group_list.parse_xml(parser, buffer)?;
                                self.referenceable_param_group_list = param_group_list.list;

                                // Copy the list to the parser
                                parser.referenceable_param_group_list =
                                    self.referenceable_param_group_list.clone();
                            }
                        }
                        b"sampleList" => {
                            if let Some(mut sample_list) =
                                SampleList::parse_start_tag(parser, &start_event)?
                            {
                                sample_list.parse_xml(parser, buffer)?;
                                self.sample_list = sample_list.list;

                                // Copy the list to the parser
                                parser.sample_list = self.sample_list.clone();
                            }
                        }
                        b"softwareList" => {
                            if let Some(mut software_list) =
                                SoftwareList::parse_start_tag(parser, &start_event)?
                            {
                                software_list.parse_xml(parser, buffer)?;
                                self.software_list = software_list.list;

                                // Copy the list to the parser
                                parser.software_list = self.software_list.clone();
                            }
                        }
                        b"scanSettingsList" => {
                            if let Some(mut scan_settings_list) =
                                ScanSettingsList::parse_start_tag(parser, &start_event)?
                            {
                                scan_settings_list.parse_xml(parser, buffer)?;
                                self.scan_settings_list = scan_settings_list;

                                // Copy the list to the parser
                                parser.scan_settings_list = self.scan_settings_list.clone();
                            }
                        }
                        b"instrumentConfigurationList" => {
                            if let Some(mut instrument_list) =
                                InstrumentConfigurationList::parse_start_tag(parser, &start_event)?
                            {
                                instrument_list.parse_xml(parser, buffer)?;
                                self.instrument_configuration_list = instrument_list.list;

                                // Copy the list to the parser
                                parser.instrument_list = self.instrument_configuration_list.clone();
                            }
                        }
                        b"dataProcessingList" => {
                            if let Some(mut data_processing_list) =
                                DataProcessingList::parse_start_tag(parser, &start_event)?
                            {
                                data_processing_list.parse_xml(parser, buffer)?;
                                self.data_processing_list = data_processing_list.list;

                                // Copy the list to the parser
                                parser.data_processing_list = self.data_processing_list.clone();
                            }
                        }
                        b"run" => {
                            if let Some(mut run) = Run::parse_start_tag(parser, &start_event)? {
                                run.parse_xml(parser, buffer)?;
                                self.run = Some(run);
                            }
                        }
                        _ => parser.errors.push_back(ParseError::UnexpectedTag(format!(
                            "{:?} unexpected when processing {:?}",
                            std::str::from_utf8(start_event.name().as_ref()),
                            Tag::MzML
                        ))),
                    }
                }
                Event::End(end_event) => {
                    if let b"mzML" = end_event.name().as_ref() {
                        // Check that all required entries are present
                        if self.file_description.is_none() {
                            return Err(FatalParseError::MissingOpeningTag(
                                "fileDescription".to_string(),
                            ));
                        }
                        if self.run.is_none() {
                            return Err(FatalParseError::MissingOpeningTag("run".to_string()));
                        }

                        parser.breadcrumbs.pop_back();

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

        Ok(())
    }

    fn write_xml<W: Write>(&self, writer: &mut Writer<W>) -> Result<(), quick_xml::Error> {
        let mut elem = BytesStart::new("mzML");
        elem.push_attribute(("xmlns", XMLNS));
        elem.push_attribute(("xmlns:xsi", XMLNS_XSI));
        elem.push_attribute(("xsi:schemaLocation", XSI_SCHEMA_LOCATION));
        elem.push_attribute(("version", MZML_VERSION));

        writer.write_event(Event::Start(elem))?;

        // cvList
        if !self.cv_list.is_empty() {
            writer.write_list("cvList", &self.cv_list)?;
        }

        // fileDescription
        if let Some(ref file_description) = self.file_description {
            file_description.write_xml(writer)?;
        }

        // referenceableParamGroupList
        if !self.referenceable_param_group_list.is_empty() {
            writer.write_arc_list(
                "referenceableParamGroupList",
                &self.referenceable_param_group_list,
            )?;
        }

        // sampleList
        if !self.sample_list.is_empty() {
            writer.write_arc_list("sampleList", &self.sample_list)?;
        }

        // softwareList
        if !self.software_list.is_empty() {
            writer.write_arc_list("softwareList", &self.software_list)?;
        }

        // scanSettingsList
        if !self.scan_settings_list.is_empty() {
            self.scan_settings_list.write_xml(writer)?;
        }

        // instrumentConfigurationList
        if !self.instrument_configuration_list.is_empty() {
            writer.write_arc_list(
                "instrumentConfigurationList",
                &self.instrument_configuration_list,
            )?;
        }

        // dataProcessingList
        if !self.data_processing_list.is_empty() {
            writer.write_arc_list("dataProcessingList", &self.data_processing_list)?;
        }

        self.run.as_ref().unwrap().write_xml(writer)?;

        writer.end_tag("mzML")
    }

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