//! Wire message metadata and codec

use std::fmt;

use rkyv::api::high::{HighDeserializer, HighSerializer, HighValidator};
use rkyv::bytecheck::CheckBytes;
use rkyv::rancor::Error as RkyvError;
use rkyv::ser::allocator::ArenaHandle;
use rkyv::util::AlignedVec;
use rkyv::{Archive, Deserialize, Serialize};

use super::compression;
use super::frame::FrameHeader;
use super::rpc::error::WireError;
use super::rpc::number::ByteCount;
use super::rpc::{Ack, Flow, In, Out, Reply};

const DEFAULT_MIN_COMPRESS_BYTES: u64 = 2048;
const DEFAULT_MAX_BODY_BYTES: u64 = 1024 * 1024;
const DEFAULT_MAX_RAW_BYTES: u64 = 8 * 1024 * 1024;
const DEFAULT_MIN_SAVED_BYTES: u64 = 64;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct MessageName(&'static str);

impl MessageName {
    pub const fn new(value: &'static str) -> Self {
        Self(value)
    }

    pub const fn as_str(self) -> &'static str {
        self.0
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct MessageSchema(u64);

impl MessageSchema {
    pub const fn from_name(name: MessageName) -> Self {
        Self(schema_id(name.as_str()))
    }

    pub const fn from_u64(value: u64) -> Self {
        Self(value)
    }

    pub const fn combine(self, name: MessageName) -> Self {
        Self(self.0 ^ schema_id(name.as_str()).rotate_left(13))
    }

    pub const fn get(self) -> u64 {
        self.0
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum MessageKind {
    Req,
    Rep,
    Frame,
    InFrame,
    OutFrame,
    Evt,
}

/// Compression policy used while encoding a wire frame.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Compression {
    Off,
    Auto { min_bytes: ByteCount },
    On,
}

impl Compression {
    pub const DEFAULT: Self = Self::Auto {
        min_bytes: ByteCount::new(DEFAULT_MIN_COMPRESS_BYTES),
    };
}

impl Default for Compression {
    fn default() -> Self {
        Self::DEFAULT
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct CompressionRatio(u16);

impl CompressionRatio {
    pub const DEFAULT: Self = Self(32);

    pub fn new(value: u16) -> Result<Self, WireCodecError> {
        if value == 0 {
            return Err(WireCodecError::new(
                "compression ratio must be greater than zero",
                "provide a ratio like 32",
            ));
        }
        Ok(Self(value))
    }

    pub const fn get(self) -> u16 {
        self.0
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WireLimits {
    pub max_body: ByteCount,
    pub max_raw: ByteCount,
    pub max_ratio: CompressionRatio,
    pub min_saved: ByteCount,
}

impl Default for WireLimits {
    fn default() -> Self {
        Self {
            max_body: ByteCount::new(DEFAULT_MAX_BODY_BYTES),
            max_raw: ByteCount::new(DEFAULT_MAX_RAW_BYTES),
            max_ratio: CompressionRatio::DEFAULT,
            min_saved: ByteCount::new(DEFAULT_MIN_SAVED_BYTES),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WireFrame(bytes::Bytes);

impl WireFrame {
    pub fn from_bytes(bytes: bytes::Bytes) -> Self {
        Self(bytes)
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_ref()
    }

    pub fn into_bytes(self) -> bytes::Bytes {
        self.0
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WireCodec {
    limits: WireLimits,
}

impl WireCodec {
    pub fn new(limits: WireLimits) -> Self {
        Self { limits }
    }

    pub fn encode<T>(&self, value: &T) -> Result<WireFrame, WireCodecError>
    where
        T: WirePayload + for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
    {
        self.encode_with(value, Compression::DEFAULT)
    }

    pub fn encode_with<T>(
        &self,
        value: &T,
        compression: Compression,
    ) -> Result<WireFrame, WireCodecError>
    where
        T: WirePayload + for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
    {
        self.encode_payload(value, compression)
    }

    pub fn to_bytes<T>(
        &self,
        value: T,
        compression: Option<Compression>,
    ) -> Result<bytes::Bytes, WireCodecError>
    where
        T: WireEncode,
    {
        value.encode_wire(self, compression.unwrap_or_default())
    }

    pub fn from_bytes<T>(&self, bytes: impl AsRef<[u8]>) -> Result<T, WireCodecError>
    where
        T: WireDecode,
    {
        T::decode_wire(self, bytes.as_ref())
    }

    fn encode_payload<T>(
        &self,
        value: &T,
        compression: Compression,
    ) -> Result<WireFrame, WireCodecError>
    where
        T: WirePayload + for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
    {
        value.validate_wire()?;
        let raw = rkyv::to_bytes::<RkyvError>(value).map_err(WireCodecError::rkyv)?;
        let body = compression::encode(raw.as_slice(), compression, self.limits)?;
        FrameHeader::write(T::schema(), raw.len(), body, self.limits)
    }

    pub fn decode<T>(&self, frame: &WireFrame) -> Result<T, WireCodecError>
    where
        T: WirePayload + Archive,
        T::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
            + Deserialize<T, HighDeserializer<RkyvError>>,
    {
        self.decode_payload(frame.as_bytes())
    }

    fn decode_payload<T>(&self, bytes: &[u8]) -> Result<T, WireCodecError>
    where
        T: WirePayload + Archive,
        T::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
            + Deserialize<T, HighDeserializer<RkyvError>>,
    {
        let (header, body) = FrameHeader::read(bytes, self.limits)?;
        header.expect_schema(T::schema(), T::name())?;
        let raw = compression::decode(header, body, self.limits)?;
        let value = rkyv::from_bytes::<T, RkyvError>(&raw).map_err(WireCodecError::rkyv)?;
        value.validate_wire()?;
        Ok(value)
    }
}

impl Default for WireCodec {
    fn default() -> Self {
        Self::new(WireLimits::default())
    }
}

pub fn to_bytes<T>(
    value: T,
    compression: Option<Compression>,
) -> Result<bytes::Bytes, WireCodecError>
where
    T: WireEncode,
{
    WireCodec::default().to_bytes(value, compression)
}

pub fn from_bytes<T>(bytes: impl AsRef<[u8]>) -> Result<T, WireCodecError>
where
    T: WireDecode,
{
    WireCodec::default().from_bytes(bytes)
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WireCodecError {
    reason: String,
    suggestion: String,
    remote: Option<Box<WireError>>,
}

impl WireCodecError {
    pub fn new(reason: impl Into<String>, suggestion: impl Into<String>) -> Self {
        Self {
            reason: reason.into(),
            suggestion: suggestion.into(),
            remote: None,
        }
    }

    pub fn remote(error: WireError) -> Self {
        Self {
            reason: format!("remote wire call failed: {}", error.message),
            suggestion: error.help.to_string(),
            remote: Some(Box::new(error)),
        }
    }

    pub fn invalid_value(error: efdn_types::TypeError) -> Self {
        Self::new(error.to_string(), error.suggestion().to_owned())
    }

    pub fn rkyv(error: impl fmt::Display) -> Self {
        Self::new(
            format!("message bytes failed rkyv validation: {error}"),
            "send bytes produced by efdn wire for this message type",
        )
    }

    pub fn reason(&self) -> &str {
        self.reason.as_str()
    }

    pub fn suggestion(&self) -> &str {
        self.suggestion.as_str()
    }

    pub fn is_remote(&self) -> bool {
        self.remote.is_some()
    }

    pub fn remote_wire_error(&self) -> Option<&WireError> {
        self.remote.as_deref()
    }
}

impl fmt::Display for WireCodecError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}\nplease try: {}", self.reason, self.suggestion)
    }
}

impl std::error::Error for WireCodecError {}

#[cfg(feature = "errors")]
impl miette::Diagnostic for WireCodecError {
    fn code<'a>(&'a self) -> Option<Box<dyn fmt::Display + 'a>> {
        Some(Box::new("wire::codec"))
    }

    fn help<'a>(&'a self) -> Option<Box<dyn fmt::Display + 'a>> {
        Some(Box::new(self.suggestion()))
    }
}

#[diagnostic::on_unimplemented(
    message = "`{Self}` is not an efdn wire value",
    label = "missing efdn wire support",
    note = "please try: add #[efdn::wire::message] to the payload type"
)]
pub trait WireValidate {
    fn validate_wire(&self) -> Result<(), WireCodecError>;
}

#[diagnostic::on_unimplemented(
    message = "`{Self}` is not an efdn wire payload",
    label = "missing #[efdn::wire::message]",
    note = "please try: add #[efdn::wire::message] to request, reply, frame, or event types"
)]
pub trait WirePayload: WireValidate {
    fn name() -> MessageName;
    fn schema() -> MessageSchema;
    fn kind() -> MessageKind;
}

#[doc(hidden)]
pub trait WireEncode {
    fn encode_wire(
        self,
        codec: &WireCodec,
        compression: Compression,
    ) -> Result<bytes::Bytes, WireCodecError>;
}

#[doc(hidden)]
pub trait WireDecode: Sized {
    fn decode_wire(codec: &WireCodec, bytes: &[u8]) -> Result<Self, WireCodecError>;
}

const fn schema_id(value: &str) -> u64 {
    let bytes = value.as_bytes();
    let mut hash = 0xcbf29ce484222325u64;
    let mut index = 0;
    while index < bytes.len() {
        hash ^= bytes[index] as u64;
        hash = hash.wrapping_mul(0x100000001b3);
        index += 1;
    }
    hash
}

impl<T: WirePayload> WirePayload for Reply<T> {
    fn name() -> MessageName {
        MessageName::new("Reply")
    }

    fn schema() -> MessageSchema {
        T::schema().combine(MessageName::new("Reply"))
    }

    fn kind() -> MessageKind {
        T::kind()
    }
}

impl<T: WirePayload> WirePayload for In<T> {
    fn name() -> MessageName {
        MessageName::new("In")
    }

    fn schema() -> MessageSchema {
        T::schema().combine(MessageName::new("In"))
    }

    fn kind() -> MessageKind {
        T::kind()
    }
}

impl<T: WirePayload> WirePayload for Out<T> {
    fn name() -> MessageName {
        MessageName::new("Out")
    }

    fn schema() -> MessageSchema {
        T::schema().combine(MessageName::new("Out"))
    }

    fn kind() -> MessageKind {
        T::kind()
    }
}

impl WirePayload for Flow {
    fn name() -> MessageName {
        MessageName::new("Flow")
    }

    fn schema() -> MessageSchema {
        MessageSchema::from_name(Self::name())
    }

    fn kind() -> MessageKind {
        MessageKind::Frame
    }
}

impl WirePayload for Ack {
    fn name() -> MessageName {
        MessageName::new("Ack")
    }

    fn schema() -> MessageSchema {
        MessageSchema::from_name(Self::name())
    }

    fn kind() -> MessageKind {
        MessageKind::Frame
    }
}

impl<T> WireEncode for &T
where
    T: WirePayload + for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
{
    fn encode_wire(
        self,
        codec: &WireCodec,
        compression: Compression,
    ) -> Result<bytes::Bytes, WireCodecError> {
        codec
            .encode_payload(self, compression)
            .map(WireFrame::into_bytes)
    }
}

impl<T, E> WireEncode for Result<T, E>
where
    T: WirePayload,
    E: Into<WireError>,
    Reply<T>:
        WirePayload + for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
{
    fn encode_wire(
        self,
        codec: &WireCodec,
        compression: Compression,
    ) -> Result<bytes::Bytes, WireCodecError> {
        if T::kind() != MessageKind::Rep {
            return Err(WireCodecError::new(
                "wire result replies can only encode Rep payloads",
                "return Result<SomethingRep, E> or encode a plain payload by reference",
            ));
        }
        codec
            .encode_payload(&Reply::from(self), compression)
            .map(WireFrame::into_bytes)
    }
}

impl<T> WireDecode for T
where
    T: WirePayload + Archive,
    T::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
        + Deserialize<T, HighDeserializer<RkyvError>>,
    Reply<T>: WirePayload + Archive,
    <Reply<T> as Archive>::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
        + Deserialize<Reply<T>, HighDeserializer<RkyvError>>,
{
    fn decode_wire(codec: &WireCodec, bytes: &[u8]) -> Result<Self, WireCodecError> {
        let (header, _) = FrameHeader::read(bytes, codec.limits)?;
        if header.schema() == T::schema() {
            return codec.decode_payload(bytes);
        }
        if T::kind() == MessageKind::Rep && header.schema() == Reply::<T>::schema() {
            return match codec.decode_payload::<Reply<T>>(bytes)? {
                Reply::Ok(value) => Ok(value),
                Reply::Err(error) => Err(WireCodecError::remote(error)),
            };
        }
        header.expect_schema(T::schema(), T::name())?;
        unreachable!("schema mismatch must return before this point")
    }
}
