//! Per-tenant identity, configuration, rate limiting, and routing.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};

use dashmap::DashMap;

use crate::schema::EntityId;

use super::shard::{RegionId, ShardId};

/// Unique identifier for a tenant.
#[derive(
    Debug,
    Clone,
    Copy,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Default,
    serde::Serialize,
    serde::Deserialize,
)]
pub struct TenantId(pub u64);

impl TenantId {
    pub const DEFAULT: Self = Self(0);
}

impl std::fmt::Display for TenantId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "tenant:{}", self.0)
    }
}

/// Per-tenant configuration and resource limits.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TenantConfig {
    pub tenant_id: TenantId,
    pub max_signals_per_sec: Option<u32>,
    pub max_entities: Option<u64>,
    pub max_storage_bytes: Option<u64>,
    pub required_regions: Vec<RegionId>,
    pub label: String,
}

impl TenantConfig {
    /// Create a default (unlimited) tenant config.
    #[must_use]
    pub fn default_tenant() -> Self {
        Self {
            tenant_id: TenantId::DEFAULT,
            max_signals_per_sec: None,
            max_entities: None,
            max_storage_bytes: None,
            required_regions: vec![],
            label: "default".to_string(),
        }
    }

    /// Create an unlimited tenant config with a custom label.
    #[must_use]
    pub fn unlimited(id: TenantId, label: impl Into<String>) -> Self {
        Self {
            tenant_id: id,
            max_signals_per_sec: None,
            max_entities: None,
            max_storage_bytes: None,
            required_regions: vec![],
            label: label.into(),
        }
    }
}

/// Token-bucket rate limiter for per-tenant signal throughput.
///
/// Uses `Ordering::Relaxed` for the CAS loop because approximate enforcement
/// is acceptable for rate limiting: a brief over-admission of 1-2 signals under
/// contention is tolerable, and the token bucket self-corrects on the next
/// refill cycle. The alternative (`AcqRel`) would add unnecessary contention
/// on the hot signal-write path.
#[derive(Debug)]
pub struct TenantRateLimiter {
    tokens_bits: AtomicU64,
    refill_rate_per_ns: f64,
    max_tokens: f64,
    last_refill_ns: AtomicU64,
}

impl TenantRateLimiter {
    /// Create a new rate limiter allowing `max_signals_per_sec` signals per second.
    #[must_use]
    pub fn new(max_signals_per_sec: u32) -> Self {
        let rate_per_ns = f64::from(max_signals_per_sec) / 1_000_000_000.0;
        let max_tokens = f64::from(max_signals_per_sec) * 2.0;
        Self {
            tokens_bits: AtomicU64::new(max_tokens.to_bits()),
            refill_rate_per_ns: rate_per_ns,
            max_tokens,
            last_refill_ns: AtomicU64::new(crate::replication::now_ns()),
        }
    }

    /// Try to consume one token from the bucket.
    ///
    /// # Errors
    ///
    /// Returns `QuotaExceeded` if the token bucket is empty.
    #[allow(clippy::cast_precision_loss)] // ns elapsed fits in f64 mantissa
    pub fn try_acquire(&self) -> crate::Result<()> {
        loop {
            // Re-read time inside the CAS loop so retries use fresh values,
            // preventing double-counting of the refill amount under contention.
            let now = crate::replication::now_ns();
            let last = self.last_refill_ns.load(Ordering::Relaxed);
            let elapsed_ns = now.saturating_sub(last);
            let refill = elapsed_ns as f64 * self.refill_rate_per_ns;

            let current = f64::from_bits(self.tokens_bits.load(Ordering::Relaxed));
            let new_tokens = (current + refill).min(self.max_tokens);
            if new_tokens < 1.0 {
                return Err(crate::schema::TidalError::QuotaExceeded(
                    "signal rate limit exceeded".into(),
                ));
            }
            let new_bits = (new_tokens - 1.0).to_bits();
            let old_bits = current.to_bits();
            // Relaxed ordering: see struct-level doc comment.
            if self
                .tokens_bits
                .compare_exchange_weak(old_bits, new_bits, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                self.last_refill_ns.store(now, Ordering::Relaxed);
                return Ok(());
            }
        }
    }
}

/// Describes the shard topology for the cluster.
#[derive(Debug, Clone)]
pub struct ClusterTopology {
    pub shards: Vec<ShardAssignment>,
}

/// Maps a shard to its region.
#[derive(Debug, Clone, Copy)]
pub struct ShardAssignment {
    pub shard_id: ShardId,
    pub region_id: RegionId,
}

impl ClusterTopology {
    /// Create a single-shard, single-region topology.
    #[must_use]
    pub fn single() -> Self {
        Self {
            shards: vec![ShardAssignment {
                shard_id: ShardId::SINGLE,
                region_id: RegionId::SINGLE,
            }],
        }
    }
}

/// Routes entities to shards based on tenant configuration and jump consistent hash.
pub struct TenantRouter {
    tenants: DashMap<TenantId, TenantConfig>,
    rate_limiters: DashMap<TenantId, Arc<TenantRateLimiter>>,
    topology: Arc<RwLock<ClusterTopology>>,
    /// Tenants currently in dual-write migration mode: maps `TenantId` to
    /// `(source_shard, target_shard)`. During dual-write, write operations
    /// target both shards; reads target the source (old) shard until finalization.
    dual_write_map: DashMap<TenantId, (ShardId, ShardId)>,
    /// Explicit per-tenant shard pin. Set after migration finalization to lock
    /// all routing for a tenant to the new target shard regardless of jump hash.
    shard_pin_map: DashMap<TenantId, ShardId>,
}

impl TenantRouter {
    /// Create a new router with the given topology.
    #[must_use]
    pub fn new(topology: Arc<RwLock<ClusterTopology>>) -> Self {
        Self {
            tenants: DashMap::new(),
            rate_limiters: DashMap::new(),
            topology,
            dual_write_map: DashMap::new(),
            shard_pin_map: DashMap::new(),
        }
    }

    /// Register a tenant and its rate limiter (if configured).
    pub fn register_tenant(&self, config: TenantConfig) {
        if let Some(rate) = config.max_signals_per_sec {
            let limiter = Arc::new(TenantRateLimiter::new(rate));
            self.rate_limiters.insert(config.tenant_id, limiter);
        }
        self.tenants.insert(config.tenant_id, config);
    }

    /// Return the rate limiter for a tenant, if one is configured.
    #[must_use]
    pub fn rate_limiter_for(&self, tenant_id: TenantId) -> Option<Arc<TenantRateLimiter>> {
        self.rate_limiters.get(&tenant_id).map(|r| Arc::clone(&*r))
    }

    /// Route an entity to a shard via jump consistent hash.
    ///
    /// If the tenant has been pinned to a shard (after migration finalization),
    /// the pinned shard is returned directly, bypassing the hash.
    ///
    /// # Errors
    ///
    /// Returns an error if no eligible shards exist for the tenant's region constraints.
    pub fn route(
        &self,
        tenant_id: TenantId,
        entity_id: EntityId,
    ) -> crate::Result<ShardAssignment> {
        let topo = self
            .topology
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);

        // After migration finalization a tenant is pinned to its new shard.
        if let Some(pinned) = self.shard_pin_map.get(&tenant_id) {
            let shard_id = *pinned;
            let assignment = topo
                .shards
                .iter()
                .find(|s| s.shard_id == shard_id)
                .copied()
                .unwrap_or(ShardAssignment {
                    shard_id,
                    region_id: RegionId::SINGLE,
                });
            return Ok(assignment);
        }

        let eligible = self.eligible_shards_for(tenant_id, &topo)?;
        drop(topo);
        Ok(jump_hash_select(entity_id.as_u64(), &eligible))
    }

    /// Return the number of registered tenants.
    #[must_use]
    pub fn tenant_count(&self) -> usize {
        self.tenants.len()
    }

    /// Enter dual-write mode for `tenant_id`: writes are routed to both
    /// `source_shard` and `target_shard` until `finalize_migration` is called.
    ///
    /// Reads continue to route to the source shard (via `route()`) until
    /// finalization pins the tenant to the target.
    pub fn set_dual_write(&self, tenant_id: TenantId, source: ShardId, target: ShardId) {
        self.dual_write_map.insert(tenant_id, (source, target));
    }

    /// Finalize migration for `tenant_id`: clear dual-write mode and pin all
    /// future routing to `target_shard`.
    ///
    /// After this call, `route()` and `write_assignments()` both return
    /// `target_shard` exclusively. The old source shard can be GC'd.
    pub fn finalize_migration(&self, tenant_id: TenantId, target: ShardId) {
        self.dual_write_map.remove(&tenant_id);
        self.shard_pin_map.insert(tenant_id, target);
    }

    /// Return all shard assignments that a write for `(tenant_id, entity_id)`
    /// must be sent to.
    ///
    /// - Normal mode: one shard (same as `route()`).
    /// - Dual-write mode: two shards — source and target — so both copies are
    ///   updated during migration with no signal loss.
    ///
    /// On a single-node deployment the caller writes to local storage once;
    /// the returned assignments carry the routing metadata for multi-node use.
    ///
    /// # Errors
    ///
    /// Propagates errors from `route()` (e.g., no eligible shards).
    pub fn write_assignments(
        &self,
        tenant_id: TenantId,
        entity_id: EntityId,
    ) -> crate::Result<Vec<ShardAssignment>> {
        if let Some(pair) = self.dual_write_map.get(&tenant_id) {
            let (source_shard, target_shard) = *pair;
            let topo = self
                .topology
                .read()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let find = |shard_id: ShardId| -> ShardAssignment {
                topo.shards
                    .iter()
                    .find(|s| s.shard_id == shard_id)
                    .copied()
                    .unwrap_or(ShardAssignment {
                        shard_id,
                        region_id: RegionId::SINGLE,
                    })
            };
            return Ok(vec![find(source_shard), find(target_shard)]);
        }
        Ok(vec![self.route(tenant_id, entity_id)?])
    }

    /// Return `true` if `tenant_id` is currently in dual-write migration mode.
    #[must_use]
    pub fn is_dual_write(&self, tenant_id: TenantId) -> bool {
        self.dual_write_map.contains_key(&tenant_id)
    }

    /// Return the pinned target shard for `tenant_id` after migration, if any.
    #[must_use]
    pub fn pinned_shard(&self, tenant_id: TenantId) -> Option<ShardId> {
        self.shard_pin_map.get(&tenant_id).map(|r| *r)
    }

    fn eligible_shards_for(
        &self,
        tenant_id: TenantId,
        topo: &ClusterTopology,
    ) -> crate::Result<Vec<ShardAssignment>> {
        let required_regions: Vec<RegionId> = self
            .tenants
            .get(&tenant_id)
            .map(|c| c.required_regions.clone())
            .unwrap_or_default();

        let eligible: Vec<ShardAssignment> = if required_regions.is_empty() {
            topo.shards.clone()
        } else {
            topo.shards
                .iter()
                .filter(|s| required_regions.contains(&s.region_id))
                .copied()
                .collect()
        };

        if eligible.is_empty() {
            return Err(crate::TidalError::internal(
                "route",
                "no eligible shards for tenant",
            ));
        }
        Ok(eligible)
    }
}

/// Jump consistent hash (Lamping & Veach, 2014).
///
/// The casts between i64/u64/f64 are inherent to the algorithm's use of
/// signed intermediates and floating-point arithmetic. They are safe because:
/// - `n` (shard count) is always small (< 2^16 in practice)
/// - `b` and `j` are bounded by `n`
/// - The final `b as usize` is non-negative since the loop ensures `b >= 0`
#[allow(
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss
)]
fn jump_hash_select(key: u64, shards: &[ShardAssignment]) -> ShardAssignment {
    debug_assert!(
        !shards.is_empty(),
        "jump_hash_select requires non-empty shards"
    );
    let n = shards.len() as u64;
    let mut k = key;
    let mut b: i64 = -1;
    let mut j: i64 = 0;
    while j < n as i64 {
        b = j;
        k = k.wrapping_mul(2_862_933_555_777_941_757).wrapping_add(1);
        j = ((b + 1) as f64 * (f64::from(1u32 << 31) / ((k >> 33) + 1) as f64)) as i64;
    }
    shards[b as usize]
}

/// Build the WAL directory path for a tenant.
#[must_use]
pub fn tenant_wal_dir(data_dir: &std::path::Path, tenant_id: TenantId) -> std::path::PathBuf {
    if tenant_id == TenantId::DEFAULT {
        data_dir.join("wal")
    } else {
        data_dir
            .join("tenants")
            .join(tenant_id.0.to_string())
            .join("wal")
    }
}

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

    #[test]
    fn tenant_id_default_is_zero() {
        assert_eq!(TenantId::DEFAULT.0, 0);
    }

    #[test]
    fn tenant_id_display() {
        assert_eq!(TenantId(42).to_string(), "tenant:42");
    }

    #[test]
    fn rate_limiter_allows_burst() {
        let limiter = TenantRateLimiter::new(100);
        for _ in 0..200 {
            assert!(limiter.try_acquire().is_ok());
        }
        assert!(limiter.try_acquire().is_err());
    }

    #[test]
    fn cluster_topology_single() {
        let t = ClusterTopology::single();
        assert_eq!(t.shards.len(), 1);
        assert_eq!(t.shards[0].shard_id, ShardId::SINGLE);
    }

    #[test]
    fn tenant_router_register_and_count() {
        let topo = Arc::new(RwLock::new(ClusterTopology::single()));
        let router = TenantRouter::new(topo);
        router.register_tenant(TenantConfig::default_tenant());
        assert_eq!(router.tenant_count(), 1);
    }

    #[test]
    fn tenant_router_rate_limiter_cached() {
        let topo = Arc::new(RwLock::new(ClusterTopology::single()));
        let router = TenantRouter::new(topo);
        let mut cfg = TenantConfig::default_tenant();
        cfg.max_signals_per_sec = Some(50);
        router.register_tenant(cfg);
        let l1 = router.rate_limiter_for(TenantId::DEFAULT).unwrap();
        let l2 = router.rate_limiter_for(TenantId::DEFAULT).unwrap();
        assert!(Arc::ptr_eq(&l1, &l2));
    }

    #[test]
    fn tenant_router_route_single_shard() {
        let topo = Arc::new(RwLock::new(ClusterTopology::single()));
        let router = TenantRouter::new(topo);
        router.register_tenant(TenantConfig::default_tenant());
        let assignment = router.route(TenantId::DEFAULT, EntityId::new(42)).unwrap();
        assert_eq!(assignment.shard_id, ShardId::SINGLE);
    }

    #[test]
    fn jump_hash_distributes_uniformly() {
        let shards: Vec<ShardAssignment> = (0..4u16)
            .map(|i| ShardAssignment {
                shard_id: ShardId(i),
                region_id: RegionId::SINGLE,
            })
            .collect();
        let mut counts = [0usize; 4];
        for i in 0..4000u64 {
            let a = jump_hash_select(i, &shards);
            counts[a.shard_id.0 as usize] += 1;
        }
        for c in counts {
            assert!(
                c > 800 && c < 1200,
                "bucket count {c} out of expected range"
            );
        }
    }

    #[test]
    fn tenant_wal_dir_default() {
        let dir = std::path::Path::new("/data");
        assert_eq!(
            tenant_wal_dir(dir, TenantId::DEFAULT),
            std::path::PathBuf::from("/data/wal")
        );
    }

    #[test]
    fn tenant_wal_dir_non_default() {
        let dir = std::path::Path::new("/data");
        assert_eq!(
            tenant_wal_dir(dir, TenantId(7)),
            std::path::PathBuf::from("/data/tenants/7/wal")
        );
    }
}
