use std::mem;
use tokio::sync::mpsc::{Receiver, Sender};

use crate::chopper::types::{Nanos, Row};

pub struct Joiner {
    rx: Vec<Receiver<Box<Row>>>,
    tx: Sender<Box<Row>>,
}

impl Joiner {
    pub fn new(rx: Vec<Receiver<Box<Row>>>, tx: Sender<Box<Row>>) -> Joiner {
        Joiner { rx, tx }
    }

    pub async fn run(&mut self) {
        let mut rows = self.get_initial_rows().await;
        loop {
            if rows.is_empty() {
                // we don't have any sources that could produce any rows
                break;
            }

            // find timestamp we are going to be processing
            let timestamp = rows.iter().map(|r| r.timestamp).min().unwrap();
            // find row with next timestamp
            let row_idx = Self::get_first_matching_row_idx(&rows, timestamp);

            // then we want to keep trying to deliver rows from same source as long as
            // it keeps producing rows with same timestamp
            loop {
                let (row, next_row_same_timestamp) = self
                    .get_row_with_replacement(&mut rows, row_idx, timestamp)
                    .await;
                self.tx.send(row).await.unwrap();

                // if next buffered row has same timestamp as what we are processing right now,
                // just keep processing this row index
                if !next_row_same_timestamp {
                    break;
                }
            }
        }
    }

    async fn get_initial_rows(&mut self) -> Vec<Box<Row>> {
        let mut rows: Vec<Box<Row>> = Vec::with_capacity(self.rx.len());
        // collect one row from every source
        let mut current_rx = 0;
        loop {
            if current_rx == self.rx.len() {
                break;
            }

            let row = match self.rx[current_rx].recv().await {
                None => {
                    // remove sources which don't have any data
                    self.rx.remove(current_rx);
                    continue;
                }
                Some(row) => row,
            };
            current_rx += 1;
            rows.push(row);
        }
        rows
    }

    fn get_first_matching_row_idx(rows: &Vec<Box<Row>>, timestamp: Nanos) -> usize {
        let mut i = 0;
        loop {
            if i == rows.len() {
                panic!("can't find row for timestamp {}", timestamp);
            }
            if rows[i].timestamp == timestamp {
                break;
            }
            i = i + 1;
        }
        i
    }

    /// returns buffered row along with boolean indicating if next row has same timestamp
    /// as the one just returned here
    async fn get_row_with_replacement(
        &mut self,
        rows: &mut Vec<Box<Row>>,
        row_idx: usize,
        timestamp: Nanos,
    ) -> (Box<Row>, bool) {
        // get next row first
        match self.rx[row_idx].recv().await {
            None => {
                // rx is out of rows, so just remove it and the associated row
                self.rx.remove(row_idx);
                (rows.remove(row_idx), false)
            }
            Some(next_row) => {
                let next_row_same_timestamp = next_row.timestamp == timestamp;
                (
                    // stick next row in place of current row and just get current one out
                    mem::replace(&mut rows[row_idx], next_row),
                    next_row_same_timestamp,
                )
            }
        }
    }
}
