use std::io;

use vfile::util::location::Location;

use crate::chopper::error::{Error, Result};
use crate::chopper::types::{Header, Row, TimestampRange};
use crate::input::dir::location_provider::LocationProvider;
use crate::input::input::InputFormat;
use crate::input::location::factory::ChopperLocationInputFactory;
use crate::source::source::Source;

/// can chain multiple sources together, ensuring they have consistent header/row format;
/// sources must have non-decreasing timestamps
pub struct MultiSource {
    input_factory: ChopperLocationInputFactory,
    loc_provider: Box<dyn LocationProvider>,
    range: TimestampRange,
    input_format: InputFormat,
    common_header: Header,
    swap_map: Vec<usize>, // empty map means pass-through
    current_source: Option<Box<dyn Source>>,
}

impl MultiSource {
    pub fn new(
        input_factory: ChopperLocationInputFactory,
        mut loc_provider: Box<dyn LocationProvider>,
        range: TimestampRange,
        input_format: InputFormat,
        external_common_header: Option<Header>,
    ) -> Result<MultiSource> {
        let (common_header, current_source, swap_map) = match loc_provider.next() {
            None => (Header::new(Vec::new(), Vec::new()), None, Vec::new()),
            Some(first_path) => {
                let first_source = Self::create_source_from_location(
                    &input_factory,
                    &first_path,
                    &range,
                    &input_format,
                )?;

                let (common_header, swap_map) = match external_common_header {
                    None => (first_source.header().clone(), Vec::new()),
                    Some(header) => {
                        let swap_map = Self::check_headers_match_and_return_swap_map(
                            &header,
                            first_source.header(),
                        )?;
                        (header, swap_map)
                    }
                };

                (common_header, Some(first_source), swap_map)
            }
        };

        Ok(MultiSource {
            input_factory,
            loc_provider,
            range,
            input_format,
            common_header,
            current_source,
            swap_map,
        })
    }

    fn update_to_next_source(&mut self) -> Result<()> {
        let next_path = self.loc_provider.next();

        let (next_source, next_swap_map) = match next_path {
            None => (None, Vec::new()),
            Some(path) => {
                let next_source = Self::create_source_from_location(
                    &mut self.input_factory,
                    &path,
                    &self.range,
                    &self.input_format,
                )?;
                let next_swap_map = Self::check_headers_match_and_return_swap_map(
                    &self.common_header,
                    next_source.header(),
                )?;
                (Some(next_source), next_swap_map)
            }
        };

        self.current_source = next_source;
        self.swap_map = next_swap_map;

        Ok(())
    }

    fn create_source_from_location(
        input_factory: &ChopperLocationInputFactory,
        loc: &Location,
        range: &TimestampRange,
        input_format: &InputFormat,
    ) -> Result<Box<dyn Source>> {
        let source = input_factory.create_source(loc, range, input_format, ())?;
        match source {
            None => {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::Other,
                    format!("failed to handle path {}", loc),
                )));
            }
            Some(source) => Ok(source),
        }
    }

    fn check_headers_match_and_return_swap_map(
        ref_header: &Header,
        new_header: &Header,
    ) -> Result<Vec<usize>> {
        if new_header.field_names().len() < ref_header.field_names().len() {
            return Err(Error::from(format!(
                "next file in a multi-file input has less columns than expected; \
                expected columns: {:?}; found columns: {:?}",
                ref_header.field_names(),
                new_header.field_names()
            )));
        }

        let mut new_header = new_header.clone();
        let mut is_identity_map = true;

        let mut swap_map = Vec::with_capacity(ref_header.field_names().len());
        for (ref_idx, ref_field_name) in ref_header.field_names().iter().enumerate() {
            let new_idx_offset = new_header.field_names()[ref_idx..]
                .iter()
                .position(|name| name == ref_field_name);

            match new_idx_offset {
                None => {
                    return Err(Error::from(format!(
                        "next file in a multi-file input is missing at least one expected column; \
                        expected column missing: {:?}",
                        ref_field_name,
                    )));
                }
                Some(new_idx_offset) => {
                    let new_idx = ref_idx + new_idx_offset;

                    let ref_field_type = &ref_header.field_types()[ref_idx];
                    let new_field_type = &new_header.field_types()[new_idx];

                    if ref_field_type != new_field_type {
                        return Err(Error::from(format!(
                            "next file in a multi-file input has unexpected column type for column {:?}; \
                            expected type: {:?}; found type: {:?}",
                            ref_field_name,
                            ref_field_type,
                            new_field_type
                        )));
                    }

                    swap_map.push(new_idx);
                    if ref_idx != new_idx {
                        new_header.field_names_mut().swap(ref_idx, new_idx);
                        new_header.field_types_mut().swap(ref_idx, new_idx);
                        is_identity_map = false;
                    }
                }
            };
        }

        if is_identity_map {
            swap_map.clear();
        }

        Ok(swap_map)
    }
}

impl Source for MultiSource {
    fn header(&self) -> &Header {
        &self.common_header
    }
}

impl Iterator for MultiSource {
    type Item = Result<Box<Row>>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let source = match &mut self.current_source {
                None => return None,
                Some(source) => source,
            };

            let next_row = match source.next() {
                None => {
                    if let Err(e) = self.update_to_next_source() {
                        return Some(Err(e));
                    }
                    continue;
                }
                Some(next_row) => next_row,
            };
            let mut next_row = match next_row {
                Ok(next_row) => next_row,
                Err(e) => return Some(Err(e)),
            };

            if !self.swap_map.is_empty() {
                for (ref_idx, &new_idx) in self.swap_map.iter().enumerate() {
                    next_row.field_values.swap(ref_idx, new_idx);
                }
                next_row.field_values.truncate(self.swap_map.len());
            }
            return Some(Ok(next_row));
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::chopper::types::{FieldType, Header};
    use crate::source::multi_source::MultiSource;

    #[test]
    fn test_check_headers_match_and_return_swap_map() {
        let ref_field_names = vec![
            "f1".to_string(),
            "f2".to_string(),
            "f3".to_string(),
            "f4".to_string(),
            "f5".to_string(),
        ];
        let ref_field_types = vec![
            FieldType::Boolean,
            FieldType::Byte,
            FieldType::String,
            FieldType::Int,
            FieldType::Long,
        ];
        let ref_header = Header::new(ref_field_names, ref_field_types);

        let new_field_names = vec![
            "f1".to_string(),
            "f2".to_string(),
            "f3".to_string(),
            "f4".to_string(),
            "f5".to_string(),
        ];
        let new_field_types = vec![
            FieldType::Boolean,
            FieldType::Byte,
            FieldType::String,
            FieldType::Int,
            FieldType::Long,
        ];
        let new_header = Header::new(new_field_names, new_field_types);
        assert!(
            MultiSource::check_headers_match_and_return_swap_map(&ref_header, &new_header)
                .unwrap()
                .is_empty()
        );

        let new_field_names = vec![
            "f1".to_string(),
            "f2".to_string(),
            "f3".to_string(),
            "f4".to_string(),
            "f6".to_string(),
        ];
        let new_field_types = vec![
            FieldType::Boolean,
            FieldType::Byte,
            FieldType::String,
            FieldType::Int,
            FieldType::Long,
        ];
        let new_header = Header::new(new_field_names, new_field_types);
        assert!(
            MultiSource::check_headers_match_and_return_swap_map(&ref_header, &new_header).is_err()
        );

        let new_field_names = vec![
            "f1".to_string(),
            "f3".to_string(),
            "f2".to_string(),
            "f4".to_string(),
            "f5".to_string(),
        ];
        let new_field_types = vec![
            FieldType::Boolean,
            FieldType::Byte,
            FieldType::String,
            FieldType::Int,
            FieldType::Long,
        ];
        let new_header = Header::new(new_field_names, new_field_types);
        assert!(
            MultiSource::check_headers_match_and_return_swap_map(&ref_header, &new_header).is_err()
        );

        let new_field_names = vec![
            "f1".to_string(),
            "f3".to_string(),
            "f2".to_string(),
            "f4".to_string(),
            "f5".to_string(),
        ];
        let new_field_types = vec![
            FieldType::Boolean,
            FieldType::String,
            FieldType::Byte,
            FieldType::Int,
            FieldType::Long,
        ];
        let new_header = Header::new(new_field_names, new_field_types);
        assert_eq!(
            MultiSource::check_headers_match_and_return_swap_map(&ref_header, &new_header).unwrap(),
            vec![0, 2, 2, 3, 4]
        );

        let new_field_names = vec![
            "f1".to_string(),
            "f3".to_string(),
            "q2".to_string(),
            "q2".to_string(),
            "f2".to_string(),
            "f4".to_string(),
            "f5".to_string(),
            "q4".to_string(),
        ];
        let new_field_types = vec![
            FieldType::Boolean,
            FieldType::String,
            FieldType::Double,
            FieldType::Double,
            FieldType::Byte,
            FieldType::Int,
            FieldType::Long,
            FieldType::Double,
        ];
        let new_header = Header::new(new_field_names, new_field_types);
        assert_eq!(
            MultiSource::check_headers_match_and_return_swap_map(&ref_header, &new_header).unwrap(),
            vec![0, 4, 4, 5, 6]
        );
    }
}
