use std::ops::Deref;

use arrow2::array::{BinaryArray, MutableArray, MutableBinaryArray, TryExtendFromSelf};
use arrow2::bitmap::Bitmap;
use arrow2::error::Error;

#[test]
fn new() {
    assert_eq!(MutableBinaryArray::<i32>::new().len(), 0);

    let a = MutableBinaryArray::<i32>::with_capacity(2);
    assert_eq!(a.len(), 0);
    assert!(a.offsets().capacity() >= 2);
    assert_eq!(a.values().capacity(), 0);

    let a = MutableBinaryArray::<i32>::with_capacities(2, 60);
    assert_eq!(a.len(), 0);
    assert!(a.offsets().capacity() >= 2);
    assert!(a.values().capacity() >= 60);
}

#[test]
fn from_iter() {
    let iter = (0..3u8).map(|x| Some(vec![x; x as usize]));
    let a: MutableBinaryArray<i32> = iter.clone().collect();
    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
    assert_eq!(a.offsets().as_slice(), &[0, 0, 1, 3]);
    assert_eq!(a.validity(), None);

    let a = unsafe { MutableBinaryArray::<i32>::from_trusted_len_iter_unchecked(iter) };
    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
    assert_eq!(a.offsets().as_slice(), &[0, 0, 1, 3]);
    assert_eq!(a.validity(), None);
}

#[test]
fn from_trusted_len_iter() {
    let data = vec![vec![0; 0], vec![1; 1], vec![2; 2]];
    let a: MutableBinaryArray<i32> = data.iter().cloned().map(Some).collect();
    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
    assert_eq!(a.offsets().as_slice(), &[0, 0, 1, 3]);
    assert_eq!(a.validity(), None);

    let a = MutableBinaryArray::<i32>::from_trusted_len_iter(data.iter().cloned().map(Some));
    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
    assert_eq!(a.offsets().as_slice(), &[0, 0, 1, 3]);
    assert_eq!(a.validity(), None);

    let a = MutableBinaryArray::<i32>::try_from_trusted_len_iter::<Error, _, _>(
        data.iter().cloned().map(Some).map(Ok),
    )
    .unwrap();
    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
    assert_eq!(a.offsets().as_slice(), &[0, 0, 1, 3]);
    assert_eq!(a.validity(), None);

    let a = MutableBinaryArray::<i32>::from_trusted_len_values_iter(data.iter().cloned());
    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
    assert_eq!(a.offsets().as_slice(), &[0, 0, 1, 3]);
    assert_eq!(a.validity(), None);
}

#[test]
fn push_null() {
    let mut array = MutableBinaryArray::<i32>::new();
    array.push::<&str>(None);

    let array: BinaryArray<i32> = array.into();
    assert_eq!(array.validity(), Some(&Bitmap::from([false])));
}

#[test]
fn pop() {
    let mut a = MutableBinaryArray::<i32>::new();
    a.push(Some(b"first"));
    a.push(Some(b"second"));
    a.push::<Vec<u8>>(None);
    a.push_null();

    assert_eq!(a.pop(), None);
    assert_eq!(a.len(), 3);
    assert_eq!(a.pop(), None);
    assert_eq!(a.len(), 2);
    assert_eq!(a.pop(), Some(b"second".to_vec()));
    assert_eq!(a.len(), 1);
    assert_eq!(a.pop(), Some(b"first".to_vec()));
    assert_eq!(a.len(), 0);
    assert_eq!(a.pop(), None);
    assert_eq!(a.len(), 0);
}

#[test]
fn pop_all_some() {
    let mut a = MutableBinaryArray::<i32>::new();
    a.push(Some(b"first"));
    a.push(Some(b"second"));
    a.push(Some(b"third"));
    a.push(Some(b"fourth"));

    for _ in 0..4 {
        a.push(Some(b"aaaa"));
    }

    a.push(Some(b"bbbb"));

    assert_eq!(a.pop(), Some(b"bbbb".to_vec()));
    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
    assert_eq!(a.len(), 5);
    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
    assert_eq!(a.pop(), Some(b"fourth".to_vec()));
    assert_eq!(a.pop(), Some(b"third".to_vec()));
    assert_eq!(a.pop(), Some(b"second".to_vec()));
    assert_eq!(a.pop(), Some(b"first".to_vec()));
    assert!(a.is_empty());
    assert_eq!(a.pop(), None);
}

#[test]
fn extend_trusted_len_values() {
    let mut array = MutableBinaryArray::<i32>::new();

    array.extend_trusted_len_values(vec![b"first".to_vec(), b"second".to_vec()].into_iter());
    array.extend_trusted_len_values(vec![b"third".to_vec()].into_iter());
    array.extend_trusted_len(vec![None, Some(b"fourth".to_vec())].into_iter());

    let array: BinaryArray<i32> = array.into();

    assert_eq!(array.values().as_slice(), b"firstsecondthirdfourth");
    assert_eq!(array.offsets().as_slice(), &[0, 5, 11, 16, 16, 22]);
    assert_eq!(
        array.validity(),
        Some(&Bitmap::from_u8_slice([0b00010111], 5))
    );
}

#[test]
fn extend_trusted_len() {
    let mut array = MutableBinaryArray::<i32>::new();

    array.extend_trusted_len(vec![Some(b"first".to_vec()), Some(b"second".to_vec())].into_iter());
    array.extend_trusted_len(vec![None, Some(b"third".to_vec())].into_iter());

    let array: BinaryArray<i32> = array.into();

    assert_eq!(array.values().as_slice(), b"firstsecondthird");
    assert_eq!(array.offsets().as_slice(), &[0, 5, 11, 11, 16]);
    assert_eq!(
        array.validity(),
        Some(&Bitmap::from_u8_slice([0b00001011], 4))
    );
}

#[test]
fn extend_from_self() {
    let mut a = MutableBinaryArray::<i32>::from([Some(b"aa"), None]);

    a.try_extend_from_self(&a.clone()).unwrap();

    assert_eq!(
        a,
        MutableBinaryArray::<i32>::from([Some(b"aa"), None, Some(b"aa"), None])
    );
}

#[test]
fn test_set_validity() {
    let mut array = MutableBinaryArray::<i32>::new();
    array.push(Some(b"first"));
    array.push(Some(b"second"));
    array.push(Some(b"third"));
    array.set_validity(Some([false, false, true].into()));

    assert!(!array.is_valid(0));
    assert!(!array.is_valid(1));
    assert!(array.is_valid(2));
}

#[test]
fn test_apply_validity() {
    let mut array = MutableBinaryArray::<i32>::new();
    array.push(Some(b"first"));
    array.push(Some(b"second"));
    array.push(Some(b"third"));
    array.set_validity(Some([true, true, true].into()));

    array.apply_validity(|mut mut_bitmap| {
        mut_bitmap.set(1, false);
        mut_bitmap.set(2, false);
        mut_bitmap
    });

    assert!(array.is_valid(0));
    assert!(!array.is_valid(1));
    assert!(!array.is_valid(2));
}

#[test]
fn test_apply_validity_with_no_validity_inited() {
    let mut array = MutableBinaryArray::<i32>::new();
    array.push(Some(b"first"));
    array.push(Some(b"second"));
    array.push(Some(b"third"));

    array.apply_validity(|mut mut_bitmap| {
        mut_bitmap.set(1, false);
        mut_bitmap.set(2, false);
        mut_bitmap
    });

    assert!(array.is_valid(0));
    assert!(array.is_valid(1));
    assert!(array.is_valid(2));
}