use std::{cmp::Ordering, collections::HashMap, sync::Arc};

use arrow::{
    array::{Array, RecordBatch, UInt32Array},
    compute::take,
    datatypes::Schema,
};
use mpera_frontend::op::{JoinKind, OpRef};
use rayon::prelude::*;

use crate::{
    op::{OpOut, StatefulBatchOp},
    rowkey,
    runtimedb::RuntimeDB,
};

pub struct JoinBatchOp {
    pub opref: OpRef,
    pub left: OpRef,
    pub right: OpRef,
    pub left_on: OpRef,
    pub right_on: OpRef,
    pub kind: JoinKind,
    left_batches: Vec<RecordBatch>,
    right_batches: Vec<RecordBatch>,
    left_on_batches: Vec<RecordBatch>,
    right_on_batches: Vec<RecordBatch>,
    emitted_left_batches: usize,
    emitted_right_batches: usize,
}

impl JoinBatchOp {
    pub fn new(
        opref: OpRef,
        left: OpRef,
        right: OpRef,
        left_on: OpRef,
        right_on: OpRef,
        kind: JoinKind,
    ) -> Self {
        Self {
            opref,
            left,
            right,
            left_on,
            right_on,
            kind,
            left_batches: Vec::new(),
            right_batches: Vec::new(),
            left_on_batches: Vec::new(),
            right_on_batches: Vec::new(),
            emitted_left_batches: 0,
            emitted_right_batches: 0,
        }
    }
}

impl StatefulBatchOp for JoinBatchOp {
    fn opref(&self) -> OpRef {
        self.opref
    }

    fn consume(&mut self, rtdb: &RuntimeDB) -> OpOut {
        self.sync_batches(rtdb)?;

        if self.left_batches.is_empty() {
            return Ok(None);
        }
        if self.right_batches.is_empty() {
            return Ok(None);
        }

        if !matches!(self.kind, JoinKind::Inner) {
            return Ok(None);
        }

        let out = self.consume_inner_ready()?;
        if out.is_empty() {
            return Ok(None);
        }
        Ok(Some(out))
    }

    fn finalize(&mut self) -> crate::Result<Vec<RecordBatch>> {
        if matches!(self.kind, JoinKind::Inner) {
            return Ok(Vec::new());
        }
        if self.left_batches.is_empty() {
            return Ok(Vec::new());
        }

        let left = merge_batches("left", &self.left_batches)?;
        let right = merge_batches("right", &self.right_batches)?;
        let left_on = merge_batches("left_on", &self.left_on_batches)?;
        let right_on = merge_batches("right_on", &self.right_on_batches)?;

        Ok(vec![join_batch(
            &left, &right, &left_on, &right_on, &self.kind,
        )?])
    }
}

impl JoinBatchOp {
    fn sync_batches(&mut self, rtdb: &RuntimeDB) -> crate::Result<()> {
        sync_op_batches(rtdb, self.left, &mut self.left_batches)?;
        sync_op_batches(rtdb, self.right, &mut self.right_batches)?;
        sync_op_batches(rtdb, self.left_on, &mut self.left_on_batches)?;
        sync_op_batches(rtdb, self.right_on, &mut self.right_on_batches)?;
        Ok(())
    }

    fn consume_inner_ready(&mut self) -> crate::Result<Vec<RecordBatch>> {
        let mut out = Vec::new();

        let new_left = &self.left_batches[self.emitted_left_batches..];
        let new_left_on = &self.left_on_batches[self.emitted_left_batches..];
        if !new_left.is_empty() {
            let left = merge_batches("left", new_left)?;
            let left_on = merge_batches("left_on", new_left_on)?;
            let right = merge_batches("right", &self.right_batches)?;
            let right_on = merge_batches("right_on", &self.right_on_batches)?;
            out.push(join_batch(&left, &right, &left_on, &right_on, &self.kind)?);
        }

        let new_right = &self.right_batches[self.emitted_right_batches..];
        let new_right_on = &self.right_on_batches[self.emitted_right_batches..];
        let old_left = &self.left_batches[..self.emitted_left_batches];
        let old_left_on = &self.left_on_batches[..self.emitted_left_batches];
        if !new_right.is_empty() && !old_left.is_empty() {
            let left = merge_batches("left", old_left)?;
            let left_on = merge_batches("left_on", old_left_on)?;
            let right = merge_batches("right", new_right)?;
            let right_on = merge_batches("right_on", new_right_on)?;
            out.push(join_batch(&left, &right, &left_on, &right_on, &self.kind)?);
        }

        self.emitted_left_batches = self.left_batches.len();
        self.emitted_right_batches = self.right_batches.len();
        Ok(out)
    }
}

fn sync_op_batches(
    rtdb: &RuntimeDB,
    opref: OpRef,
    out: &mut Vec<RecordBatch>,
) -> crate::Result<()> {
    let Some(batches) = rtdb.read_batches(opref) else {
        return Ok(());
    };
    if batches.len() <= out.len() {
        return Ok(());
    }
    out.extend(batches[out.len()..].iter().cloned());
    Ok(())
}

fn merge_batches(label: &str, batches: &[RecordBatch]) -> crate::Result<RecordBatch> {
    let Some(first) = batches.first() else {
        return Err(autark_error::Error::Reader(format!(
            "join `{label}` produced no record batches"
        )));
    };
    if batches.len() == 1 {
        return Ok(first.clone());
    }
    Ok(arrow::compute::concat_batches(&first.schema(), batches)?)
}

fn join_batch(
    left: &RecordBatch,
    right: &RecordBatch,
    left_on: &RecordBatch,
    right_on: &RecordBatch,
    kind: &JoinKind,
) -> crate::Result<RecordBatch> {
    ensure_aligned("left", left, left_on)?;
    ensure_aligned("right", right, right_on)?;

    let (left_ix, right_ix) = match kind {
        JoinKind::Inner => equality_join_indices(left_on, right_on, false)?,
        JoinKind::LeftOuter => equality_join_indices(left_on, right_on, true)?,
        JoinKind::AsOfBackward => asof_join_indices(left_on, right_on, false)?,
        JoinKind::AsOfForward => asof_join_indices(left_on, right_on, true)?,
    };

    build_joined_batch(left, right, &left_ix, &right_ix)
}

fn ensure_aligned(side: &str, batch: &RecordBatch, on: &RecordBatch) -> crate::Result<()> {
    if batch.num_rows() != on.num_rows() {
        return Err(autark_error::Error::Reader(format!(
            "{side} join input requires matching row counts: batch={} on={}",
            batch.num_rows(),
            on.num_rows()
        )));
    }
    Ok(())
}

fn equality_join_indices(
    left_on: &RecordBatch,
    right_on: &RecordBatch,
    include_unmatched_left: bool,
) -> crate::Result<(Vec<u32>, Vec<Option<u32>>)> {
    let mut right_map = HashMap::<String, Vec<u32>>::new();
    for row in 0..right_on.num_rows() {
        right_map
            .entry(rowkey::row_key("join", right_on, row)?)
            .or_default()
            .push(row as u32);
    }

    let pairs = (0..left_on.num_rows())
        .into_par_iter()
        .map(|row| equality_row_pairs(left_on, row, &right_map, include_unmatched_left))
        .collect::<crate::Result<Vec<_>>>()?;

    Ok(flatten_pairs(pairs))
}

fn asof_join_indices(
    left_on: &RecordBatch,
    right_on: &RecordBatch,
    forward: bool,
) -> crate::Result<(Vec<u32>, Vec<Option<u32>>)> {
    if left_on.num_columns() != 1 || right_on.num_columns() != 1 {
        return Err(autark_error::Error::Reader(
            "asof join currently requires a single join key column".to_string(),
        ));
    }

    let left_key = left_on.column(0).as_ref();
    let right_key = right_on.column(0).as_ref();
    if left_key.data_type() != right_key.data_type() {
        return Err(autark_error::Error::Reader(format!(
            "asof join requires matching key dtypes: left={} right={}",
            left_key.data_type(),
            right_key.data_type()
        )));
    }

    let right_ix = (0..left_on.num_rows())
        .into_par_iter()
        .map(|left_row| best_asof_match(left_key, left_row, right_key, forward))
        .collect::<crate::Result<Vec<_>>>()?;
    let left_ix = (0..left_on.num_rows()).map(|row| row as u32).collect();

    Ok((left_ix, right_ix))
}

fn best_asof_match(
    left_key: &dyn Array,
    left_row: usize,
    right_key: &dyn Array,
    forward: bool,
) -> crate::Result<Option<u32>> {
    let mut best = None::<usize>;
    for right_row in 0..right_key.len() {
        let ord = rowkey::compare_scalar("join", left_key, left_row, right_key, right_row)?;
        let qualifies = if forward {
            ord != Ordering::Greater
        } else {
            ord != Ordering::Less
        };
        if !qualifies {
            continue;
        }

        match best {
            None => best = Some(right_row),
            Some(current) => {
                let better = if forward {
                    rowkey::compare_scalar("join", right_key, right_row, right_key, current)?
                        == Ordering::Less
                } else {
                    rowkey::compare_scalar("join", right_key, right_row, right_key, current)?
                        == Ordering::Greater
                };
                if better {
                    best = Some(right_row);
                }
            }
        }
    }
    Ok(best.map(|ix| ix as u32))
}

fn equality_row_pairs(
    left_on: &RecordBatch,
    row: usize,
    right_map: &HashMap<String, Vec<u32>>,
    include_unmatched_left: bool,
) -> crate::Result<Vec<(u32, Option<u32>)>> {
    let Some(matches) = right_map.get(&rowkey::row_key("join", left_on, row)?) else {
        if include_unmatched_left {
            return Ok(vec![(row as u32, None)]);
        }
        return Ok(Vec::new());
    };

    Ok(matches
        .iter()
        .copied()
        .map(|right_row| (row as u32, Some(right_row)))
        .collect())
}

fn flatten_pairs(pairs: Vec<Vec<(u32, Option<u32>)>>) -> (Vec<u32>, Vec<Option<u32>>) {
    let len = pairs.iter().map(Vec::len).sum();
    let mut left_ix = Vec::with_capacity(len);
    let mut right_ix = Vec::with_capacity(len);

    for pair_group in pairs {
        for (left, right) in pair_group {
            left_ix.push(left);
            right_ix.push(right);
        }
    }

    (left_ix, right_ix)
}

fn build_joined_batch(
    left: &RecordBatch,
    right: &RecordBatch,
    left_ix: &[u32],
    right_ix: &[Option<u32>],
) -> crate::Result<RecordBatch> {
    let left_ix = UInt32Array::from(left_ix.to_vec());
    let right_ix = UInt32Array::from(right_ix.to_vec());

    let mut fields = left.schema().fields().iter().cloned().collect::<Vec<_>>();
    fields.extend(right.schema().fields().iter().cloned());

    let mut columns = left
        .columns()
        .iter()
        .map(|column| Ok(take(column.as_ref(), &left_ix, None)?))
        .collect::<crate::Result<Vec<_>>>()?;
    columns.extend(
        right
            .columns()
            .iter()
            .map(|column| Ok(take(column.as_ref(), &right_ix, None)?))
            .collect::<crate::Result<Vec<_>>>()?,
    );

    Ok(RecordBatch::try_new(
        Arc::new(Schema::new(fields)),
        columns,
    )?)
}
