pub mod aggregator;
pub mod dedup;
pub mod gaps;
pub mod quality;

use std::sync::Arc;

use arrow::array::UInt32Array;
use arrow::record_batch::RecordBatch;
use bt_data::{DataStore, DatasetVersion, WriteOptions};
use bt_kernel::{Result, SymbolId, TimeRange};
use quality::{generate_quality_reports, DataQualityReport};

use crate::aggregator::aggregate_raw_trades_to_1s_bars;
use crate::dedup::dedup_and_sort_raw_trades;
use crate::gaps::fill_gaps_1s;

#[derive(Debug)]
pub struct EtlIngestResult {
    pub deduped_raw_trades: RecordBatch,
    pub bars: RecordBatch,
    pub version: DatasetVersion,
    pub quality_reports: Vec<DataQualityReport>,
    pub warnings: Vec<String>,
}

pub fn transform_raw_trades(
    raw_trades: &RecordBatch,
    symbol: SymbolId,
    range: Option<TimeRange>,
) -> Result<(
    RecordBatch,
    RecordBatch,
    Vec<DataQualityReport>,
    Vec<String>,
)> {
    let deduped = dedup_and_sort_raw_trades(raw_trades)?;
    let bars = aggregate_raw_trades_to_1s_bars(&deduped, symbol)?;
    let gap_result = fill_gaps_1s(&bars, range)?;
    let reports = generate_quality_reports(&gap_result.batch)?;

    Ok((deduped, gap_result.batch, reports, gap_result.warnings))
}

/// Ingest pre-built bars (e.g. klines from exchange API) directly into the store.
///
/// Skips trade dedup/aggregation and gap-filling — klines are already complete bars.
/// Only replaces the `symbol_id` column (providers set it to 0) and runs quality reports.
#[tracing::instrument(level = "info", skip_all, fields(symbol_id = symbol.0))]
pub fn ingest_bars_to_store(
    store: &dyn DataStore,
    bars: &RecordBatch,
    symbol: SymbolId,
    write_opts: WriteOptions,
) -> Result<EtlIngestResult> {
    let bars = replace_symbol_id(bars, symbol)?;
    let quality_reports = generate_quality_reports(&bars)?;
    let version = store.write_bars(symbol, &bars, write_opts)?;

    Ok(EtlIngestResult {
        deduped_raw_trades: RecordBatch::new_empty(bars.schema()),
        bars,
        version,
        quality_reports,
        warnings: Vec::new(),
    })
}

/// Replace the `symbol_id` column in a RecordBatch with the given symbol id.
fn replace_symbol_id(batch: &RecordBatch, symbol: SymbolId) -> Result<RecordBatch> {
    let schema = batch.schema();
    let sid_idx = schema
        .index_of("symbol_id")
        .map_err(|_| bt_kernel::BtError::Data("missing symbol_id column".to_string()))?;

    let new_sid = Arc::new(UInt32Array::from(vec![symbol.0; batch.num_rows()]));
    let mut columns: Vec<Arc<dyn arrow::array::Array>> = batch.columns().to_vec();
    columns[sid_idx] = new_sid;

    Ok(RecordBatch::try_new(schema, columns)?)
}

#[tracing::instrument(level = "info", skip_all, fields(symbol_id = symbol.0))]
pub fn ingest_raw_trades_to_store(
    store: &dyn DataStore,
    raw_trades: &RecordBatch,
    symbol: SymbolId,
    range: Option<TimeRange>,
    write_opts: WriteOptions,
) -> Result<EtlIngestResult> {
    let (deduped, bars, quality_reports, warnings) =
        transform_raw_trades(raw_trades, symbol, range)?;
    let version = store.write_bars(symbol, &bars, write_opts)?;

    Ok(EtlIngestResult {
        deduped_raw_trades: deduped,
        bars,
        version,
        quality_reports,
        warnings,
    })
}
