// 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.

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

use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, Schema, TimeUnit};
use datafusion::{
    common::JoinSide, logical_expr::Operator, parquet::arrow::PARQUET_FIELD_ID_META_KEY,
    scalar::ScalarValue,
};
use datafusion_ext_plans::{agg::AggFunction, joins::join_utils::JoinType};

use crate::error::PlanSerDeError;

// include the generated protobuf source as a submodule
#[allow(clippy::all)]
pub mod protobuf {
    include!(concat!(env!("OUT_DIR"), "/plan.protobuf.rs"));
}

pub mod error;
pub mod planner;

pub(crate) fn proto_error<S: Into<String>>(message: S) -> PlanSerDeError {
    PlanSerDeError::General(message.into())
}

#[macro_export]
macro_rules! convert_required {
    ($PB:expr) => {{
        if let Some(field) = $PB.as_ref() {
            field.try_into()
        } else {
            Err(proto_error("Missing required field in protobuf"))
        }
    }};
}

#[macro_export]
macro_rules! into_required {
    ($PB:expr) => {{
        if let Some(field) = $PB.as_ref() {
            Ok(field.into())
        } else {
            Err(proto_error("Missing required field in protobuf"))
        }
    }};
}

#[macro_export]
macro_rules! convert_box_required {
    ($self:expr, $PB:expr) => {{
        if let Some(field) = $PB.as_ref() {
            $self.create_plan(field.as_ref())
        } else {
            Err(proto_error("Missing required field in protobuf"))
        }
    }};
}

pub fn from_proto_binary_op(op: &str) -> Result<Operator, PlanSerDeError> {
    match op {
        "And" => Ok(Operator::And),
        "Or" => Ok(Operator::Or),
        "Eq" => Ok(Operator::Eq),
        "NotEq" => Ok(Operator::NotEq),
        "LtEq" => Ok(Operator::LtEq),
        "Lt" => Ok(Operator::Lt),
        "Gt" => Ok(Operator::Gt),
        "GtEq" => Ok(Operator::GtEq),
        "Plus" => Ok(Operator::Plus),
        "Minus" => Ok(Operator::Minus),
        "Multiply" => Ok(Operator::Multiply),
        "Divide" => Ok(Operator::Divide),
        "Modulo" => Ok(Operator::Modulo),
        "IsDistinctFrom" => Ok(Operator::IsDistinctFrom),
        "IsNotDistinctFrom" => Ok(Operator::IsNotDistinctFrom),
        "BitwiseAnd" => Ok(Operator::BitwiseAnd),
        "BitwiseOr" => Ok(Operator::BitwiseOr),
        "BitwiseXor" => Ok(Operator::BitwiseXor),
        "BitwiseShiftLeft" => Ok(Operator::BitwiseShiftLeft),
        "BitwiseShiftRight" => Ok(Operator::BitwiseShiftRight),
        "RegexIMatch" => Ok(Operator::RegexIMatch),
        "RegexMatch" => Ok(Operator::RegexMatch),
        "RegexNotIMatch" => Ok(Operator::RegexNotIMatch),
        "RegexNotMatch" => Ok(Operator::RegexNotMatch),
        "StringConcat" => Ok(Operator::StringConcat),
        other => Err(proto_error(format!(
            "Unsupported binary operator '{other:?}'"
        ))),
    }
}

impl From<protobuf::JoinType> for JoinType {
    fn from(t: protobuf::JoinType) -> Self {
        match t {
            protobuf::JoinType::Inner => JoinType::Inner,
            protobuf::JoinType::Left => JoinType::Left,
            protobuf::JoinType::Right => JoinType::Right,
            protobuf::JoinType::Full => JoinType::Full,
            protobuf::JoinType::Semi => JoinType::LeftSemi,
            protobuf::JoinType::Anti => JoinType::LeftAnti,
            protobuf::JoinType::Existence => JoinType::Existence,
        }
    }
}

impl From<protobuf::JoinSide> for JoinSide {
    fn from(t: protobuf::JoinSide) -> Self {
        match t {
            protobuf::JoinSide::LeftSide => JoinSide::Left,
            protobuf::JoinSide::RightSide => JoinSide::Right,
        }
    }
}

impl From<protobuf::AggFunction> for AggFunction {
    fn from(agg_fun: protobuf::AggFunction) -> AggFunction {
        match agg_fun {
            protobuf::AggFunction::Min => AggFunction::Min,
            protobuf::AggFunction::Max => AggFunction::Max,
            protobuf::AggFunction::Sum => AggFunction::Sum,
            protobuf::AggFunction::Avg => AggFunction::Avg,
            protobuf::AggFunction::Count => AggFunction::Count,
            protobuf::AggFunction::CollectList => AggFunction::CollectList,
            protobuf::AggFunction::CollectSet => AggFunction::CollectSet,
            protobuf::AggFunction::First => AggFunction::First,
            protobuf::AggFunction::FirstIgnoresNull => AggFunction::FirstIgnoresNull,
            protobuf::AggFunction::BloomFilter => AggFunction::BloomFilter,
            protobuf::AggFunction::BrickhouseCollect => AggFunction::BrickhouseCollect,
            protobuf::AggFunction::BrickhouseCombineUnique => AggFunction::BrickhouseCombineUnique,
            protobuf::AggFunction::Udaf => AggFunction::Udaf,
        }
    }
}

impl protobuf::TimeUnit {
    pub fn from_arrow_time_unit(val: &TimeUnit) -> Self {
        match val {
            TimeUnit::Second => protobuf::TimeUnit::Second,
            TimeUnit::Millisecond => protobuf::TimeUnit::Millisecond,
            TimeUnit::Microsecond => protobuf::TimeUnit::Microsecond,
            TimeUnit::Nanosecond => protobuf::TimeUnit::Nanosecond,
        }
    }
    pub fn try_from_to_arrow(time_unit_i32: i32) -> Result<TimeUnit, PlanSerDeError> {
        let pb_time_unit = protobuf::TimeUnit::try_from(time_unit_i32).expect("invalid TimeUnit");
        Ok(match pb_time_unit {
            protobuf::TimeUnit::Second => TimeUnit::Second,
            protobuf::TimeUnit::Millisecond => TimeUnit::Millisecond,
            protobuf::TimeUnit::Microsecond => TimeUnit::Microsecond,
            protobuf::TimeUnit::Nanosecond => TimeUnit::Nanosecond,
        })
    }
}

impl protobuf::IntervalUnit {
    pub fn from_arrow_interval_unit(interval_unit: &IntervalUnit) -> Self {
        match interval_unit {
            IntervalUnit::YearMonth => protobuf::IntervalUnit::YearMonth,
            IntervalUnit::DayTime => protobuf::IntervalUnit::DayTime,
            IntervalUnit::MonthDayNano => protobuf::IntervalUnit::MonthDayNano,
        }
    }

    pub fn try_from_to_arrow(interval_unit_i32: i32) -> Result<IntervalUnit, PlanSerDeError> {
        let pb_interval_unit =
            protobuf::IntervalUnit::try_from(interval_unit_i32).expect("invalid IntervalUnit");
        Ok(match pb_interval_unit {
            protobuf::IntervalUnit::YearMonth => IntervalUnit::YearMonth,
            protobuf::IntervalUnit::DayTime => IntervalUnit::DayTime,
            protobuf::IntervalUnit::MonthDayNano => IntervalUnit::MonthDayNano,
        })
    }
}

impl TryInto<arrow::datatypes::DataType> for &protobuf::arrow_type::ArrowTypeEnum {
    type Error = PlanSerDeError;
    fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
        use protobuf::arrow_type;
        Ok(match self {
            arrow_type::ArrowTypeEnum::None(_) => DataType::Null,
            arrow_type::ArrowTypeEnum::Bool(_) => DataType::Boolean,
            arrow_type::ArrowTypeEnum::Uint8(_) => DataType::UInt8,
            arrow_type::ArrowTypeEnum::Int8(_) => DataType::Int8,
            arrow_type::ArrowTypeEnum::Uint16(_) => DataType::UInt16,
            arrow_type::ArrowTypeEnum::Int16(_) => DataType::Int16,
            arrow_type::ArrowTypeEnum::Uint32(_) => DataType::UInt32,
            arrow_type::ArrowTypeEnum::Int32(_) => DataType::Int32,
            arrow_type::ArrowTypeEnum::Uint64(_) => DataType::UInt64,
            arrow_type::ArrowTypeEnum::Int64(_) => DataType::Int64,
            arrow_type::ArrowTypeEnum::Float16(_) => DataType::Float16,
            arrow_type::ArrowTypeEnum::Float32(_) => DataType::Float32,
            arrow_type::ArrowTypeEnum::Float64(_) => DataType::Float64,
            arrow_type::ArrowTypeEnum::Utf8(_) => DataType::Utf8,
            arrow_type::ArrowTypeEnum::LargeUtf8(_) => DataType::LargeUtf8,
            arrow_type::ArrowTypeEnum::Binary(_) => DataType::Binary,
            arrow_type::ArrowTypeEnum::FixedSizeBinary(size) => DataType::FixedSizeBinary(*size),
            arrow_type::ArrowTypeEnum::LargeBinary(_) => DataType::LargeBinary,
            arrow_type::ArrowTypeEnum::Date32(_) => DataType::Date32,
            arrow_type::ArrowTypeEnum::Date64(_) => DataType::Date64,
            arrow_type::ArrowTypeEnum::Duration(time_unit) => {
                DataType::Duration(protobuf::TimeUnit::try_from_to_arrow(*time_unit)?)
            }
            arrow_type::ArrowTypeEnum::Timestamp(protobuf::Timestamp {
                time_unit,
                timezone,
            }) => DataType::Timestamp(
                protobuf::TimeUnit::try_from_to_arrow(*time_unit)?,
                match timezone.len() {
                    0 => None,
                    _ => Some(timezone.to_owned().into()),
                },
            ),
            arrow_type::ArrowTypeEnum::Time32(time_unit) => {
                DataType::Time32(protobuf::TimeUnit::try_from_to_arrow(*time_unit)?)
            }
            arrow_type::ArrowTypeEnum::Time64(time_unit) => {
                DataType::Time64(protobuf::TimeUnit::try_from_to_arrow(*time_unit)?)
            }
            arrow_type::ArrowTypeEnum::Interval(interval_unit) => {
                DataType::Interval(protobuf::IntervalUnit::try_from_to_arrow(*interval_unit)?)
            }
            arrow_type::ArrowTypeEnum::Decimal(protobuf::Decimal { whole, fractional }) => {
                DataType::Decimal128(*whole as u8, *fractional as i8)
            }
            arrow_type::ArrowTypeEnum::List(list) => {
                let list_type: &protobuf::Field = list
                    .as_ref()
                    .field_type
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: List message missing required field 'field_type'"))?
                    .as_ref();
                DataType::List(Arc::new(list_type.try_into()?))
            }
            arrow_type::ArrowTypeEnum::LargeList(list) => {
                let list_type: &protobuf::Field = list
                    .as_ref()
                    .field_type
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: List message missing required field 'field_type'"))?
                    .as_ref();
                DataType::LargeList(Arc::new(list_type.try_into()?))
            }
            arrow_type::ArrowTypeEnum::FixedSizeList(list) => {
                let list_type: &protobuf::Field = list
                    .as_ref()
                    .field_type
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: List message missing required field 'field_type'"))?
                    .as_ref();
                let list_size = list.list_size;
                DataType::FixedSizeList(Arc::new(list_type.try_into()?), list_size)
            }
            arrow_type::ArrowTypeEnum::Struct(strct) => DataType::Struct(
                strct
                    .sub_field_types
                    .iter()
                    .map(|field| Ok(Arc::new(field.try_into()?)))
                    .collect::<Result<Fields, Self::Error>>()?,
            ),
            arrow_type::ArrowTypeEnum::Union(_union) => {
                // let union_mode = protobuf::UnionMode::try_from(union.union_mode)
                //     .ok_or_else(|| {
                //         proto_error(
                //             "Protobuf deserialization error: Unknown union mode type",
                //         )
                //     })?;
                // let union_mode = match union_mode {
                //     protobuf::UnionMode::Dense => UnionMode::Dense,
                //     protobuf::UnionMode::Sparse => UnionMode::Sparse,
                // };
                // let union_types = union
                //     .union_types
                //     .iter()
                //     .map(|field| field.try_into())
                //     .collect::<Result<Vec<_>, _>>()?;
                // DataType::Union(union_types, _, union_mode)
                unimplemented!()
            }
            arrow_type::ArrowTypeEnum::Map(map) => {
                let key_type: &protobuf::Field = map
                    .as_ref()
                    .key_type
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: Map message missing required field 'key_type'"))?
                    .as_ref();
                let value_type: &protobuf::Field = map
                    .as_ref()
                    .value_type
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: Map message missing required field 'value_type'"))?
                    .as_ref();

                let vec_field = vec![
                    Arc::new(key_type.try_into()?),
                    Arc::new(value_type.try_into()?),
                ];
                let fields = Arc::new(Field::new(
                    "entries",
                    DataType::Struct(vec_field.into()),
                    false,
                ));
                DataType::Map(fields, false)
            }
            arrow_type::ArrowTypeEnum::Dictionary(dict) => {
                let pb_key_datatype = dict
                    .as_ref()
                    .key
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: Dictionary message missing required field 'key'"))?;
                let pb_value_datatype = dict
                    .as_ref()
                    .value
                    .as_ref()
                    .ok_or_else(|| proto_error("Protobuf deserialization error: Dictionary message missing required field 'key'"))?;
                let key_datatype: DataType = pb_key_datatype.as_ref().try_into()?;
                let value_datatype: DataType = pb_value_datatype.as_ref().try_into()?;
                DataType::Dictionary(Box::new(key_datatype), Box::new(value_datatype))
            }
        })
    }
}

#[allow(clippy::from_over_into)]
impl Into<arrow::datatypes::DataType> for protobuf::PrimitiveScalarType {
    fn into(self) -> arrow::datatypes::DataType {
        match self {
            protobuf::PrimitiveScalarType::Bool => DataType::Boolean,
            protobuf::PrimitiveScalarType::Uint8 => DataType::UInt8,
            protobuf::PrimitiveScalarType::Int8 => DataType::Int8,
            protobuf::PrimitiveScalarType::Uint16 => DataType::UInt16,
            protobuf::PrimitiveScalarType::Int16 => DataType::Int16,
            protobuf::PrimitiveScalarType::Uint32 => DataType::UInt32,
            protobuf::PrimitiveScalarType::Int32 => DataType::Int32,
            protobuf::PrimitiveScalarType::Uint64 => DataType::UInt64,
            protobuf::PrimitiveScalarType::Int64 => DataType::Int64,
            protobuf::PrimitiveScalarType::Float32 => DataType::Float32,
            protobuf::PrimitiveScalarType::Float64 => DataType::Float64,
            protobuf::PrimitiveScalarType::Utf8 => DataType::Utf8,
            protobuf::PrimitiveScalarType::LargeUtf8 => DataType::LargeUtf8,
            protobuf::PrimitiveScalarType::Date32 => DataType::Date32,
            protobuf::PrimitiveScalarType::Null => DataType::Null,
            protobuf::PrimitiveScalarType::Decimal128 => DataType::Decimal128(1, 0),
            protobuf::PrimitiveScalarType::Date64 => DataType::Date64,
            protobuf::PrimitiveScalarType::TimestampSecond => {
                DataType::Timestamp(TimeUnit::Second, None)
            }
            protobuf::PrimitiveScalarType::TimestampMillisecond => {
                DataType::Timestamp(TimeUnit::Millisecond, None)
            }
            protobuf::PrimitiveScalarType::TimestampMicrosecond => {
                DataType::Timestamp(TimeUnit::Microsecond, None)
            }
            protobuf::PrimitiveScalarType::TimestampNanosecond => {
                DataType::Timestamp(TimeUnit::Nanosecond, None)
            }
            protobuf::PrimitiveScalarType::IntervalYearmonth => {
                DataType::Interval(IntervalUnit::YearMonth)
            }
            protobuf::PrimitiveScalarType::IntervalDaytime => {
                DataType::Interval(IntervalUnit::DayTime)
            }
        }
    }
}

impl TryInto<DataType> for &protobuf::ArrowType {
    type Error = PlanSerDeError;
    fn try_into(self) -> Result<DataType, Self::Error> {
        let pb_arrow_type = self.arrow_type_enum.as_ref().ok_or_else(|| {
            proto_error(
                "Protobuf deserialization error: ArrowType missing required field 'data_type'",
            )
        })?;
        pb_arrow_type.try_into()
    }
}

impl TryInto<DataType> for &Box<protobuf::List> {
    type Error = PlanSerDeError;
    fn try_into(self) -> Result<DataType, Self::Error> {
        let list_ref = self.as_ref();
        match &list_ref.field_type {
            Some(pb_field) => {
                let pb_field_ref = pb_field.as_ref();
                let arrow_field: Field = pb_field_ref.try_into()?;
                Ok(DataType::List(Arc::new(arrow_field)))
            }
            None => Err(proto_error(
                "List message missing required field 'field_type'",
            )),
        }
    }
}

impl TryInto<Field> for &protobuf::Field {
    type Error = PlanSerDeError;
    fn try_into(self) -> Result<Field, Self::Error> {
        build_arrow_field(self)
    }
}

fn build_arrow_field(field: &protobuf::Field) -> Result<Field, PlanSerDeError> {
    let pb_datatype = field.arrow_type.as_ref().ok_or_else(|| {
        proto_error(
            "Protobuf deserialization error: Field message missing required field 'arrow_type'",
        )
    })?;
    let arrow_field = Field::new(
        field.name.as_str(),
        pb_datatype.as_ref().try_into()?,
        field.nullable,
    );

    if field.field_id == 0 {
        Ok(arrow_field)
    } else {
        Ok(arrow_field.with_metadata(HashMap::from([(
            PARQUET_FIELD_ID_META_KEY.to_string(),
            field.field_id.to_string(),
        )])))
    }
}

impl TryInto<Schema> for &protobuf::Schema {
    type Error = PlanSerDeError;

    fn try_into(self) -> Result<Schema, PlanSerDeError> {
        let fields = self
            .columns
            .iter()
            .map(build_arrow_field)
            .collect::<Result<Vec<_>, _>>()?;
        Ok(Schema::new(fields))
    }
}

impl TryInto<datafusion::scalar::ScalarValue> for &protobuf::ScalarValue {
    type Error = PlanSerDeError;
    fn try_into(self) -> Result<datafusion::scalar::ScalarValue, Self::Error> {
        let mut ipc_reader =
            arrow::ipc::reader::StreamReader::try_new(&self.ipc_bytes[..], None)
                .map_err(|e| proto_error(format!("error deserializing arrow stream: {e}")))?;
        let batch = ipc_reader.next().expect("missing record batch")?;
        let scalar = ScalarValue::try_from_array(batch.column(0), 0)?;
        Ok(scalar)
    }
}
