use core::fmt;
use std::hash::{Hash, Hasher};

use anyhow::{Context, Result, bail};
use base64::Engine;
use hex;
use serde::{Deserialize, de::Error};

pub type RawKey = Vec<u8>;

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct ShardingKey {
    key: Vec<u8>,
}

impl<'de> Deserialize<'de> for ShardingKey {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        ShardingKey::from_base64(&s).map_err(D::Error::custom)
    }
}

impl std::fmt::Display for ShardingKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.hex())
    }
}

impl std::fmt::Debug for ShardingKey {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self) // Use Display
    }
}

impl Hash for ShardingKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.key.hash(state);
    }
}

impl ShardingKey {
    pub fn from_base64(base64_key: &str) -> Result<Self> {
        let key = base64::engine::general_purpose::STANDARD
            .decode(base64_key)
            .with_context(|| format!("Base64 decode failed, key={}", base64_key))?;
        Ok(ShardingKey { key })
    }

    pub fn from_raw(raw_key: Vec<u8>) -> Self {
        ShardingKey { key: raw_key }
    }

    pub fn from_hex(hex_key: &str) -> Result<Self> {
        let key =
            hex::decode(hex_key).with_context(|| format!("Hex decode failed, key={}", hex_key))?;
        Ok(ShardingKey { key })
    }

    pub fn hex(&self) -> String {
        hex::encode(&self.key)
    }

    pub fn raw(&self) -> &[u8] {
        &self.key
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.key
    }

    pub fn next_lexicographic(v: &[u8]) -> Vec<u8> {
        let mut res = v.to_vec();
        for i in (0..res.len()).rev() {
            if res[i] < 0xFF {
                res[i] += 1;
                res.truncate(i + 1);
                return res;
            }
        }
        // All bytes are 0xFF, append a 0x00
        res.push(0x00);
        res
    }

    #[cfg(test)]
    pub fn base64(&self) -> String {
        base64::engine::general_purpose::STANDARD.encode(&self.key)
    }
}

#[derive(Clone)]
pub struct ShardingKeyRange {
    start: ShardingKey,
    end: ShardingKey,
}

impl ShardingKeyRange {
    // Create range `[start, end)`
    // If start == end, it is a empty range.
    pub fn half_open(start: Vec<u8>, end: Vec<u8>) -> Result<Self> {
        if start > end {
            bail!(
                "Invalid key range: start cannot be greater than end={} > {}",
                hex::encode(start),
                hex::encode(end)
            );
        }
        Ok(ShardingKeyRange {
            start: ShardingKey::from_raw(start),
            end: ShardingKey::from_raw(end),
        })
    }

    pub fn contains(&self, key: &[u8]) -> bool {
        key >= self.start.raw() && key < self.end.raw()
    }

    pub fn intersect(&self, range: &ShardingKeyRange) -> Self {
        let intersect_start = std::cmp::max(&self.start, &range.start);
        let intersect_end = std::cmp::min(&self.end, &range.end);
        if intersect_start >= intersect_end {
            // empty intersection
            return ShardingKeyRange {
                start: intersect_start.clone(),
                end: intersect_start.clone(),
            };
        }
        ShardingKeyRange {
            start: intersect_start.clone(),
            end: intersect_end.clone(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.start == self.end
    }

    // Check if self is a subset of range.
    pub fn is_subset_of(&self, range: &ShardingKeyRange) -> bool {
        // In theory, empty set is a subset of any set. But in our scenario, subset
        // implies that there is no need to check whether the key meets the
        // requirements of the read ranges, so false is returned here.
        if self.is_empty() {
            return false;
        }
        // `self` is not empty set. It is not subset of empty set.
        if range.is_empty() {
            return false;
        }
        range.start <= self.start && self.end <= range.end
    }

    pub fn is_subset_of_any(&self, ranges: &[ShardingKeyRange]) -> bool {
        ranges.iter().any(|range| self.is_subset_of(range))
    }

    // Check if self is disjoint with range.
    pub fn is_disjoint_with(&self, range: &ShardingKeyRange) -> bool {
        // Empty set is disjoint with any range.
        if self.is_empty() || range.is_empty() {
            return true;
        }
        self.start >= range.end || self.end <= range.start
    }

    pub fn is_disjoint_with_all(&self, ranges: &[ShardingKeyRange]) -> bool {
        ranges.iter().all(|range| self.is_disjoint_with(range))
    }

    // Check if self is intersect with range.
    pub fn is_intersect_with(&self, range: &ShardingKeyRange) -> bool {
        !self.is_disjoint_with(range)
    }

    pub fn start(&self) -> &ShardingKey {
        &self.start
    }

    pub fn end(&self) -> &ShardingKey {
        &self.end
    }
}

impl std::fmt::Debug for ShardingKeyRange {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "[{}, {})", self.start, self.end)
    }
}

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

    #[test]
    fn test_key_from_and_to() {
        let raw = b"t_123_abc";
        let h = ShardingKey::from_raw(raw.to_vec());
        assert_eq!(h.raw(), raw);

        let hex_str = h.hex();
        let h2 = ShardingKey::from_hex(&hex_str).unwrap();
        assert_eq!(h2.raw(), raw);

        let base64_str = h.base64();
        let h3 = ShardingKey::from_base64(&base64_str).unwrap();
        assert_eq!(h3.raw(), raw);
    }

    #[test]
    fn test_key_display_and_debug() {
        let raw = b"hello";
        let h = ShardingKey::from_raw(raw.to_vec());
        // hex of "hello" = 68656c6c6f
        assert_eq!(h.to_string(), "68656c6c6f");
        assert_eq!(format!("{:?}", h), "68656c6c6f");
    }

    #[test]
    fn test_key_eq_ord_hash() {
        use std::collections::hash_map::DefaultHasher;

        let h1 = ShardingKey::from_raw(b"a".to_vec());
        let h2 = ShardingKey::from_raw(b"a".to_vec());
        let h3 = ShardingKey::from_raw(b"b".to_vec());

        assert_eq!(h1, h2);
        assert!(h1 < h3);

        let mut hasher1 = DefaultHasher::new();
        h1.hash(&mut hasher1);
        let mut hasher2 = DefaultHasher::new();
        h2.hash(&mut hasher2);
        assert_eq!(hasher1.finish(), hasher2.finish());
    }

    #[test]
    fn test_key_range_basic() {
        let r = ShardingKeyRange::half_open(b"a".to_vec(), b"z".to_vec()).unwrap();
        assert!(r.contains(b"a"));
        assert!(r.contains(b"y"));
        assert!(!r.contains(b"z")); // end is exclusive
    }

    #[test]
    fn test_key_range_invalid() {
        let err = ShardingKeyRange::half_open(b"z".to_vec(), b"a".to_vec());
        assert!(err.is_err());
    }

    #[test]
    fn test_key_range_empty() {
        let r = ShardingKeyRange::half_open(b"".to_vec(), b"".to_vec()).unwrap();
        assert!(!r.contains(b"anything"));
        assert_eq!(r.start().raw(), b"");
        assert_eq!(r.end().raw(), b"");
    }

    #[test]
    fn test_intersect_non_empty() {
        // [a, m) ∩ [f, z) = [f, m)
        let r1 = ShardingKeyRange::half_open(b"a".to_vec(), b"m".to_vec()).unwrap();
        let r2 = ShardingKeyRange::half_open(b"f".to_vec(), b"z".to_vec()).unwrap();
        let inter = r1.intersect(&r2);
        assert_eq!(inter.start().raw(), b"f");
        assert_eq!(inter.end().raw(), b"m");
    }

    #[test]
    fn test_intersect_empty() {
        // [a, c) ∩ [x, z) = ∅
        let r1 = ShardingKeyRange::half_open(b"a".to_vec(), b"c".to_vec()).unwrap();
        let r2 = ShardingKeyRange::half_open(b"x".to_vec(), b"z".to_vec()).unwrap();
        let inter = r1.intersect(&r2);
        assert_eq!(inter.start().raw(), b"x");
        assert_eq!(inter.end().raw(), b"x");
    }

    #[test]
    fn test_subset_and_disjoint() {
        let a = ShardingKeyRange::half_open(b"a".to_vec(), b"d".to_vec()).unwrap();
        let b = ShardingKeyRange::half_open(b"b".to_vec(), b"c".to_vec()).unwrap();
        let c = ShardingKeyRange::half_open(b"d".to_vec(), b"f".to_vec()).unwrap();

        assert!(b.is_subset_of(&a));
        assert!(b.is_subset_of_any(&[c.clone(), a.clone()]));
        assert!(!a.is_subset_of(&b));

        assert!(a.is_disjoint_with(&c));
        assert!(c.is_disjoint_with(&a));
        assert!(!a.is_disjoint_with(&b));
        assert!(b.is_disjoint_with_all(&[c.clone()]));
        assert!(!b.is_disjoint_with_all(&[a.clone(), c.clone()]));
    }

    #[test]
    fn test_next_lexicographic_basic() {
        // Normal case: simply increment the last byte
        assert_eq!(ShardingKey::next_lexicographic(b"abc"), b"abd");

        // The last byte is 0xFF → carry and truncate
        assert_eq!(ShardingKey::next_lexicographic(&[0x01, 0xFF]), vec![0x02]);

        // Multiple trailing 0xFF bytes → propagate carry
        assert_eq!(
            ShardingKey::next_lexicographic(&[0x01, 0x02, 0xFF, 0xFF]),
            vec![0x01, 0x03]
        );

        // All bytes are 0xFF → append a 0x00
        assert_eq!(
            ShardingKey::next_lexicographic(&[0xFF, 0xFF]),
            vec![0xFF, 0xFF, 0x00]
        );

        // Empty input → return [0x00]
        assert_eq!(ShardingKey::next_lexicographic(&[]), vec![0x00]);
    }

    #[test]
    fn test_next_lexicographic_monotonicity() {
        // Verify that the result is always lexicographically greater than the input
        let samples = vec![
            b"a".to_vec(),
            b"\x00\x7F".to_vec(),
            b"\x00\xFF".to_vec(),
            b"\xFF\xFF".to_vec(),
        ];
        for s in samples {
            let next = ShardingKey::next_lexicographic(&s);
            assert!(s < next, "expected {:?} < {:?}", s, next);
        }
    }

    #[test]
    fn test_next_lexicographic_truncation() {
        // Check truncation behavior: 0x01 0xFE → 0x01 0xFF
        assert_eq!(
            ShardingKey::next_lexicographic(&[0x01, 0xFE]),
            vec![0x01, 0xFF]
        );

        // 0x01 0xFF → 0x02 (the tail is truncated)
        assert_eq!(ShardingKey::next_lexicographic(&[0x01, 0xFF]), vec![0x02]);
    }

    #[test]
    fn test_subset_and_disjoint_empty_cases() {
        let empty = ShardingKeyRange::half_open(b"".to_vec(), b"".to_vec()).unwrap();
        let a = ShardingKeyRange::half_open(b"a".to_vec(), b"z".to_vec()).unwrap();

        // Empty set is not a subset of any in our code.
        assert!(!empty.is_subset_of(&a));
        assert!(!empty.is_subset_of(&empty));

        // Non-empty set is not a subset of empty set
        assert!(!a.is_subset_of(&empty));

        // Empty set is disjoint with any set
        assert!(empty.is_disjoint_with(&a));
        assert!(empty.is_disjoint_with(&empty));
        assert!(a.is_disjoint_with(&empty));
    }

    #[test]
    fn test_intersect_with_empty() {
        let empty = ShardingKeyRange::half_open(b"".to_vec(), b"".to_vec()).unwrap();
        let a = ShardingKeyRange::half_open(b"a".to_vec(), b"m".to_vec()).unwrap();

        // Empty ∩ [a,m) = empty
        let inter = empty.intersect(&a);
        assert!(inter.is_empty());

        // [a,m) ∩ empty = empty
        let inter2 = a.intersect(&empty);
        assert!(inter2.is_empty());

        // empty ∩ empty = empty
        let inter3 = empty.intersect(&empty);
        assert!(inter3.is_empty());
    }

    #[test]
    fn test_is_intersect_with() {
        let a = ShardingKeyRange::half_open(b"a".to_vec(), b"d".to_vec()).unwrap();
        let b = ShardingKeyRange::half_open(b"b".to_vec(), b"c".to_vec()).unwrap();
        let c = ShardingKeyRange::half_open(b"d".to_vec(), b"e".to_vec()).unwrap();
        let empty = ShardingKeyRange::half_open(b"".to_vec(), b"".to_vec()).unwrap();

        assert!(a.is_intersect_with(&b)); // overlap
        assert!(!a.is_intersect_with(&c)); // touch at boundary only
        assert!(!a.is_intersect_with(&empty)); // empty range never intersects
    }

    #[test]
    fn test_intersect_touching() {
        // [a, b) and [b, c) share boundary but no overlap
        let r1 = ShardingKeyRange::half_open(b"a".to_vec(), b"b".to_vec()).unwrap();
        let r2 = ShardingKeyRange::half_open(b"b".to_vec(), b"c".to_vec()).unwrap();
        assert!(r1.is_disjoint_with(&r2));
        let inter = r1.intersect(&r2);
        assert!(inter.is_empty());
    }
}