wedb_embed 0.1.1

pub mod codec;
pub mod conf;

pub use codec::{
    BASE32_ALPHABET, BASE32_DECODE_TABLE, D_R, EARTH_RADIUS_METERS, GEO_LAT_MAX, GEO_LAT_MIN,
    GEO_LAT_RANGE, GEO_LAT_RANGE_STANDARD, GEO_LON_MAX, GEO_LON_MIN, GEO_LON_RANGE,
    GEO_LON_RANGE_STANDARD, GEO_STEP_MAX, MERCATOR_MAX, align_52bits, base32_to_coords,
    bounding_box, convert_meters_to_unit, convert_unit_to_meters, coords_to_base32, decode_geohash,
    deinterleave64, encode_geohash, encode_geohash_bytes, encode_geohash_string,
    estimate_steps_by_radius, geohash_decode, geohash_decode_area_to_long_lat,
    geohash_decode_wgs84, geohash_encode, geohash_encode_wgs84, geohash_move_x, geohash_move_y,
    geohash_neighbors, geohash_to_base32, get_areas_by_shape_wgs84, haversine_distance,
    interleave64, scores_of_geohash_box, validate_long_lat,
};
pub use conf::{
    DistanceSort, DistanceUnit, GeoHashArea, GeoHashBits, GeoHashNeighbors, GeoHashRadius,
    GeoHashRange, GeoPoint, GeoRadiusOption, GeoSearch, GeoSearchStoreOption, GeoShape,
    GeoShapeType, OriginPoint,
};

use crate::db::WeDb;
use crate::error::{Error, Result};
use crate::zset::conf::{RangeScoreSpec, ZAdd};

#[inline]
fn sort_and_truncate_points(
    points: &mut Vec<GeoPoint>,
    mut sort: DistanceSort,
    count: Option<usize>,
    any: bool,
) {
    if count.is_some() && !any && sort == DistanceSort::None {
        sort = DistanceSort::Asc;
    }

    if let Some(limit) = count {
        if limit == 0 {
            points.clear();
            return;
        }
        match sort {
            DistanceSort::Asc => {
                if limit < points.len() {
                    points.select_nth_unstable_by(limit, |a, b| a.dist.total_cmp(&b.dist));
                    points.truncate(limit);
                }
                points.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist));
            }
            DistanceSort::Desc => {
                if limit < points.len() {
                    points.select_nth_unstable_by(limit, |a, b| b.dist.total_cmp(&a.dist));
                    points.truncate(limit);
                }
                points.sort_unstable_by(|a, b| b.dist.total_cmp(&a.dist));
            }
            DistanceSort::None => {
                points.truncate(limit);
            }
        }
    } else {
        match sort {
            DistanceSort::Asc => {
                points.sort_unstable_by(|a, b| a.dist.total_cmp(&b.dist));
            }
            DistanceSort::Desc => {
                points.sort_unstable_by(|a, b| b.dist.total_cmp(&a.dist));
            }
            DistanceSort::None => {}
        }
    }
}

impl WeDb {
    /// GEOADD key longitude latitude member [longitude latitude member ...]
    pub fn geoadd<K: AsRef<[u8]>, M: AsRef<[u8]>>(
        &self,
        key: K,
        items: &[(f64, f64, M)],
    ) -> Result<usize> {
        self.geoadd_opts(key, items, [])
    }

    /// GEOADD 带选项版本（支持 NX, XX, CH 等，对标 Kvrocks Geo::Add）
    pub fn geoadd_opts<K: AsRef<[u8]>, M: AsRef<[u8]>>(
        &self,
        key: K,
        items: &[(f64, f64, M)],
        options: impl AsRef<[ZAdd]>,
    ) -> Result<usize> {
        let mut score_members = Vec::with_capacity(items.len());
        for (lon, lat, member) in items {
            validate_long_lat(*lon, *lat)?;
            let bits = geohash_encode_wgs84(*lon, *lat, GEO_STEP_MAX)
                .ok_or_else(|| Error::invalid_data("ERR invalid longitude/latitude coordinates"))?;
            let score = align_52bits(bits) as f64;
            score_members.push((score, member));
        }
        self.zadd(key, &score_members, options)
    }

    /// GEODIST key member1 member2 [unit]（对标 Kvrocks Geo::Dist）
    pub fn geodist<K: AsRef<[u8]>, M1: AsRef<[u8]>, M2: AsRef<[u8]>>(
        &self,
        key: K,
        member1: M1,
        member2: M2,
        unit: Option<&str>,
    ) -> Result<Option<f64>> {
        let u = match unit {
            Some(s) => DistanceUnit::parse(s).ok_or_else(|| {
                Error::invalid_data("ERR unsupported unit provided. please use M, KM, FT, MI")
            })?,
            None => DistanceUnit::Meters,
        };

        let scores = self.zmscore(&key, &[member1.as_ref(), member2.as_ref()])?;
        if let (Some(Some(s1)), Some(Some(s2))) = (scores.first().copied(), scores.get(1).copied())
        {
            let (lon1, lat1) = geohash_decode_wgs84(GeoHashBits {
                bits: s1 as u64,
                step: GEO_STEP_MAX,
            });
            let (lon2, lat2) = geohash_decode_wgs84(GeoHashBits {
                bits: s2 as u64,
                step: GEO_STEP_MAX,
            });
            let d_meters = haversine_distance(lon1, lat1, lon2, lat2);
            Ok(Some(u.from_meters(d_meters)))
        } else {
            Ok(None)
        }
    }

    /// GEOPOS key member [member ...]（对标 Kvrocks Geo::Pos）
    pub fn geopos<K: AsRef<[u8]>, M: AsRef<[u8]>>(
        &self,
        key: K,
        members: &[M],
    ) -> Result<Vec<Option<(f64, f64)>>> {
        let scores = self.zmscore(&key, members)?;
        let results = scores
            .into_iter()
            .map(|score| {
                score.map(|s| {
                    let bits = GeoHashBits {
                        bits: s as u64,
                        step: GEO_STEP_MAX,
                    };
                    geohash_decode_wgs84(bits)
                })
            })
            .collect();
        Ok(results)
    }

    /// GEOHASH key member [member ...]（对标 Kvrocks Geo::Hash）
    pub fn geohash<K: AsRef<[u8]>, M: AsRef<[u8]>>(
        &self,
        key: K,
        members: &[M],
    ) -> Result<Vec<Option<String>>> {
        let scores = self.zmscore(&key, members)?;
        let results = scores
            .into_iter()
            .map(|score| {
                score.map(|s| {
                    let bits = GeoHashBits {
                        bits: s as u64,
                        step: GEO_STEP_MAX,
                    };
                    let (lon, lat) = geohash_decode_wgs84(bits);
                    encode_geohash_string(lon, lat)
                })
            })
            .collect();
        Ok(results)
    }

    /// GEORADIUS key longitude latitude radius m|km|ft|mi [WITHCOORD] [WITHDIST] [WITHHASH] [COUNT count [ANY]] [ASC|DESC] [STORE key] [STOREDIST key]（对标 Kvrocks Geo::Radius）
    pub fn georadius<K: AsRef<[u8]>>(
        &self,
        key: K,
        longitude: f64,
        latitude: f64,
        radius: f64,
        opt: &GeoRadiusOption,
    ) -> Result<Vec<GeoPoint>> {
        validate_long_lat(longitude, latitude)?;
        if radius.is_nan() || radius.is_infinite() || radius < 0.0 {
            return Err(Error::invalid_data(
                "ERR radius must be greater than or equal to 0",
            ));
        }

        if self.zcard(key.as_ref())? == 0 {
            if let Some(ref store_k) = opt.store_key {
                self.del(&[store_k])?;
            }
            if let Some(ref store_dist_k) = opt.store_dist_key {
                self.del(&[store_dist_k])?;
            }
            return Ok(Vec::new());
        }

        let mut shape = GeoShape::new_circular_with_unit(longitude, latitude, radius, opt.unit);
        let mut points = self.search_shape(&key, &mut shape, opt.unit)?;

        sort_and_truncate_points(&mut points, opt.sort, opt.count, opt.any);

        if let Some(ref store_k) = opt.store_key {
            self.del(&[store_k])?;
            if !points.is_empty() {
                let items: Vec<(f64, &str)> = points
                    .iter()
                    .map(|p| (p.score, p.member.as_str()))
                    .collect();
                self.zadd(store_k, &items, [])?;
            }
        }

        if let Some(ref store_dist_k) = opt.store_dist_key {
            self.del(&[store_dist_k])?;
            if !points.is_empty() {
                let items: Vec<(f64, &str)> =
                    points.iter().map(|p| (p.dist, p.member.as_str())).collect();
                self.zadd(store_dist_k, &items, [])?;
            }
        }

        Ok(points)
    }

    /// GEORADIUSBYMEMBER key member radius m|km|ft|mi [WITHCOORD] [WITHDIST] [WITHHASH] [COUNT count [ANY]] [ASC|DESC] [STORE key] [STOREDIST key]（对标 Kvrocks Geo::RadiusByMember）
    pub fn georadiusbymember<K: AsRef<[u8]>, M: AsRef<[u8]>>(
        &self,
        key: K,
        member: M,
        radius: f64,
        opt: &GeoRadiusOption,
    ) -> Result<Vec<GeoPoint>> {
        let k_ref = key.as_ref();
        let m_ref = member.as_ref();

        // 检查 key 是否存在（对标 Kvrocks: key 不存在时清理目标 key 并返回空列表）
        if self.zcard(k_ref)? == 0 {
            if let Some(ref store_k) = opt.store_key {
                self.del(&[store_k])?;
            }
            if let Some(ref store_dist_k) = opt.store_dist_key {
                self.del(&[store_dist_k])?;
            }
            return Ok(Vec::new());
        }

        let pos = self.geopos(k_ref, &[m_ref])?;
        if let Some(Some((lon, lat))) = pos.into_iter().next() {
            self.georadius(k_ref, lon, lat, radius, opt)
        } else {
            Err(Error::invalid_data(
                "ERR could not decode requested zset member",
            ))
        }
    }

    /// GEOSEARCH key [FROMMEMBER member | FROMLONLAT lon lat] [BYRADIUS radius m|km|ft|mi | BYBOX width height m|km|ft|mi] [ASC|DESC] [COUNT count [ANY]] [WITHCOORD] [WITHDIST] [WITHHASH]（对标 Kvrocks Geo::Search）
    pub fn geosearch<K: AsRef<[u8]>>(
        &self,
        key: K,
        origin: &OriginPoint,
        shape: &mut GeoShape,
        opt: &GeoSearch,
    ) -> Result<Vec<GeoPoint>> {
        let k_ref = key.as_ref();
        if self.zcard(k_ref)? == 0 {
            return Ok(Vec::new());
        }

        let (lon, lat) = match origin {
            OriginPoint::Coord { lon, lat } => {
                validate_long_lat(*lon, *lat)?;
                (*lon, *lat)
            }
            OriginPoint::Member(m) => {
                let pos = self.geopos(k_ref, &[m.as_bytes()])?;
                match pos.into_iter().next().flatten() {
                    Some((l, t)) => (l, t),
                    None => {
                        return Err(Error::invalid_data(
                            "ERR could not decode requested zset member",
                        ));
                    }
                }
            }
        };

        shape.center_lon = lon;
        shape.center_lat = lat;
        bounding_box(shape);

        let mut points = self.search_shape(k_ref, shape, opt.unit)?;

        let sort_order = if opt.asc || opt.sort == DistanceSort::Asc {
            DistanceSort::Asc
        } else {
            opt.sort
        };

        sort_and_truncate_points(&mut points, sort_order, opt.count, opt.any);

        Ok(points)
    }

    /// GEOSEARCHSTORE destination source [FROMMEMBER member | FROMLONLAT lon lat] [BYRADIUS radius m|km|ft|mi | BYBOX width height m|km|ft|mi] [ASC|DESC] [COUNT count [ANY]] [STOREDIST]（对标 Kvrocks Geo::SearchStore）
    pub fn geosearchstore<K: AsRef<[u8]>, D: AsRef<[u8]>>(
        &self,
        destination: D,
        source: K,
        origin: &OriginPoint,
        shape: &mut GeoShape,
        opt: &GeoSearchStoreOption,
    ) -> Result<usize> {
        let dest_ref = destination.as_ref();
        let src_ref = source.as_ref();

        let (lon, lat) = match origin {
            OriginPoint::Coord { lon, lat } => {
                validate_long_lat(*lon, *lat)?;
                if self.zcard(src_ref)? == 0 {
                    self.del(&[dest_ref])?;
                    return Ok(0);
                }
                (*lon, *lat)
            }
            OriginPoint::Member(m) => {
                if self.zcard(src_ref)? == 0 {
                    self.del(&[dest_ref])?;
                    return Ok(0);
                }
                let pos = self.geopos(src_ref, &[m.as_bytes()])?;
                match pos.into_iter().next().flatten() {
                    Some((l, t)) => (l, t),
                    None => {
                        // member 不存在时报错，且不影响 destination
                        return Err(Error::invalid_data(
                            "ERR could not decode requested zset member",
                        ));
                    }
                }
            }
        };

        shape.center_lon = lon;
        shape.center_lat = lat;
        bounding_box(shape);

        let mut points = self.search_shape(src_ref, shape, opt.unit)?;

        sort_and_truncate_points(&mut points, opt.sort, opt.count, opt.any);

        self.del(&[dest_ref])?;
        if points.is_empty() {
            return Ok(0);
        }

        let items: Vec<(f64, &str)> = points
            .iter()
            .map(|p| {
                let score = if opt.store_dist { p.dist } else { p.score };
                (score, p.member.as_str())
            })
            .collect();

        self.zadd(dest_ref, &items, [])?;
        Ok(items.len())
    }

    /// 内部核心方法：扫描 GeoShape 外接的所有 Geohash 块并精确过滤（对标 Kvrocks membersOfAllNeighbors + appendIfWithinShape）
    fn search_shape<K: AsRef<[u8]>>(
        &self,
        key: K,
        geo_shape: &mut GeoShape,
        unit: DistanceUnit,
    ) -> Result<Vec<GeoPoint>> {
        if self.zcard(key.as_ref())? == 0 {
            return Ok(Vec::new());
        }

        let georadius = get_areas_by_shape_wgs84(geo_shape);
        let raw_neighbors = [
            georadius.hash,
            georadius.neighbors.north,
            georadius.neighbors.south,
            georadius.neighbors.east,
            georadius.neighbors.west,
            georadius.neighbors.north_east,
            georadius.neighbors.north_west,
            georadius.neighbors.south_east,
            georadius.neighbors.south_west,
        ];

        // 预先去重并过滤无效块（栈上固定数组，零堆分配）
        let mut unique_hashes = [GeoHashBits::default(); 9];
        let mut unique_count = 0usize;
        for hash in raw_neighbors {
            if hash.is_zero() {
                continue;
            }
            if !unique_hashes[..unique_count].contains(&hash) {
                unique_hashes[unique_count] = hash;
                unique_count += 1;
            }
        }

        let mut points = Vec::new();
        let center_lon = geo_shape.center_lon;
        let center_lat = geo_shape.center_lat;
        let max_radius_meters = geo_shape.radius * geo_shape.conversion;
        let bounds = geo_shape.bounds;
        let shape_type = geo_shape.shape_type;

        for &hash in &unique_hashes[..unique_count] {
            let (min_bits, max_bits) = scores_of_geohash_box(hash);
            let spec = RangeScoreSpec {
                min: min_bits as f64,
                max: max_bits as f64,
                minex: false,
                maxex: true, // [min, max)
                offset: 0,
                count: None,
            };

            let range_members = self.zrangebyscore(&key, &spec)?;
            for (member_bytes, score) in range_members {
                let bits = GeoHashBits {
                    bits: score as u64,
                    step: GEO_STEP_MAX,
                };
                let (pt_lon, pt_lat) = geohash_decode_wgs84(bits);
                let d_meters = haversine_distance(center_lon, center_lat, pt_lon, pt_lat);

                let is_inside = match shape_type {
                    GeoShapeType::Circular => d_meters <= max_radius_meters,
                    GeoShapeType::Rectangular => {
                        pt_lon >= bounds[0]
                            && pt_lon <= bounds[2]
                            && pt_lat >= bounds[1]
                            && pt_lat <= bounds[3]
                    }
                    GeoShapeType::None => false,
                };

                if is_inside {
                    let member_str = String::from_utf8(member_bytes)
                        .unwrap_or_else(|e| String::from_utf8_lossy(e.as_bytes()).into_owned());
                    let dist = unit.from_meters(d_meters);
                    points.push(GeoPoint {
                        longitude: pt_lon,
                        latitude: pt_lat,
                        member: member_str,
                        dist,
                        score,
                    });
                }
            }
        }

        Ok(points)
    }
}