use arrow::record_batch::RecordBatch;
use bt_data::{DataStore, DatasetVersion, WriteOptions};
use bt_kernel::{BtError, Result, SymbolId, TimeRange, Timestamp};

pub struct InMemoryStore {
    bars: RecordBatch,
    version_id: String,
}

impl InMemoryStore {
    pub fn new(bars: RecordBatch, version_id: &str) -> Self {
        Self {
            bars,
            version_id: version_id.to_string(),
        }
    }

    fn version(&self) -> DatasetVersion {
        DatasetVersion {
            id: self.version_id.clone(),
            dataset: "bars_1m".to_string(),
            created_at: Timestamp::from_nanos(1),
            schema_hash: "memory".to_string(),
            row_count: self.bars.num_rows() as u64,
            time_range: TimeRange::new(Timestamp::from_nanos(0), Timestamp::from_nanos(i64::MAX))
                .expect("range"),
            symbols: vec![SymbolId(1)],
            source_hash: "memory".to_string(),
            metadata: serde_json::json!({}),
        }
    }
}

impl DataStore for InMemoryStore {
    fn load_bars(
        &self,
        _symbol: SymbolId,
        _range: TimeRange,
        _version: Option<&str>,
    ) -> Result<RecordBatch> {
        Ok(self.bars.clone())
    }

    fn load_features(
        &self,
        _feature_set: &str,
        _symbol: SymbolId,
        _range: TimeRange,
        _columns: Option<&[&str]>,
    ) -> Result<RecordBatch> {
        Err(BtError::Unsupported("in-memory store".to_string()))
    }

    fn load_exo(&self, _name: &str, _range: TimeRange) -> Result<RecordBatch> {
        Err(BtError::Unsupported("in-memory store".to_string()))
    }

    fn list_versions(&self, _dataset: &str) -> Result<Vec<DatasetVersion>> {
        Ok(vec![self.version()])
    }

    fn active_version(&self, _dataset: &str) -> Result<DatasetVersion> {
        Ok(self.version())
    }

    fn write_bars(
        &self,
        _symbol: SymbolId,
        _batch: &RecordBatch,
        _opts: WriteOptions,
    ) -> Result<DatasetVersion> {
        Err(BtError::Unsupported("in-memory store".to_string()))
    }
}
