use std::collections::HashMap;

use arrow::array::{Array, ArrayRef, StringArray, TimestampNanosecondArray, UInt32Array};
use arrow::compute::take;
use arrow::record_batch::RecordBatch;
use bt_kernel::{BtError, Result};

pub fn dedup_and_sort_raw_trades(batch: &RecordBatch) -> Result<RecordBatch> {
    if batch.num_rows() == 0 {
        return Ok(RecordBatch::new_empty(batch.schema()));
    }

    let provider = string_column(batch, "provider")?;
    let symbol_raw = string_column(batch, "symbol_raw")?;
    let trade_id = string_column(batch, "trade_id")?;
    let timestamp = timestamp_column(batch, "timestamp")?;
    let symbol_id = u32_column(batch, "symbol_id")?;

    let mut keep_by_key: HashMap<(String, String, String), usize> = HashMap::new();
    for idx in 0..batch.num_rows() {
        ensure_not_null(provider, "provider", idx)?;
        ensure_not_null(symbol_raw, "symbol_raw", idx)?;
        ensure_not_null(trade_id, "trade_id", idx)?;
        ensure_not_null(timestamp, "timestamp", idx)?;
        ensure_not_null(symbol_id, "symbol_id", idx)?;

        let key = (
            provider.value(idx).to_string(),
            symbol_raw.value(idx).to_string(),
            trade_id.value(idx).to_string(),
        );

        let should_replace = match keep_by_key.get(&key) {
            Some(existing_idx) => timestamp.value(idx) < timestamp.value(*existing_idx),
            None => true,
        };
        if should_replace {
            keep_by_key.insert(key, idx);
        }
    }

    let mut indices = keep_by_key.into_values().collect::<Vec<_>>();
    indices.sort_by(|a, b| {
        symbol_id
            .value(*a)
            .cmp(&symbol_id.value(*b))
            .then_with(|| timestamp.value(*a).cmp(&timestamp.value(*b)))
            .then_with(|| trade_id.value(*a).cmp(trade_id.value(*b)))
    });

    take_rows(batch, &indices)
}

fn take_rows(batch: &RecordBatch, indices: &[usize]) -> Result<RecordBatch> {
    let index_array = UInt32Array::from(
        indices
            .iter()
            .map(|idx| {
                u32::try_from(*idx)
                    .map_err(|_| BtError::Data(format!("row index too large to fit in u32: {idx}")))
            })
            .collect::<Result<Vec<_>>>()?,
    );

    let mut columns: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
    for column in batch.columns() {
        columns.push(take(column.as_ref(), &index_array, None)?);
    }

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

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

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

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