csv-codegen 0.2.3

use crate::predicate::PredicateComparison;
use csv::StringRecord;
use syn::LitStr;

/// Information about a pivot field for a specific record
#[derive(Debug, Clone)]
pub struct PivotField {
    pub key: String,
    pub value: String,
}

/// Iterator that transforms records through pivoting, yielding (record, pivot_info) pairs
pub struct PivotIterator<'a, I> {
    source_iter: I,
    pivot_fields: Vec<(&'a str, usize)>,
    pivot_key_filter: Option<&'a PredicateComparison>,
    current_record: Option<&'a StringRecord>,
    current_pivot_index: usize,
}

impl<'a, I> PivotIterator<'a, I>
where
    I: Iterator<Item = &'a StringRecord>,
{
    pub fn new(
        source_iter: I,
        pivot_fields: Vec<(&'a str, usize)>,
        pivot_key_filter: Option<&'a PredicateComparison>,
    ) -> Self {
        Self {
            source_iter,
            pivot_fields,
            pivot_key_filter,
            current_record: None,
            current_pivot_index: 0,
        }
    }
}

impl<'a, I> Iterator for PivotIterator<'a, I>
where
    I: Iterator<Item = &'a StringRecord>,
{
    type Item = (&'a StringRecord, PivotField);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // If we don't have a current record, get the next one
            if self.current_record.is_none() {
                self.current_record = self.source_iter.next();
                self.current_pivot_index = 0;
            }

            // If we still don't have a record, we're done
            let record = self.current_record?;

            // Try to get the next valid pivot field for this record
            while self.current_pivot_index < self.pivot_fields.len() {
                let (pivot_heading, pivot_index) = self.pivot_fields[self.current_pivot_index];
                self.current_pivot_index += 1;

                // Apply pivot key filter if present
                if let Some(key_comparison) = &self.pivot_key_filter {
                    let pivot_heading_lit =
                        LitStr::new(pivot_heading, key_comparison.field.span()).into();
                    if !key_comparison.evaluate(pivot_heading_lit) {
                        continue;
                    }
                }

                let pivot_value = record.get(pivot_index).expect("pivot field in record");
                let pivot_field = PivotField {
                    key: pivot_heading.to_string(),
                    value: pivot_value.to_string(),
                };

                return Some((record, pivot_field));
            }

            // We've exhausted all pivot fields for this record, move to the next record
            self.current_record = None;
        }
    }
}

/// Type alias for the non-pivoted pass through data
type PassThroughMapper<'a, I> =
    std::iter::Map<I, fn(&'a StringRecord) -> (&'a StringRecord, Option<PivotField>)>;

/// Unified iterator that can handle both pivoted and non-pivoted records
pub enum UnifiedRecordIterator<'a, I> {
    Pivoted(PivotIterator<'a, I>),
    PassThrough(PassThroughMapper<'a, I>),
}

impl<'a, I> Iterator for UnifiedRecordIterator<'a, I>
where
    I: Iterator<Item = &'a StringRecord>,
{
    type Item = (&'a StringRecord, Option<PivotField>);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            UnifiedRecordIterator::Pivoted(iter) => {
                iter.next().map(|(record, pivot)| (record, Some(pivot)))
            }
            UnifiedRecordIterator::PassThrough(iter) => iter.next(),
        }
    }
}

impl<'a, I> UnifiedRecordIterator<'a, I>
where
    I: Iterator<Item = &'a StringRecord>,
{
    pub fn new_pivoted(
        source_iter: I,
        pivot_fields: Vec<(&'a str, usize)>,
        pivot_key_filter: Option<&'a PredicateComparison>,
    ) -> Self {
        Self::Pivoted(PivotIterator::new(
            source_iter,
            pivot_fields,
            pivot_key_filter,
        ))
    }

    pub fn new_pass_through(source_iter: I) -> Self {
        fn pass_through_fn(record: &StringRecord) -> (&StringRecord, Option<PivotField>) {
            (record, None)
        }
        Self::PassThrough(source_iter.map(pass_through_fn))
    }
}