use std::collections::HashMap;
use super::SignalValue;

/// Columnar signal storage: maps signal IDs to vectors of values per timestep.
#[derive(Clone, Debug)]
pub struct SignalStore {
    data: HashMap<usize, Vec<SignalValue>>,
}

impl SignalStore {
    pub fn new() -> Self {
        Self {
            data: HashMap::new(),
        }
    }

    pub fn get(&self, signal_id: usize, index: usize) -> SignalValue {
        self.data
            .get(&signal_id)
            .and_then(|v| v.get(index).cloned())
            .unwrap_or(SignalValue::Unknown)
    }

    pub fn set(&mut self, signal_id: usize, index: usize, value: SignalValue) {
        let col = self.data.entry(signal_id).or_default();
        if index >= col.len() {
            col.resize(index + 1, SignalValue::Unknown);
        }
        col[index] = value;
    }

    pub fn ensure_length(&mut self, signal_id: usize, len: usize) {
        let col = self.data.entry(signal_id).or_default();
        if col.len() < len {
            col.resize(len, SignalValue::Unknown);
        }
    }

    pub fn get_range(&self, signal_id: usize, start: usize, end: usize) -> Vec<SignalValue> {
        self.data
            .get(&signal_id)
            .map(|v| {
                (start..end)
                    .map(|i| v.get(i).cloned().unwrap_or(SignalValue::Unknown))
                    .collect()
            })
            .unwrap_or_else(|| vec![SignalValue::Unknown; end - start])
    }

    pub fn signal_ids(&self) -> Vec<usize> {
        self.data.keys().copied().collect()
    }

    /// Forward-fill all signals from index-1 to index.
    /// Only copies values for signals that have data at index-1
    /// but not yet at index. O(num_signals) per call.
    pub fn forward_fill(&mut self, index: usize) {
        if index == 0 {
            return;
        }
        let prev = index - 1;
        // Collect (id, value) pairs to avoid borrow conflict
        let fills: Vec<(usize, SignalValue)> = self.data.iter()
            .filter_map(|(&id, col)| {
                if let Some(val) = col.get(prev) {
                    if *val != SignalValue::Unknown {
                        // Only fill if current index doesn't have a value yet
                        let needs_fill = col.get(index)
                            .map(|v| *v == SignalValue::Unknown)
                            .unwrap_or(true);
                        if needs_fill {
                            return Some((id, val.clone()));
                        }
                    }
                }
                None
            })
            .collect();
        for (id, val) in fills {
            self.set(id, index, val);
        }
    }
}

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