/* 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 super::PartitioningKind;
use crate::{IggyByteSize, Sizeable, error::IggyError};
use serde::{Deserialize, Serialize};
use serde_with::base64::Base64;
use serde_with::serde_as;
use std::{
    fmt::Display,
    hash::{Hash, Hasher},
};

/// `Partitioning` is used to specify to which partition the messages should be sent.
/// It has the following kinds:
/// - `Balanced` - the partition ID is calculated by the server using the round-robin algorithm.
/// - `PartitionId` - the partition ID is provided by the client.
/// - `MessagesKey` - the partition ID is calculated by the server using the hash of the provided messages key.
#[serde_as]
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
pub struct Partitioning {
    /// The kind of partitioning.
    pub kind: PartitioningKind,
    #[serde(skip)]
    /// The length of the value payload.
    pub length: u8,
    #[serde_as(as = "Base64")]
    /// The binary value payload.
    pub value: Vec<u8>,
}

impl Display for Partitioning {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.kind {
            PartitioningKind::Balanced => write!(f, "{}|0", self.kind),
            PartitioningKind::PartitionId => {
                if let Some(bytes) = self.value.get(..4).and_then(|s| s.try_into().ok()) {
                    write!(f, "{}|{}", self.kind, u32::from_le_bytes(bytes))
                } else {
                    write!(f, "{}|<invalid>", self.kind)
                }
            }
            PartitioningKind::MessagesKey => {
                write!(f, "{}|{}", self.kind, String::from_utf8_lossy(&self.value))
            }
        }
    }
}

impl Partitioning {
    /// Partition the messages using the balanced round-robin algorithm on the server.
    pub fn balanced() -> Self {
        Partitioning {
            kind: PartitioningKind::Balanced,
            length: 0,
            value: vec![],
        }
    }

    /// Partition the messages using the provided partition ID.
    pub fn partition_id(partition_id: u32) -> Self {
        Partitioning {
            kind: PartitioningKind::PartitionId,
            length: 4,
            value: partition_id.to_le_bytes().to_vec(),
        }
    }

    /// Partition the messages using the provided messages key.
    pub fn messages_key(value: &[u8]) -> Result<Self, IggyError> {
        let length = value.len();
        if length == 0 || length > 255 {
            return Err(IggyError::InvalidCommand);
        }

        Ok(Partitioning {
            kind: PartitioningKind::MessagesKey,
            #[allow(clippy::cast_possible_truncation)]
            length: length as u8,
            value: value.to_vec(),
        })
    }

    /// Partition the messages using the provided messages key as str.
    pub fn messages_key_str(value: &str) -> Result<Self, IggyError> {
        Self::messages_key(value.as_bytes())
    }

    /// Partition the messages using the provided messages key as u32.
    pub fn messages_key_u32(value: u32) -> Self {
        Partitioning {
            kind: PartitioningKind::MessagesKey,
            length: 4,
            value: value.to_le_bytes().to_vec(),
        }
    }

    /// Partition the messages using the provided messages key as u64.
    pub fn messages_key_u64(value: u64) -> Self {
        Partitioning {
            kind: PartitioningKind::MessagesKey,
            length: 8,
            value: value.to_le_bytes().to_vec(),
        }
    }

    /// Partition the messages using the provided messages key as u128.
    pub fn messages_key_u128(value: u128) -> Self {
        Partitioning {
            kind: PartitioningKind::MessagesKey,
            length: 16,
            value: value.to_le_bytes().to_vec(),
        }
    }

    /// Create the partitioning from the provided partitioning.
    pub fn from_partitioning(partitioning: &Partitioning) -> Self {
        Partitioning {
            kind: partitioning.kind,
            length: partitioning.length,
            value: partitioning.value.clone(),
        }
    }

    /// Create the partitioning from a raw byte slice.
    pub fn from_raw_bytes(bytes: &[u8]) -> Result<Self, IggyError> {
        let kind = PartitioningKind::from_code(*bytes.first().ok_or(IggyError::InvalidCommand)?)?;
        let length = *bytes.get(1).ok_or(IggyError::InvalidCommand)?;
        if kind == PartitioningKind::PartitionId && length != 4 {
            return Err(IggyError::InvalidCommand);
        }
        if kind == PartitioningKind::Balanced && length != 0 {
            return Err(IggyError::InvalidCommand);
        }

        let value = bytes
            .get(2..2 + length as usize)
            .ok_or(IggyError::InvalidCommand)?
            .to_vec();

        Ok(Partitioning {
            kind,
            length,
            value,
        })
    }

    /// Maximum size of the Partitioning struct
    pub const fn maximum_byte_size() -> usize {
        2 + 255
    }
}

impl Hash for Partitioning {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.kind.hash(state);
        self.length.hash(state);
        self.value.hash(state);
    }
}

impl Default for Partitioning {
    fn default() -> Self {
        Partitioning::balanced()
    }
}

impl Sizeable for Partitioning {
    fn get_size_bytes(&self) -> IggyByteSize {
        IggyByteSize::from(u64::from(self.length) + 2)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn from_raw_bytes_should_reject_partition_id_with_wrong_length() {
        for bad_len in [0u8, 1, 2, 3, 5, 255] {
            let mut buf = vec![PartitioningKind::PartitionId.as_code(), bad_len];
            buf.extend(vec![0u8; bad_len as usize]);
            assert!(
                Partitioning::from_raw_bytes(&buf).is_err(),
                "expected error for PartitionId with length={bad_len}"
            );
        }
    }

    #[test]
    fn from_raw_bytes_should_accept_valid_partition_id() {
        let mut buf = vec![PartitioningKind::PartitionId.as_code(), 4];
        buf.extend(42u32.to_le_bytes());
        let p = Partitioning::from_raw_bytes(&buf).unwrap();
        assert_eq!(p.kind, PartitioningKind::PartitionId);
        assert_eq!(p.length, 4);
        assert_eq!(p.value, 42u32.to_le_bytes());
    }

    #[test]
    fn from_raw_bytes_should_fail_on_truncated_input() {
        let p = Partitioning::partition_id(99);
        let mut full = vec![p.kind.as_code(), p.length];
        full.extend_from_slice(&p.value);
        for i in 0..full.len() {
            assert!(
                Partitioning::from_raw_bytes(&full[..i]).is_err(),
                "expected error for truncation at byte {i}"
            );
        }
    }

    #[test]
    fn display_should_not_panic_on_malformed_partition_id() {
        let malformed = Partitioning {
            kind: PartitioningKind::PartitionId,
            length: 0,
            value: vec![],
        };
        let s = format!("{malformed}");
        assert!(s.contains("<invalid>"));
    }

    #[test]
    fn from_raw_bytes_should_reject_balanced_with_nonzero_length() {
        for bad_len in [1u8, 2, 4, 255] {
            let mut buf = vec![PartitioningKind::Balanced.as_code(), bad_len];
            buf.extend(vec![0u8; bad_len as usize]);
            assert!(
                Partitioning::from_raw_bytes(&buf).is_err(),
                "expected error for Balanced with length={bad_len}"
            );
        }
    }
}
