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, TimeRange, NANOS_PER_SECOND};

#[derive(Debug)]
pub struct GapFillResult {
    pub batch: RecordBatch,
    pub warnings: Vec<String>,
    pub gap_bar_count: usize,
}

#[tracing::instrument(level = "debug", skip_all)]
pub fn fill_gaps_1s(bars: &RecordBatch, range: Option<TimeRange>) -> Result<GapFillResult> {
    if bars.num_rows() == 0 {
        return Ok(GapFillResult {
            batch: RecordBatch::new_empty(canonical_bars_1m_schema()),
            warnings: Vec::new(),
            gap_bar_count: 0,
        });
    }

    let input = extract_bar_columns(bars)?;
    let mut by_symbol: BTreeMap<u32, BTreeMap<i64, BarRow>> = BTreeMap::new();
    for idx in 0..bars.num_rows() {
        input.ensure_required_not_null(idx)?;
        let symbol = input.symbol_id.value(idx);
        let timestamp = input.timestamp.value(idx);
        by_symbol.entry(symbol).or_default().insert(
            timestamp,
            BarRow {
                open: input.open.value(idx),
                high: input.high.value(idx),
                low: input.low.value(idx),
                close: input.close.value(idx),
                vwap: input.vwap.value(idx),
                volume: input.volume.value(idx),
                buy_volume: value_or_none(input.buy_volume, idx),
                sell_volume: value_or_none(input.sell_volume, idx),
                trade_count: input.trade_count.value(idx),
                bid: value_or_none(input.bid, idx),
                ask: value_or_none(input.ask, idx),
                spread: value_or_none(input.spread, idx),
                is_gap: input.is_gap.value(idx),
                gap_fill_method: input.gap_fill_method.value(idx),
            },
        );
    }

    let mut out_timestamp = Vec::new();
    let mut out_symbol_id = Vec::new();
    let mut out_open = Vec::new();
    let mut out_high = Vec::new();
    let mut out_low = Vec::new();
    let mut out_close = Vec::new();
    let mut out_vwap = Vec::new();
    let mut out_volume = Vec::new();
    let mut out_buy_volume = Vec::new();
    let mut out_sell_volume = Vec::new();
    let mut out_trade_count = Vec::new();
    let mut out_bid: Vec<Option<f64>> = Vec::new();
    let mut out_ask: Vec<Option<f64>> = Vec::new();
    let mut out_spread: Vec<Option<f64>> = Vec::new();
    let mut out_is_gap = Vec::new();
    let mut out_gap_fill_method = Vec::new();
    let mut warnings = Vec::new();
    let mut gap_bar_count = 0_usize;

    for (symbol_id, rows) in by_symbol {
        let symbol_start = rows
            .first_key_value()
            .map(|(ts, _)| *ts)
            .ok_or_else(|| BtError::Data("expected at least one row per symbol".to_string()))?;
        let symbol_end = rows
            .last_key_value()
            .map(|(ts, _)| {
                ts.checked_add(NANOS_PER_SECOND)
                    .ok_or_else(|| BtError::TimeRange("symbol end overflow".to_string()))
            })
            .ok_or_else(|| BtError::Data("expected at least one row per symbol".to_string()))??;

        let (start_ns, end_ns) = match range {
            Some(custom) => (align_to_second(custom.start.0), custom.end.0),
            None => (symbol_start, symbol_end),
        };
        if start_ns >= end_ns {
            continue;
        }

        let mut cursor = start_ns;
        let mut previous_close: Option<f64> = None;
        let mut consecutive_gaps = 0_usize;
        let mut warned_for_gap_run = false;

        while cursor < end_ns {
            if let Some(row) = rows.get(&cursor) {
                push_row(
                    symbol_id,
                    cursor,
                    row,
                    &mut out_timestamp,
                    &mut out_symbol_id,
                    &mut out_open,
                    &mut out_high,
                    &mut out_low,
                    &mut out_close,
                    &mut out_vwap,
                    &mut out_volume,
                    &mut out_buy_volume,
                    &mut out_sell_volume,
                    &mut out_trade_count,
                    &mut out_bid,
                    &mut out_ask,
                    &mut out_spread,
                    &mut out_is_gap,
                    &mut out_gap_fill_method,
                );
                previous_close = Some(row.close);

                if row.is_gap {
                    gap_bar_count = gap_bar_count.saturating_add(1);
                    consecutive_gaps = consecutive_gaps.saturating_add(1);
                } else {
                    consecutive_gaps = 0;
                    warned_for_gap_run = false;
                }
            } else {
                let (fill_price, method) = match previous_close {
                    Some(close) => (close, 1_u8),
                    None => (0.0, 3_u8),
                };
                let gap_row = BarRow {
                    open: fill_price,
                    high: fill_price,
                    low: fill_price,
                    close: fill_price,
                    vwap: fill_price,
                    volume: 0.0,
                    buy_volume: Some(0.0),
                    sell_volume: Some(0.0),
                    trade_count: 0,
                    bid: None,
                    ask: None,
                    spread: None,
                    is_gap: true,
                    gap_fill_method: method,
                };
                push_row(
                    symbol_id,
                    cursor,
                    &gap_row,
                    &mut out_timestamp,
                    &mut out_symbol_id,
                    &mut out_open,
                    &mut out_high,
                    &mut out_low,
                    &mut out_close,
                    &mut out_vwap,
                    &mut out_volume,
                    &mut out_buy_volume,
                    &mut out_sell_volume,
                    &mut out_trade_count,
                    &mut out_bid,
                    &mut out_ask,
                    &mut out_spread,
                    &mut out_is_gap,
                    &mut out_gap_fill_method,
                );
                gap_bar_count = gap_bar_count.saturating_add(1);
                consecutive_gaps = consecutive_gaps.saturating_add(1);
            }

            if consecutive_gaps > 60 && !warned_for_gap_run {
                warnings.push(format!(
                    "symbol {symbol_id}: consecutive gap run exceeded 60 bars at {cursor}"
                ));
                warned_for_gap_run = true;
            }

            cursor = cursor
                .checked_add(NANOS_PER_SECOND)
                .ok_or_else(|| BtError::TimeRange("gap fill cursor overflow".to_string()))?;
        }
    }

    let arrays: Vec<ArrayRef> = vec![
        Arc::new(TimestampNanosecondArray::from(out_timestamp).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)),
    ];

    tracing::debug!(gap_count = gap_bar_count, "filled gaps");
    Ok(GapFillResult {
        batch: RecordBatch::try_new(canonical_bars_1m_schema(), arrays)?,
        warnings,
        gap_bar_count,
    })
}

#[derive(Clone, Debug)]
struct BarColumns<'a> {
    timestamp: &'a TimestampNanosecondArray,
    symbol_id: &'a UInt32Array,
    open: &'a Float64Array,
    high: &'a Float64Array,
    low: &'a Float64Array,
    close: &'a Float64Array,
    vwap: &'a Float64Array,
    volume: &'a Float64Array,
    buy_volume: &'a Float64Array,
    sell_volume: &'a Float64Array,
    trade_count: &'a UInt32Array,
    bid: &'a Float64Array,
    ask: &'a Float64Array,
    spread: &'a Float64Array,
    is_gap: &'a BooleanArray,
    gap_fill_method: &'a UInt8Array,
}

impl<'a> BarColumns<'a> {
    fn ensure_required_not_null(&self, idx: usize) -> Result<()> {
        ensure_not_null(self.timestamp, "timestamp", idx)?;
        ensure_not_null(self.symbol_id, "symbol_id", idx)?;
        ensure_not_null(self.open, "open", idx)?;
        ensure_not_null(self.high, "high", idx)?;
        ensure_not_null(self.low, "low", idx)?;
        ensure_not_null(self.close, "close", idx)?;
        ensure_not_null(self.vwap, "vwap", idx)?;
        ensure_not_null(self.volume, "volume", idx)?;
        ensure_not_null(self.trade_count, "trade_count", idx)?;
        ensure_not_null(self.is_gap, "is_gap", idx)?;
        ensure_not_null(self.gap_fill_method, "gap_fill_method", idx)?;
        Ok(())
    }
}

#[derive(Clone, Debug)]
struct BarRow {
    open: f64,
    high: f64,
    low: f64,
    close: f64,
    vwap: f64,
    volume: f64,
    buy_volume: Option<f64>,
    sell_volume: Option<f64>,
    trade_count: u32,
    bid: Option<f64>,
    ask: Option<f64>,
    spread: Option<f64>,
    is_gap: bool,
    gap_fill_method: u8,
}

fn extract_bar_columns(batch: &RecordBatch) -> Result<BarColumns<'_>> {
    Ok(BarColumns {
        timestamp: timestamp_column(batch, "timestamp")?,
        symbol_id: u32_column(batch, "symbol_id")?,
        open: float64_column(batch, "open")?,
        high: float64_column(batch, "high")?,
        low: float64_column(batch, "low")?,
        close: float64_column(batch, "close")?,
        vwap: float64_column(batch, "vwap")?,
        volume: float64_column(batch, "volume")?,
        buy_volume: float64_column(batch, "buy_volume")?,
        sell_volume: float64_column(batch, "sell_volume")?,
        trade_count: u32_column(batch, "trade_count")?,
        bid: float64_column(batch, "bid")?,
        ask: float64_column(batch, "ask")?,
        spread: float64_column(batch, "spread")?,
        is_gap: bool_column(batch, "is_gap")?,
        gap_fill_method: u8_column(batch, "gap_fill_method")?,
    })
}

#[allow(clippy::too_many_arguments)]
fn push_row(
    symbol_id: u32,
    timestamp: i64,
    row: &BarRow,
    out_timestamp: &mut Vec<i64>,
    out_symbol_id: &mut Vec<u32>,
    out_open: &mut Vec<f64>,
    out_high: &mut Vec<f64>,
    out_low: &mut Vec<f64>,
    out_close: &mut Vec<f64>,
    out_vwap: &mut Vec<f64>,
    out_volume: &mut Vec<f64>,
    out_buy_volume: &mut Vec<Option<f64>>,
    out_sell_volume: &mut Vec<Option<f64>>,
    out_trade_count: &mut Vec<u32>,
    out_bid: &mut Vec<Option<f64>>,
    out_ask: &mut Vec<Option<f64>>,
    out_spread: &mut Vec<Option<f64>>,
    out_is_gap: &mut Vec<bool>,
    out_gap_fill_method: &mut Vec<u8>,
) {
    out_timestamp.push(timestamp);
    out_symbol_id.push(symbol_id);
    out_open.push(row.open);
    out_high.push(row.high);
    out_low.push(row.low);
    out_close.push(row.close);
    out_vwap.push(row.vwap);
    out_volume.push(row.volume);
    out_buy_volume.push(row.buy_volume);
    out_sell_volume.push(row.sell_volume);
    out_trade_count.push(row.trade_count);
    out_bid.push(row.bid);
    out_ask.push(row.ask);
    out_spread.push(row.spread);
    out_is_gap.push(row.is_gap);
    out_gap_fill_method.push(row.gap_fill_method);
}

fn align_to_second(ts: i64) -> i64 {
    ts.div_euclid(NANOS_PER_SECOND) * NANOS_PER_SECOND
}

fn value_or_none(array: &Float64Array, idx: usize) -> Option<f64> {
    if array.is_null(idx) {
        None
    } else {
        Some(array.value(idx))
    }
}

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 bar 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 bar column '{name}'")))?;
    column
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be Float64")))
}

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

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 bar column '{name}'")))?;
    column
        .as_any()
        .downcast_ref::<UInt8Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be UInt8")))
}

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

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(())
}
