//! Codec type for serializing/deserializing Hamelin UDFs and physical plan
//! extensions in DataFusion plans.
//!
//! Embedders that ship Hamelin plans over the wire (e.g. phantom) hold a
//! [`HamelinPhysicalCodec`] as a field of their own `PhysicalExtensionCodec`
//! impl and delegate to it before falling through to their own logic.
//!
//! Hamelin doesn't expose a logical codec: phantom (and other distributed
//! callers) lower to physical plans before serializing, so logical-extension
//! round-trips aren't on the hot path. If that ever changes, add a
//! `HamelinLogicalCodec` here mirroring the physical one.
//!
//! Each method follows the DataFusion convention of returning
//! `DataFusionError::NotImplemented` when the input isn't a Hamelin extension —
//! callers chain via:
//!
//! ```ignore
//! match self.hamelin.try_encode(node, buf) {
//!     Err(DataFusionError::NotImplemented(_)) => {
//!         // fall through to embedder-specific encoding
//!     }
//!     other => other,
//! }
//! ```

use std::sync::Arc;

use datafusion::arrow::datatypes::DataType;
use datafusion::common::{not_impl_err, DataFusionError};
use datafusion::execution::TaskContext;
use datafusion::logical_expr::{AggregateUDF, ScalarUDF};
use datafusion::physical_plan::ExecutionPlan;
use datafusion_proto::physical_plan::PhysicalExtensionCodec;
use prost::Message;
use serde::{Deserialize, Serialize};

use crate::exec::TrimStringsExec;
use crate::udf;

type DfResult<T> = Result<T, DataFusionError>;

/// Magic prefix on every Hamelin physical-extension wire payload. Lets us tell
/// "this is a Hamelin extension" from "this is something the embedder wrote"
/// without trying to decode the buffer first.
const HAMELIN_PREFIX: &[u8] = b"hmln:";

/// Wire envelope for Hamelin physical-extension nodes. Encoded after the
/// `hmln:` magic prefix as protobuf via prost.
///
/// Adding a new variant: append a new `oneof` arm with a fresh proto field
/// number. Old decoders that don't know the variant will return
/// `NotImplemented`; new decoders dispatch on the variant tag.
#[derive(Clone, PartialEq, Message)]
struct HamelinPhysicalExtension {
    #[prost(oneof = "hamelin_physical_extension::Kind", tags = "1")]
    kind: Option<hamelin_physical_extension::Kind>,
}

mod hamelin_physical_extension {
    use prost::Oneof;

    #[derive(Clone, PartialEq, Oneof)]
    pub enum Kind {
        #[prost(message, tag = "1")]
        TrimStringsExec(super::TrimStringsExecProto),
    }
}

/// Body of a `TrimStringsExec` extension node.
#[derive(Clone, PartialEq, Message)]
struct TrimStringsExecProto {
    #[prost(uint64, tag = "1")]
    byte_len: u64,
}

/// Payload for the parameterized `hamelin_array_cast` UDF.
#[derive(Serialize, Deserialize)]
struct ArrayCastPayload {
    target_type: DataType,
    cast_descriptor: udf::CastDescriptor,
}

// ============================================================================
// Public codec types
// ============================================================================

/// Codec for Hamelin UDFs and physical plan extensions.
///
/// Implements DataFusion's [`PhysicalExtensionCodec`]. Embedders compose this
/// into their own codec by holding it as a field and delegating; methods return
/// `DataFusionError::NotImplemented` for nodes/UDFs that aren't Hamelin's so
/// the embedder can fall through to its own logic.
#[derive(Debug, Default, Clone, Copy)]
pub struct HamelinPhysicalCodec;

impl PhysicalExtensionCodec for HamelinPhysicalCodec {
    fn try_decode(
        &self,
        buf: &[u8],
        inputs: &[Arc<dyn ExecutionPlan>],
        _ctx: &TaskContext,
    ) -> DfResult<Arc<dyn ExecutionPlan>> {
        // Strip the `hmln:` prefix. If it isn't there, the buffer belongs to
        // somebody else's codec — return NotImplemented so an embedder chaining
        // codecs can fall through to its own logic.
        let Some(proto_bytes) = buf.strip_prefix(HAMELIN_PREFIX) else {
            return not_impl_err!(
                "HamelinPhysicalCodec: buffer is not a Hamelin physical extension"
            );
        };
        let envelope = HamelinPhysicalExtension::decode(proto_bytes).map_err(|e| {
            DataFusionError::Internal(format!(
                "HamelinPhysicalCodec: malformed Hamelin extension payload: {e}"
            ))
        })?;
        match envelope.kind {
            Some(hamelin_physical_extension::Kind::TrimStringsExec(body)) => {
                decode_trim_strings_exec(body.byte_len, inputs)
            }
            None => {
                not_impl_err!("HamelinPhysicalCodec: Hamelin extension envelope had no kind set")
            }
        }
    }

    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> DfResult<()> {
        if let Some(trim) = node.as_any().downcast_ref::<TrimStringsExec>() {
            return encode_trim_strings_exec(trim, buf);
        }
        not_impl_err!(
            "HamelinPhysicalCodec: not a Hamelin physical extension node ({})",
            node.name()
        )
    }

    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> DfResult<Arc<ScalarUDF>> {
        decode_udf(name, buf)
    }

    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> DfResult<()> {
        encode_udf(node, buf)
    }

    fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> DfResult<Arc<AggregateUDF>> {
        decode_udaf(name, buf)
    }

    fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> DfResult<()> {
        encode_udaf(node, buf)
    }
}

// ============================================================================
// Internal: extension-node round-trip
// ============================================================================

fn encode_trim_strings_exec(node: &TrimStringsExec, buf: &mut Vec<u8>) -> DfResult<()> {
    let envelope = HamelinPhysicalExtension {
        kind: Some(hamelin_physical_extension::Kind::TrimStringsExec(
            TrimStringsExecProto {
                byte_len: node.byte_len(),
            },
        )),
    };
    buf.extend_from_slice(HAMELIN_PREFIX);
    envelope
        .encode(buf)
        .map_err(|e| DataFusionError::Internal(format!("Failed to encode TrimStringsExec: {e}")))
}

fn decode_trim_strings_exec(
    byte_len: u64,
    inputs: &[Arc<dyn ExecutionPlan>],
) -> DfResult<Arc<dyn ExecutionPlan>> {
    let input = inputs.first().cloned().ok_or_else(|| {
        DataFusionError::Plan("TrimStringsExec expects exactly one input".to_string())
    })?;
    Ok(Arc::new(TrimStringsExec::new(input, byte_len)))
}

// ============================================================================
// Internal: UDF/UDAF round-trip
// ============================================================================

fn decode_udf(name: &str, buf: &[u8]) -> DfResult<Arc<ScalarUDF>> {
    let udf = match name {
        "hamelin_array_avg" => udf::array_avg_udf(),
        "hamelin_array_coalesce_default" => udf::array_coalesce_default_udf(),
        "hamelin_array_get_field" => udf::array_get_field_udf(),
        "hamelin_filter_null" => udf::filter_null_udf(),
        "hamelin_array_length" => udf::hamelin_array_length_udf(),
        "hamelin_array_sum" => udf::array_sum_udf(),
        "hamelin_array_variant_get" => udf::array_variant_get_udf(),
        "hamelin_array_variant_to_json" => udf::array_variant_to_json_udf(),
        "hamelin_zip" => udf::array_zip_udf(),
        "hamelin_map_variant_get_values" => udf::map_variant_get_values_udf(),
        "hamelin_map_variant_to_json_values" => udf::map_variant_to_json_values_udf(),
        "hamelin_is_ipv4" => udf::is_ipv4_udf(),
        "hamelin_is_ipv6" => udf::is_ipv6_udf(),
        "hamelin_cidr_contains" => udf::cidr_contains_udf(),
        "hamelin_json_to_variant" => udf::json_to_variant_udf(),
        "hamelin_map_extract" => udf::map_extract_udf(),
        "hamelin_map_from_arrays" => udf::map_from_arrays_udf(),
        "hamelin_map_from_entries" => udf::map_from_entries_udf(),
        "hamelin_parse_timestamp" => udf::parse_timestamp_udf(),
        "hamelin_regexp_extract_all" => udf::regexp_extract_all_udf(),
        "hamelin_regexp_split" => udf::regexp_split_udf(),
        "hamelin_cast_to_variant" => udf::cast_to_variant_udf(),
        "hamelin_url_decode" => udf::url_decode_udf(),
        "hamelin_url_encode" => udf::url_encode_udf(),
        "hamelin_uuid5" => udf::uuid5_udf(),
        "hamelin_variant_get" => udf::variant_get_udf(),
        "hamelin_to_json_string" => udf::variant_to_json_udf(),
        "hamelin_width_bucket" => udf::width_bucket_array_udf(),
        "hamelin_to_millis" => udf::to_millis_udf(),
        "hamelin_to_nanos" => udf::to_nanos_udf(),
        "hamelin_from_millis" => udf::from_millis_udf(),
        "hamelin_from_nanos" => udf::from_nanos_udf(),
        "hamelin_to_months" => udf::to_months_udf(),
        "hamelin_from_months" => udf::from_months_udf(),

        "hamelin_from_variant" => {
            let target_type: DataType = serde_json::from_slice(buf).map_err(|e| {
                DataFusionError::Internal(format!("Failed to decode hamelin_from_variant: {e}"))
            })?;
            udf::from_variant_udf(target_type)
        }

        "hamelin_array_cast" => {
            let payload: ArrayCastPayload = serde_json::from_slice(buf).map_err(|e| {
                DataFusionError::Internal(format!("Failed to decode hamelin_array_cast: {e}"))
            })?;
            udf::array_cast_udf(payload.target_type, payload.cast_descriptor)
        }

        other => {
            return not_impl_err!("HamelinCodec: not a Hamelin scalar UDF ({other})");
        }
    };
    Ok(Arc::new(udf))
}

fn encode_udf(node: &ScalarUDF, buf: &mut Vec<u8>) -> DfResult<()> {
    match node.name() {
        "hamelin_from_variant" => {
            let inner = node
                .inner()
                .as_any()
                .downcast_ref::<udf::FromVariantUdf>()
                .ok_or_else(|| {
                    DataFusionError::Internal("Failed to downcast hamelin_from_variant".to_string())
                })?;
            serde_json::to_writer(buf, inner.target_type()).map_err(|e| {
                DataFusionError::Internal(format!("Failed to encode hamelin_from_variant: {e}"))
            })?;
            Ok(())
        }

        "hamelin_array_cast" => {
            let inner = node
                .inner()
                .as_any()
                .downcast_ref::<udf::ArrayCastUdf>()
                .ok_or_else(|| {
                    DataFusionError::Internal("Failed to downcast hamelin_array_cast".to_string())
                })?;
            let payload = ArrayCastPayload {
                target_type: inner.target_type().clone(),
                cast_descriptor: inner.cast_descriptor().clone(),
            };
            serde_json::to_writer(buf, &payload).map_err(|e| {
                DataFusionError::Internal(format!("Failed to encode hamelin_array_cast: {e}"))
            })?;
            Ok(())
        }

        // Stateless singletons — recognized but nothing to serialize.
        "hamelin_array_avg"
        | "hamelin_array_coalesce_default"
        | "hamelin_array_get_field"
        | "hamelin_array_length"
        | "hamelin_filter_null"
        | "hamelin_array_sum"
        | "hamelin_array_variant_get"
        | "hamelin_array_variant_to_json"
        | "hamelin_zip"
        | "hamelin_map_variant_get_values"
        | "hamelin_map_variant_to_json_values"
        | "hamelin_is_ipv4"
        | "hamelin_is_ipv6"
        | "hamelin_cidr_contains"
        | "hamelin_json_to_variant"
        | "hamelin_map_extract"
        | "hamelin_map_from_arrays"
        | "hamelin_map_from_entries"
        | "hamelin_parse_timestamp"
        | "hamelin_regexp_extract_all"
        | "hamelin_regexp_split"
        | "hamelin_cast_to_variant"
        | "hamelin_url_decode"
        | "hamelin_url_encode"
        | "hamelin_uuid5"
        | "hamelin_variant_get"
        | "hamelin_to_json_string"
        | "hamelin_width_bucket"
        | "hamelin_to_millis"
        | "hamelin_to_nanos"
        | "hamelin_from_millis"
        | "hamelin_from_nanos"
        | "hamelin_to_months"
        | "hamelin_from_months" => Ok(()),

        other => not_impl_err!("HamelinCodec: not a Hamelin scalar UDF ({other})"),
    }
}

fn decode_udaf(name: &str, _buf: &[u8]) -> DfResult<Arc<AggregateUDF>> {
    let udaf = match name {
        "hamelin_any_value" => udf::any_value_udaf(),
        "hamelin_map_agg" => udf::map_agg_udaf(),
        "hamelin_multimap_agg" => udf::multimap_agg_udaf(),
        "hamelin_set_agg" => udf::set_agg_udaf(),
        "hamelin_sliding_array_agg" => udf::sliding_array_agg_udaf(),
        other => {
            return not_impl_err!("HamelinCodec: not a Hamelin aggregate UDF ({other})");
        }
    };
    Ok(Arc::new(udaf))
}

fn encode_udaf(node: &AggregateUDF, _buf: &mut Vec<u8>) -> DfResult<()> {
    match node.name() {
        "hamelin_any_value"
        | "hamelin_map_agg"
        | "hamelin_multimap_agg"
        | "hamelin_set_agg"
        | "hamelin_sliding_array_agg" => Ok(()),
        other => not_impl_err!("HamelinCodec: not a Hamelin aggregate UDF ({other})"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::exec::TrimStringsExec;
    use datafusion::arrow::datatypes::{Field, Schema};
    use datafusion::physical_plan::empty::EmptyExec;
    use datafusion::prelude::SessionContext;

    fn dummy_input() -> Arc<dyn ExecutionPlan> {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "a",
            datafusion::arrow::datatypes::DataType::Utf8,
            false,
        )]));
        Arc::new(EmptyExec::new(schema))
    }

    /// A non-Hamelin extension buffer (no `hmln:` prefix) must surface as
    /// `NotImplemented` so a chained codec can fall through.
    #[test]
    fn try_decode_returns_not_implemented_for_non_hamelin_buffer() {
        let codec = HamelinPhysicalCodec;
        let ctx = SessionContext::new();
        let task_ctx = ctx.task_ctx();
        let inputs: Vec<Arc<dyn ExecutionPlan>> = vec![dummy_input()];

        let result = codec.try_decode(b"phantom-payload-xyz", &inputs, &task_ctx);
        match result {
            Err(DataFusionError::NotImplemented(_)) => {}
            other => panic!("expected NotImplemented, got {other:?}"),
        }
    }

    /// Round-trip: encode a TrimStringsExec, then decode. Should produce a node
    /// with the same byte_len.
    #[test]
    fn try_decode_round_trip_for_trim_strings_exec() {
        let codec = HamelinPhysicalCodec;
        let ctx = SessionContext::new();
        let task_ctx = ctx.task_ctx();
        let input = dummy_input();
        let original = Arc::new(TrimStringsExec::new(input.clone(), 64));

        let mut buf = Vec::new();
        codec
            .try_encode(original.clone() as Arc<dyn ExecutionPlan>, &mut buf)
            .expect("encode");
        assert!(
            buf.starts_with(HAMELIN_PREFIX),
            "encoded buffer must start with hmln: prefix, got {buf:?}"
        );

        let decoded = codec
            .try_decode(&buf, &[input], &task_ctx)
            .expect("decode round-trip");
        let trim = decoded
            .as_any()
            .downcast_ref::<TrimStringsExec>()
            .expect("decoded back to TrimStringsExec");
        assert_eq!(trim.byte_len(), 64);
    }
}
