use std::collections::BTreeMap;
use std::sync::Arc;

use arrow::array::{
    Array, ArrayRef, BooleanArray, Float64Array, TimestampNanosecondArray, UInt32Array, UInt8Array,
};
use arrow::record_batch::RecordBatch;
use bt_data::schema::canonical_bars_1m_schema;
use bt_kernel::{BtError, Result, SymbolId, NANOS_PER_SECOND};

#[tracing::instrument(level = "debug", skip_all)]
pub fn aggregate_raw_trades_to_1s_bars(
    raw_trades: &RecordBatch,
    symbol: SymbolId,
) -> Result<RecordBatch> {
    if raw_trades.num_rows() == 0 {
        return Ok(RecordBatch::new_empty(canonical_bars_1m_schema()));
    }

    let timestamps = timestamp_column(raw_trades, "timestamp")?;
    let prices = float64_column(raw_trades, "price")?;
    let quantities = float64_column(raw_trades, "quantity")?;
    let sides = u8_column(raw_trades, "side")?;

    let mut points = Vec::with_capacity(raw_trades.num_rows());
    for idx in 0..raw_trades.num_rows() {
        ensure_not_null(timestamps, "timestamp", idx)?;
        ensure_not_null(prices, "price", idx)?;
        ensure_not_null(quantities, "quantity", idx)?;
        ensure_not_null(sides, "side", idx)?;
        points.push(TradePoint {
            timestamp_ns: timestamps.value(idx),
            price: prices.value(idx),
            quantity: quantities.value(idx),
            side: sides.value(idx),
        });
    }

    points.sort_by(|a, b| a.timestamp_ns.cmp(&b.timestamp_ns));

    let mut by_second: BTreeMap<i64, BarAccumulator> = BTreeMap::new();
    for point in points {
        let second = point.timestamp_ns.div_euclid(NANOS_PER_SECOND) * NANOS_PER_SECOND;
        by_second
            .entry(second)
            .or_insert_with(BarAccumulator::new)
            .update(point.price, point.quantity, point.side);
    }

    let row_count = by_second.len();
    let mut out_timestamps = Vec::with_capacity(row_count);
    let mut out_symbol_id = Vec::with_capacity(row_count);
    let mut out_open = Vec::with_capacity(row_count);
    let mut out_high = Vec::with_capacity(row_count);
    let mut out_low = Vec::with_capacity(row_count);
    let mut out_close = Vec::with_capacity(row_count);
    let mut out_vwap = Vec::with_capacity(row_count);
    let mut out_volume = Vec::with_capacity(row_count);
    let mut out_buy_volume = Vec::with_capacity(row_count);
    let mut out_sell_volume = Vec::with_capacity(row_count);
    let mut out_trade_count = Vec::with_capacity(row_count);
    let mut out_bid: Vec<Option<f64>> = Vec::with_capacity(row_count);
    let mut out_ask: Vec<Option<f64>> = Vec::with_capacity(row_count);
    let mut out_spread: Vec<Option<f64>> = Vec::with_capacity(row_count);
    let mut out_is_gap = Vec::with_capacity(row_count);
    let mut out_gap_fill_method = Vec::with_capacity(row_count);

    for (second, acc) in by_second {
        out_timestamps.push(second);
        out_symbol_id.push(symbol.0);
        out_open.push(acc.open);
        out_high.push(acc.high);
        out_low.push(acc.low);
        out_close.push(acc.close);
        out_volume.push(acc.volume);
        out_vwap.push(if acc.volume > 0.0 {
            acc.weighted_sum / acc.volume
        } else {
            acc.close
        });
        out_buy_volume.push(Some(acc.buy_volume));
        out_sell_volume.push(Some(acc.sell_volume));
        out_trade_count.push(acc.trade_count);
        out_bid.push(None);
        out_ask.push(None);
        out_spread.push(None);
        out_is_gap.push(false);
        out_gap_fill_method.push(0_u8);
    }

    let arrays: Vec<ArrayRef> = vec![
        Arc::new(TimestampNanosecondArray::from(out_timestamps).with_timezone("UTC")),
        Arc::new(UInt32Array::from(out_symbol_id)),
        Arc::new(Float64Array::from(out_open)),
        Arc::new(Float64Array::from(out_high)),
        Arc::new(Float64Array::from(out_low)),
        Arc::new(Float64Array::from(out_close)),
        Arc::new(Float64Array::from(out_vwap)),
        Arc::new(Float64Array::from(out_volume)),
        Arc::new(Float64Array::from(out_buy_volume)),
        Arc::new(Float64Array::from(out_sell_volume)),
        Arc::new(UInt32Array::from(out_trade_count)),
        Arc::new(Float64Array::from(out_bid)),
        Arc::new(Float64Array::from(out_ask)),
        Arc::new(Float64Array::from(out_spread)),
        Arc::new(BooleanArray::from(out_is_gap)),
        Arc::new(UInt8Array::from(out_gap_fill_method)),
    ];

    let result = RecordBatch::try_new(canonical_bars_1m_schema(), arrays)?;
    tracing::debug!(
        output_rows = result.num_rows(),
        "aggregated trades to 1s bars"
    );
    Ok(result)
}

#[derive(Copy, Clone, Debug)]
struct TradePoint {
    timestamp_ns: i64,
    price: f64,
    quantity: f64,
    side: u8,
}

#[derive(Clone, Debug)]
struct BarAccumulator {
    open: f64,
    high: f64,
    low: f64,
    close: f64,
    volume: f64,
    weighted_sum: f64,
    buy_volume: f64,
    sell_volume: f64,
    trade_count: u32,
    initialized: bool,
}

impl BarAccumulator {
    fn new() -> Self {
        Self {
            open: 0.0,
            high: f64::MIN,
            low: f64::MAX,
            close: 0.0,
            volume: 0.0,
            weighted_sum: 0.0,
            buy_volume: 0.0,
            sell_volume: 0.0,
            trade_count: 0,
            initialized: false,
        }
    }

    fn update(&mut self, price: f64, quantity: f64, side: u8) {
        if !self.initialized {
            self.open = price;
            self.close = price;
            self.high = price;
            self.low = price;
            self.initialized = true;
        } else {
            self.high = self.high.max(price);
            self.low = self.low.min(price);
            self.close = price;
        }

        self.volume += quantity;
        self.weighted_sum += price * quantity;
        self.trade_count = self.trade_count.saturating_add(1);

        match side {
            1 => self.buy_volume += quantity,
            2 => self.sell_volume += quantity,
            _ => {}
        }
    }
}

fn timestamp_column<'a>(
    batch: &'a RecordBatch,
    name: &str,
) -> Result<&'a TimestampNanosecondArray> {
    let column = batch
        .column_by_name(name)
        .ok_or_else(|| BtError::Data(format!("missing raw trade column '{name}'")))?;
    column
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be Timestamp(ns)")))
}

fn float64_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a Float64Array> {
    let column = batch
        .column_by_name(name)
        .ok_or_else(|| BtError::Data(format!("missing raw trade column '{name}'")))?;
    column
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be Float64")))
}

fn u8_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a UInt8Array> {
    let column = batch
        .column_by_name(name)
        .ok_or_else(|| BtError::Data(format!("missing raw trade column '{name}'")))?;
    column
        .as_any()
        .downcast_ref::<UInt8Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be UInt8")))
}

fn ensure_not_null(array: &dyn Array, name: &str, idx: usize) -> Result<()> {
    if array.is_null(idx) {
        return Err(BtError::Data(format!(
            "column '{name}' contains null at row {idx}"
        )));
    }
    Ok(())
}
