// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Common code used in the integration test binaries

use hex::decode;
use serde_json::Value;

use arrow::util::integration_util::ArrowJsonBatch;

use arrow::array::*;
use arrow::datatypes::{DataType, Field, IntervalUnit, Schema};
use arrow::error::{ArrowError, Result};
use arrow::record_batch::RecordBatch;
use arrow::{
    buffer::Buffer,
    buffer::MutableBuffer,
    datatypes::ToByteSlice,
    util::{bit_util, integration_util::*},
};

use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::sync::Arc;

/// The expected username for the basic auth integration test.
pub const AUTH_USERNAME: &str = "arrow";
/// The expected password for the basic auth integration test.
pub const AUTH_PASSWORD: &str = "flight";

pub mod flight_client_scenarios;
pub mod flight_server_scenarios;

pub struct ArrowFile {
    pub schema: Schema,
    // we can evolve this into a concrete Arrow type
    // this is temporarily not being read from
    pub _dictionaries: HashMap<i64, ArrowJsonDictionaryBatch>,
    pub batches: Vec<RecordBatch>,
}

pub fn read_json_file(json_name: &str) -> Result<ArrowFile> {
    let json_file = File::open(json_name)?;
    let reader = BufReader::new(json_file);
    let arrow_json: Value = serde_json::from_reader(reader).unwrap();
    let schema = Schema::from(&arrow_json["schema"])?;
    // read dictionaries
    let mut dictionaries = HashMap::new();
    if let Some(dicts) = arrow_json.get("dictionaries") {
        for d in dicts
            .as_array()
            .expect("Unable to get dictionaries as array")
        {
            let json_dict: ArrowJsonDictionaryBatch = serde_json::from_value(d.clone())
                .expect("Unable to get dictionary from JSON");
            // TODO: convert to a concrete Arrow type
            dictionaries.insert(json_dict.id, json_dict);
        }
    }

    let mut batches = vec![];
    for b in arrow_json["batches"].as_array().unwrap() {
        let json_batch: ArrowJsonBatch = serde_json::from_value(b.clone()).unwrap();
        let batch = record_batch_from_json(&schema, json_batch, Some(&dictionaries))?;
        batches.push(batch);
    }
    Ok(ArrowFile {
        schema,
        _dictionaries: dictionaries,
        batches,
    })
}

fn record_batch_from_json(
    schema: &Schema,
    json_batch: ArrowJsonBatch,
    json_dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
) -> Result<RecordBatch> {
    let mut columns = vec![];

    for (field, json_col) in schema.fields().iter().zip(json_batch.columns) {
        let col = array_from_json(field, json_col, json_dictionaries)?;
        columns.push(col);
    }

    RecordBatch::try_new(Arc::new(schema.clone()), columns)
}

/// Construct an Arrow array from a partially typed JSON column
fn array_from_json(
    field: &Field,
    json_col: ArrowJsonColumn,
    dictionaries: Option<&HashMap<i64, ArrowJsonDictionaryBatch>>,
) -> Result<ArrayRef> {
    match field.data_type() {
        DataType::Null => Ok(Arc::new(NullArray::new(json_col.count))),
        DataType::Boolean => {
            let mut b = BooleanBuilder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_bool().unwrap()),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Int8 => {
            let mut b = Int8Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_i64().ok_or_else(|| {
                        ArrowError::JsonError(format!(
                            "Unable to get {:?} as int64",
                            value
                        ))
                    })? as i8),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Int16 => {
            let mut b = Int16Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_i64().unwrap() as i16),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Int32
        | DataType::Date32
        | DataType::Time32(_)
        | DataType::Interval(IntervalUnit::YearMonth) => {
            let mut b = Int32Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_i64().unwrap() as i32),
                    _ => b.append_null(),
                }?;
            }
            let array = Arc::new(b.finish()) as ArrayRef;
            arrow::compute::cast(&array, field.data_type())
        }
        DataType::Int64
        | DataType::Date64
        | DataType::Time64(_)
        | DataType::Timestamp(_, _)
        | DataType::Duration(_)
        | DataType::Interval(IntervalUnit::DayTime) => {
            let mut b = Int64Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(match value {
                        Value::Number(n) => n.as_i64().unwrap(),
                        Value::String(s) => {
                            s.parse().expect("Unable to parse string as i64")
                        }
                        _ => panic!("Unable to parse {:?} as number", value),
                    }),
                    _ => b.append_null(),
                }?;
            }
            let array = Arc::new(b.finish()) as ArrayRef;
            arrow::compute::cast(&array, field.data_type())
        }
        DataType::UInt8 => {
            let mut b = UInt8Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_u64().unwrap() as u8),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::UInt16 => {
            let mut b = UInt16Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_u64().unwrap() as u16),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::UInt32 => {
            let mut b = UInt32Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_u64().unwrap() as u32),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::UInt64 => {
            let mut b = UInt64Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(
                        value
                            .as_str()
                            .unwrap()
                            .parse()
                            .expect("Unable to parse string as u64"),
                    ),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Float32 => {
            let mut b = Float32Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_f64().unwrap() as f32),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Float64 => {
            let mut b = Float64Builder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_f64().unwrap()),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Binary => {
            let mut b = BinaryBuilder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => {
                        let v = decode(value.as_str().unwrap()).unwrap();
                        b.append_value(&v)
                    }
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::LargeBinary => {
            let mut b = LargeBinaryBuilder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => {
                        let v = decode(value.as_str().unwrap()).unwrap();
                        b.append_value(&v)
                    }
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::Utf8 => {
            let mut b = StringBuilder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_str().unwrap()),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::LargeUtf8 => {
            let mut b = LargeStringBuilder::new(json_col.count);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => b.append_value(value.as_str().unwrap()),
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::FixedSizeBinary(len) => {
            let mut b = FixedSizeBinaryBuilder::new(json_col.count, *len);
            for (is_valid, value) in json_col
                .validity
                .as_ref()
                .unwrap()
                .iter()
                .zip(json_col.data.unwrap())
            {
                match is_valid {
                    1 => {
                        let v = hex::decode(value.as_str().unwrap()).unwrap();
                        b.append_value(&v)
                    }
                    _ => b.append_null(),
                }?;
            }
            Ok(Arc::new(b.finish()))
        }
        DataType::List(child_field) => {
            let null_buf = create_null_buf(&json_col);
            let children = json_col.children.clone().unwrap();
            let child_array = array_from_json(
                &child_field,
                children.get(0).unwrap().clone(),
                dictionaries,
            )?;
            let offsets: Vec<i32> = json_col
                .offset
                .unwrap()
                .iter()
                .map(|v| v.as_i64().unwrap() as i32)
                .collect();
            let list_data = ArrayData::builder(field.data_type().clone())
                .len(json_col.count)
                .offset(0)
                .add_buffer(Buffer::from(&offsets.to_byte_slice()))
                .add_child_data(child_array.data().clone())
                .null_bit_buffer(null_buf)
                .build();
            Ok(Arc::new(ListArray::from(list_data)))
        }
        DataType::LargeList(child_field) => {
            let null_buf = create_null_buf(&json_col);
            let children = json_col.children.clone().unwrap();
            let child_array = array_from_json(
                &child_field,
                children.get(0).unwrap().clone(),
                dictionaries,
            )?;
            let offsets: Vec<i64> = json_col
                .offset
                .unwrap()
                .iter()
                .map(|v| match v {
                    Value::Number(n) => n.as_i64().unwrap(),
                    Value::String(s) => s.parse::<i64>().unwrap(),
                    _ => panic!("64-bit offset must be either string or number"),
                })
                .collect();
            let list_data = ArrayData::builder(field.data_type().clone())
                .len(json_col.count)
                .offset(0)
                .add_buffer(Buffer::from(&offsets.to_byte_slice()))
                .add_child_data(child_array.data().clone())
                .null_bit_buffer(null_buf)
                .build();
            Ok(Arc::new(LargeListArray::from(list_data)))
        }
        DataType::FixedSizeList(child_field, _) => {
            let children = json_col.children.clone().unwrap();
            let child_array = array_from_json(
                &child_field,
                children.get(0).unwrap().clone(),
                dictionaries,
            )?;
            let null_buf = create_null_buf(&json_col);
            let list_data = ArrayData::builder(field.data_type().clone())
                .len(json_col.count)
                .add_child_data(child_array.data().clone())
                .null_bit_buffer(null_buf)
                .build();
            Ok(Arc::new(FixedSizeListArray::from(list_data)))
        }
        DataType::Struct(fields) => {
            // construct struct with null data
            let null_buf = create_null_buf(&json_col);
            let mut array_data = ArrayData::builder(field.data_type().clone())
                .len(json_col.count)
                .null_bit_buffer(null_buf);

            for (field, col) in fields.iter().zip(json_col.children.unwrap()) {
                let array = array_from_json(field, col, dictionaries)?;
                array_data = array_data.add_child_data(array.data().clone());
            }

            let array = StructArray::from(array_data.build());
            Ok(Arc::new(array))
        }
        DataType::Dictionary(key_type, value_type) => {
            let dict_id = field.dict_id().ok_or_else(|| {
                ArrowError::JsonError(format!(
                    "Unable to find dict_id for field {:?}",
                    field
                ))
            })?;
            // find dictionary
            let dictionary = dictionaries
                .ok_or_else(|| {
                    ArrowError::JsonError(format!(
                        "Unable to find any dictionaries for field {:?}",
                        field
                    ))
                })?
                .get(&dict_id);
            match dictionary {
                Some(dictionary) => dictionary_array_from_json(
                    field, json_col, key_type, value_type, dictionary,
                ),
                None => Err(ArrowError::JsonError(format!(
                    "Unable to find dictionary for field {:?}",
                    field
                ))),
            }
        }
        t => Err(ArrowError::JsonError(format!(
            "data type {:?} not supported",
            t
        ))),
    }
}

fn dictionary_array_from_json(
    field: &Field,
    json_col: ArrowJsonColumn,
    dict_key: &DataType,
    dict_value: &DataType,
    dictionary: &ArrowJsonDictionaryBatch,
) -> Result<ArrayRef> {
    match dict_key {
        DataType::Int8
        | DataType::Int16
        | DataType::Int32
        | DataType::Int64
        | DataType::UInt8
        | DataType::UInt16
        | DataType::UInt32
        | DataType::UInt64 => {
            let null_buf = create_null_buf(&json_col);

            // build the key data into a buffer, then construct values separately
            let key_field = Field::new_dict(
                "key",
                dict_key.clone(),
                field.is_nullable(),
                field
                    .dict_id()
                    .expect("Dictionary fields must have a dict_id value"),
                field
                    .dict_is_ordered()
                    .expect("Dictionary fields must have a dict_is_ordered value"),
            );
            let keys = array_from_json(&key_field, json_col, None)?;
            // note: not enough info on nullability of dictionary
            let value_field = Field::new("value", dict_value.clone(), true);
            println!("dictionary value type: {:?}", dict_value);
            let values =
                array_from_json(&value_field, dictionary.data.columns[0].clone(), None)?;

            // convert key and value to dictionary data
            let dict_data = ArrayData::builder(field.data_type().clone())
                .len(keys.len())
                .add_buffer(keys.data().buffers()[0].clone())
                .null_bit_buffer(null_buf)
                .add_child_data(values.data().clone())
                .build();

            let array = match dict_key {
                DataType::Int8 => {
                    Arc::new(Int8DictionaryArray::from(dict_data)) as ArrayRef
                }
                DataType::Int16 => Arc::new(Int16DictionaryArray::from(dict_data)),
                DataType::Int32 => Arc::new(Int32DictionaryArray::from(dict_data)),
                DataType::Int64 => Arc::new(Int64DictionaryArray::from(dict_data)),
                DataType::UInt8 => Arc::new(UInt8DictionaryArray::from(dict_data)),
                DataType::UInt16 => Arc::new(UInt16DictionaryArray::from(dict_data)),
                DataType::UInt32 => Arc::new(UInt32DictionaryArray::from(dict_data)),
                DataType::UInt64 => Arc::new(UInt64DictionaryArray::from(dict_data)),
                _ => unreachable!(),
            };
            Ok(array)
        }
        _ => Err(ArrowError::JsonError(format!(
            "Dictionary key type {:?} not supported",
            dict_key
        ))),
    }
}

/// A helper to create a null buffer from a Vec<bool>
fn create_null_buf(json_col: &ArrowJsonColumn) -> Buffer {
    let num_bytes = bit_util::ceil(json_col.count, 8);
    let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, false);
    json_col
        .validity
        .clone()
        .unwrap()
        .iter()
        .enumerate()
        .for_each(|(i, v)| {
            let null_slice = null_buf.as_slice_mut();
            if *v != 0 {
                bit_util::set_bit(null_slice, i);
            }
        });
    null_buf.into()
}