use std::sync::Arc;

use arrow::array::{Float64Array, StringArray, TimestampNanosecondArray, UInt8Array};
use arrow::record_batch::RecordBatch;
use bt_connectors::raw_trades_schema;
use bt_data::{DataStore, ParquetDataStore, WriteOptions};
use bt_etl::aggregator::aggregate_raw_trades_to_1s_bars;
use bt_etl::dedup::dedup_and_sort_raw_trades;
use bt_etl::gaps::fill_gaps_1s;
use bt_etl::ingest_raw_trades_to_store;
use bt_kernel::{SymbolId, TimeRange, Timestamp, NANOS_PER_SECOND};
use tempfile::TempDir;

fn raw_trade_batch(
    symbol_id: SymbolId,
    timestamps: &[i64],
    prices: &[f64],
    quantities: &[f64],
    sides: &[u8],
    trade_ids: &[&str],
) -> RecordBatch {
    let row_count = timestamps.len();
    assert_eq!(prices.len(), row_count);
    assert_eq!(quantities.len(), row_count);
    assert_eq!(sides.len(), row_count);
    assert_eq!(trade_ids.len(), row_count);

    RecordBatch::try_new(
        raw_trades_schema(),
        vec![
            Arc::new(TimestampNanosecondArray::from(timestamps.to_vec()).with_timezone("UTC")),
            Arc::new(TimestampNanosecondArray::from(timestamps.to_vec()).with_timezone("UTC")),
            Arc::new(StringArray::from(vec!["binance"; row_count])),
            Arc::new(StringArray::from(vec!["BTCUSDT"; row_count])),
            Arc::new(arrow::array::UInt32Array::from(vec![
                symbol_id.0;
                row_count
            ])),
            Arc::new(Float64Array::from(prices.to_vec())),
            Arc::new(Float64Array::from(quantities.to_vec())),
            Arc::new(UInt8Array::from(sides.to_vec())),
            Arc::new(StringArray::from(trade_ids.to_vec())),
            Arc::new(StringArray::from(vec![Option::<&str>::None; row_count])),
        ],
    )
    .expect("build raw trade batch")
}

#[test]
fn dedup_and_sort_removes_duplicate_trade_ids() {
    let symbol = SymbolId(7);
    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let raw = raw_trade_batch(
        symbol,
        &[base + NANOS_PER_SECOND, base, base + 2 * NANOS_PER_SECOND],
        &[102.0, 100.0, 101.0],
        &[1.0, 1.0, 1.0],
        &[1, 1, 1],
        &["t2", "t1", "t1"],
    );

    let deduped = dedup_and_sort_raw_trades(&raw).expect("dedup raw trades");
    assert_eq!(deduped.num_rows(), 2);

    let timestamps = deduped
        .column_by_name("timestamp")
        .expect("timestamp column")
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .expect("timestamp array");
    assert_eq!(timestamps.value(0), base);
    assert_eq!(timestamps.value(1), base + NANOS_PER_SECOND);
}

#[test]
fn aggregate_raw_trades_to_1s_bars_computes_ohlcv() {
    let symbol = SymbolId(9);
    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let raw = raw_trade_batch(
        symbol,
        &[base, base + 100_000_000, base + NANOS_PER_SECOND],
        &[100.0, 105.0, 110.0],
        &[1.0, 2.0, 3.0],
        &[1, 2, 1],
        &["a", "b", "c"],
    );

    let deduped = dedup_and_sort_raw_trades(&raw).expect("dedup");
    let bars = aggregate_raw_trades_to_1s_bars(&deduped, symbol).expect("aggregate bars");
    assert_eq!(bars.num_rows(), 2);

    let open = bars
        .column_by_name("open")
        .expect("open")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("open array");
    let close = bars
        .column_by_name("close")
        .expect("close")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("close array");
    let vwap = bars
        .column_by_name("vwap")
        .expect("vwap")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("vwap array");

    assert_eq!(open.value(0), 100.0);
    assert_eq!(close.value(0), 105.0);
    assert!((vwap.value(0) - 103.333_333_333_3).abs() < 1e-9);
}

#[test]
fn fill_gaps_inserts_missing_second_with_ffill() {
    let symbol = SymbolId(12);
    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let raw = raw_trade_batch(
        symbol,
        &[base, base + 2 * NANOS_PER_SECOND],
        &[100.0, 105.0],
        &[1.0, 1.0],
        &[1, 1],
        &["x", "y"],
    );

    let bars = aggregate_raw_trades_to_1s_bars(&raw, symbol).expect("aggregate");
    let gap_filled = fill_gaps_1s(&bars, None).expect("fill gaps");
    assert_eq!(gap_filled.batch.num_rows(), 3);
    assert_eq!(gap_filled.gap_bar_count, 1);

    let is_gap = gap_filled
        .batch
        .column_by_name("is_gap")
        .expect("is_gap")
        .as_any()
        .downcast_ref::<arrow::array::BooleanArray>()
        .expect("bool is_gap");
    assert!(!is_gap.value(0));
    assert!(is_gap.value(1));
    assert!(!is_gap.value(2));

    let close = gap_filled
        .batch
        .column_by_name("close")
        .expect("close")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("close array");
    assert_eq!(close.value(1), close.value(0));
}

#[test]
fn ingest_pipeline_writes_dataset_version() {
    let symbol = SymbolId(21);
    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let raw = raw_trade_batch(
        symbol,
        &[base, base + 2 * NANOS_PER_SECOND],
        &[200.0, 205.0],
        &[1.0, 1.0],
        &[1, 2],
        &["a1", "a2"],
    );

    let tmp = TempDir::new().expect("create tempdir");
    let store = ParquetDataStore::new(tmp.path().join("data"), tmp.path().join("meta.sqlite"))
        .expect("create parquet datastore");
    let range = TimeRange::new(Timestamp(base), Timestamp(base + 3 * NANOS_PER_SECOND))
        .expect("valid time range");

    let result =
        ingest_raw_trades_to_store(&store, &raw, symbol, Some(range), WriteOptions::default())
            .expect("run etl ingest pipeline");

    assert_eq!(result.version.dataset, "bars_1m");
    assert!(!result.quality_reports.is_empty());
    assert_eq!(result.bars.num_rows(), 3);

    let loaded = store
        .load_bars(symbol, range, Some(result.version.id.as_str()))
        .expect("load stored bars");
    assert_eq!(loaded.num_rows(), 3);
}
