use std::collections::BTreeMap;

use arrow::array::{Array, BooleanArray, TimestampNanosecondArray, UInt32Array};
use arrow::record_batch::RecordBatch;
use bt_kernel::{BtError, Result, SymbolId};
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct DataQualityReport {
    pub symbol_id: SymbolId,
    pub row_count: u64,
    pub gap_bar_count: u64,
    pub max_consecutive_gap: u64,
    pub gap_ratio: f64,
    pub warnings: Vec<String>,
}

pub fn generate_quality_reports(bars: &RecordBatch) -> Result<Vec<DataQualityReport>> {
    if bars.num_rows() == 0 {
        return Ok(Vec::new());
    }

    let symbol_id = u32_column(bars, "symbol_id")?;
    let timestamp = timestamp_column(bars, "timestamp")?;
    let is_gap = bool_column(bars, "is_gap")?;

    let mut by_symbol: BTreeMap<u32, Vec<(i64, bool)>> = BTreeMap::new();
    for idx in 0..bars.num_rows() {
        ensure_not_null(symbol_id, "symbol_id", idx)?;
        ensure_not_null(timestamp, "timestamp", idx)?;
        ensure_not_null(is_gap, "is_gap", idx)?;

        by_symbol
            .entry(symbol_id.value(idx))
            .or_default()
            .push((timestamp.value(idx), is_gap.value(idx)));
    }

    let mut reports = Vec::with_capacity(by_symbol.len());
    for (symbol, mut rows) in by_symbol {
        rows.sort_by_key(|row| row.0);
        let row_count = rows.len() as u64;
        let gap_bar_count = rows.iter().filter(|(_, gap)| *gap).count() as u64;

        let mut max_consecutive_gap = 0_u64;
        let mut current_consecutive_gap = 0_u64;
        for (_, gap) in rows {
            if gap {
                current_consecutive_gap = current_consecutive_gap.saturating_add(1);
                max_consecutive_gap = max_consecutive_gap.max(current_consecutive_gap);
            } else {
                current_consecutive_gap = 0;
            }
        }

        let mut warnings = Vec::new();
        if max_consecutive_gap > 60 {
            warnings.push(format!(
                "symbol {symbol}: max consecutive gap ({max_consecutive_gap}) exceeded threshold 60"
            ));
        }

        reports.push(DataQualityReport {
            symbol_id: SymbolId(symbol),
            row_count,
            gap_bar_count,
            max_consecutive_gap,
            gap_ratio: if row_count == 0 {
                0.0
            } else {
                gap_bar_count as f64 / row_count as f64
            },
            warnings,
        });
    }

    Ok(reports)
}

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