use std::sync::Arc;

use arrow::array::{BooleanArray, Float64Array, TimestampNanosecondArray, UInt32Array, UInt8Array};
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use arrow::record_batch::RecordBatch;
use bt_data::{DataStore, DefaultPartitionScheme, ParquetDataStore, PartitionScheme, WriteOptions};
use bt_kernel::{BtError, SymbolId, TimeRange, Timestamp, NANOS_PER_SECOND};
use tempfile::TempDir;

fn make_store() -> (TempDir, ParquetDataStore) {
    let tmp = TempDir::new().expect("create temp dir");
    let data_root = tmp.path().join("data");
    let metadata_path = tmp.path().join("metadata.sqlite");

    let store = ParquetDataStore::new(&data_root, &metadata_path).expect("create parquet store");
    (tmp, store)
}

fn bars_batch(symbol: SymbolId, timestamps: &[i64]) -> RecordBatch {
    let len = timestamps.len();

    let open = (0..len).map(|v| v as f64 + 1.0).collect::<Vec<_>>();
    let high = open.iter().map(|v| v + 1.0).collect::<Vec<_>>();
    let low = open.iter().map(|v| v - 1.0).collect::<Vec<_>>();
    let close = open.iter().map(|v| v + 0.5).collect::<Vec<_>>();
    let vwap = close.clone();
    let volume = vec![10.0; len];
    let buy_volume = vec![Some(4.0); len];
    let sell_volume = vec![Some(6.0); len];
    let trade_count = vec![3_u32; len];
    let bid = vec![Some(100.0); len];
    let ask = vec![Some(101.0); len];
    let spread = vec![Some(1.0); len];
    let is_gap = vec![false; len];
    let gap_fill_method = vec![0_u8; len];

    let schema = bt_data::schema::canonical_bars_1m_schema();
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(TimestampNanosecondArray::from(timestamps.to_vec()).with_timezone("UTC")),
            Arc::new(UInt32Array::from(vec![symbol.0; len])),
            Arc::new(Float64Array::from(open)),
            Arc::new(Float64Array::from(high)),
            Arc::new(Float64Array::from(low)),
            Arc::new(Float64Array::from(close)),
            Arc::new(Float64Array::from(vwap)),
            Arc::new(Float64Array::from(volume)),
            Arc::new(Float64Array::from(buy_volume)),
            Arc::new(Float64Array::from(sell_volume)),
            Arc::new(UInt32Array::from(trade_count)),
            Arc::new(Float64Array::from(bid)),
            Arc::new(Float64Array::from(ask)),
            Arc::new(Float64Array::from(spread)),
            Arc::new(BooleanArray::from(is_gap)),
            Arc::new(UInt8Array::from(gap_fill_method)),
        ],
    )
    .expect("build bars batch")
}

#[test]
fn write_then_load_roundtrip_works() {
    let (_tmp, store) = make_store();
    let symbol = SymbolId(1);

    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let batch = bars_batch(symbol, &[base, base + NANOS_PER_SECOND]);

    let version = store
        .write_bars(symbol, &batch, WriteOptions::default())
        .expect("write bars");

    let range = TimeRange::new(Timestamp(base), Timestamp(base + 2 * NANOS_PER_SECOND))
        .expect("valid range");
    let loaded = store
        .load_bars(symbol, range, Some(&version.id))
        .expect("load bars");

    assert_eq!(loaded.num_rows(), 2);

    let active = store.active_version("bars_1m").expect("active version");
    assert_eq!(active.id, version.id);
}

#[test]
fn invalid_schema_is_rejected() {
    let (_tmp, store) = make_store();

    let schema = Arc::new(Schema::new(vec![Field::new(
        "timestamp",
        DataType::Timestamp(TimeUnit::Nanosecond, None),
        false,
    )]));

    let batch = RecordBatch::try_new(
        schema,
        vec![Arc::new(TimestampNanosecondArray::from(vec![
            1_704_067_200_i64 * NANOS_PER_SECOND,
        ]))],
    )
    .expect("invalid schema batch");

    let err = store
        .write_bars(SymbolId(1), &batch, WriteOptions::default())
        .expect_err("expected schema mismatch");

    assert!(matches!(err, BtError::SchemaMismatch { .. }));
}

#[test]
fn non_aligned_timestamp_is_rejected() {
    let (_tmp, store) = make_store();

    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let batch = bars_batch(SymbolId(1), &[base + 123]);

    let err = store
        .write_bars(SymbolId(1), &batch, WriteOptions::default())
        .expect_err("expected timestamp alignment error");

    assert!(matches!(err, BtError::Data(_)));
}

#[test]
fn partitioning_handles_midnight_boundary() {
    let (tmp, store) = make_store();
    let symbol = SymbolId(7);

    let ts_before_midnight = 1_704_153_599_i64 * NANOS_PER_SECOND;
    let ts_at_midnight = 1_704_153_600_i64 * NANOS_PER_SECOND;

    let batch = bars_batch(symbol, &[ts_before_midnight, ts_at_midnight]);
    store
        .write_bars(symbol, &batch, WriteOptions::default())
        .expect("write bars across midnight");

    let scheme = DefaultPartitionScheme::new(tmp.path().join("data"));
    let first = scheme.partition_path("bars_1m", symbol, Timestamp(ts_before_midnight));
    let second = scheme.partition_path("bars_1m", symbol, Timestamp(ts_at_midnight));

    assert!(first.exists());
    assert!(second.exists());
    assert_ne!(first, second);
}

#[test]
fn malformed_parquet_returns_graceful_error() {
    let (tmp, store) = make_store();
    let symbol = SymbolId(9);

    let ts = 1_704_067_200_i64 * NANOS_PER_SECOND;
    let scheme = DefaultPartitionScheme::new(tmp.path().join("data"));
    let path = scheme.partition_path("bars_1m", symbol, Timestamp(ts));
    std::fs::create_dir_all(path.parent().expect("path has parent")).expect("create dirs");
    std::fs::write(&path, b"not a parquet file").expect("write malformed parquet");

    let range = TimeRange::new(Timestamp(ts), Timestamp(ts + NANOS_PER_SECOND)).expect("range");
    let err = store
        .load_bars(symbol, range, None)
        .expect_err("expected malformed parquet error");

    assert!(matches!(err, BtError::Parquet(_) | BtError::Arrow(_)));
}

#[test]
fn latest_active_version_is_returned() {
    let (_tmp, store) = make_store();
    let symbol = SymbolId(11);
    let base = 1_704_067_200_i64 * NANOS_PER_SECOND;

    let first = bars_batch(symbol, &[base]);
    let second = bars_batch(symbol, &[base + NANOS_PER_SECOND]);

    let first_v = store
        .write_bars(symbol, &first, WriteOptions::default())
        .expect("write first");
    let second_v = store
        .write_bars(symbol, &second, WriteOptions::default())
        .expect("write second");

    let listed = store.list_versions("bars_1m").expect("list versions");
    assert_eq!(listed.len(), 2);

    let active = store.active_version("bars_1m").expect("active version");
    assert_ne!(active.id, first_v.id);
    assert_eq!(active.id, second_v.id);
}
