use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use crank_core::{RateLimitBucketState, RateLimitStateStore};
use thiserror::Error;

const STALE_KEY_TTL: Duration = Duration::from_secs(300);
const TOKEN_SCALE: u64 = 1_000_000;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RequestRateLimitConfig {
    pub requests_per_second: u32,
    pub burst: u32,
}

impl RequestRateLimitConfig {
    pub fn new(requests_per_second: u32, burst: u32) -> Result<Self, RequestRateLimitConfigError> {
        if requests_per_second == 0 {
            return Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(
                requests_per_second,
            ));
        }
        if burst == 0 {
            return Err(RequestRateLimitConfigError::InvalidBurst(burst));
        }

        Ok(Self {
            requests_per_second,
            burst,
        })
    }
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum RequestRateLimitConfigError {
    #[error("requests_per_second must be greater than zero, got {0}")]
    InvalidRequestsPerSecond(u32),
    #[error("burst must be greater than zero, got {0}")]
    InvalidBurst(u32),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RateLimitRejection {
    pub retry_after_ms: u64,
}

#[derive(Clone)]
pub struct RequestRateLimiter {
    config: RequestRateLimitConfig,
    backend: RequestRateLimiterBackend,
}

#[derive(Clone)]
enum RequestRateLimiterBackend {
    Local {
        states: Arc<Mutex<HashMap<String, LocalBucketState>>>,
    },
    Shared {
        store: Arc<dyn RateLimitStateStore>,
    },
}

#[derive(Clone, Copy, Debug)]
struct LocalBucketState {
    tokens: f64,
    last_refill: Instant,
    last_seen: Instant,
}

impl RequestRateLimiter {
    pub fn new(config: RequestRateLimitConfig) -> Self {
        Self {
            config,
            backend: RequestRateLimiterBackend::Local {
                states: Arc::new(Mutex::new(HashMap::new())),
            },
        }
    }

    pub fn new_shared(config: RequestRateLimitConfig, store: Arc<dyn RateLimitStateStore>) -> Self {
        Self {
            config,
            backend: RequestRateLimiterBackend::Shared { store },
        }
    }

    pub async fn check(&self, key: &str) -> Result<(), RateLimitRejection> {
        match &self.backend {
            RequestRateLimiterBackend::Local { .. } => self.check_local_at(key, Instant::now()),
            RequestRateLimiterBackend::Shared { store } => {
                self.check_shared_at(store.as_ref(), key, now_unix_ms())
                    .await
            }
        }
    }

    fn check_local_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> {
        let RequestRateLimiterBackend::Local { states } = &self.backend else {
            panic!("check_local_at called for non-local limiter");
        };

        let mut states = states
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        states.retain(|_, state| now.duration_since(state.last_seen) <= STALE_KEY_TTL);

        let burst = self.config.burst as f64;
        let requests_per_second = self.config.requests_per_second as f64;
        let state = states.entry(key.to_owned()).or_insert(LocalBucketState {
            tokens: burst,
            last_refill: now,
            last_seen: now,
        });

        let elapsed = now.duration_since(state.last_refill).as_secs_f64();
        state.tokens = (state.tokens + elapsed * requests_per_second).min(burst);
        state.last_refill = now;
        state.last_seen = now;

        if state.tokens >= 1.0 {
            state.tokens -= 1.0;
            return Ok(());
        }

        let missing_tokens = (1.0 - state.tokens).max(0.0);
        let retry_after_ms = ((missing_tokens / requests_per_second) * 1000.0)
            .ceil()
            .max(1.0) as u64;
        Err(RateLimitRejection { retry_after_ms })
    }

    async fn check_shared_at(
        &self,
        store: &dyn RateLimitStateStore,
        key: &str,
        now_unix_ms: i64,
    ) -> Result<(), RateLimitRejection> {
        let burst_tokens = u64::from(self.config.burst) * TOKEN_SCALE;
        let refill_per_second = u64::from(self.config.requests_per_second) * TOKEN_SCALE;
        let mut state =
            store
                .get_bucket(key)
                .await
                .unwrap_or(None)
                .unwrap_or(RateLimitBucketState {
                    tokens_micros: burst_tokens,
                    last_refill_unix_ms: now_unix_ms,
                });

        let elapsed_ms = (now_unix_ms - state.last_refill_unix_ms).max(0) as u64;
        let replenished =
            state.tokens_micros + (elapsed_ms.saturating_mul(refill_per_second) / 1000);
        state.tokens_micros = replenished.min(burst_tokens);
        state.last_refill_unix_ms = now_unix_ms;

        if state.tokens_micros >= TOKEN_SCALE {
            state.tokens_micros -= TOKEN_SCALE;
            let _ = store.put_bucket(key, state, STALE_KEY_TTL).await;
            return Ok(());
        }

        let missing_tokens = TOKEN_SCALE.saturating_sub(state.tokens_micros);
        let retry_after_ms = missing_tokens.div_ceil(refill_per_second).max(1);
        let _ = store.put_bucket(key, state, STALE_KEY_TTL).await;
        Err(RateLimitRejection { retry_after_ms })
    }
}

fn now_unix_ms() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_else(|_| Duration::from_secs(0))
        .as_millis()
        .try_into()
        .unwrap_or(i64::MAX)
}

#[cfg(test)]
mod tests {
    use std::{
        sync::Arc,
        time::{Duration, Instant},
    };

    use crate::InMemoryRateLimitStateStore;

    use super::{RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter};

    #[test]
    fn rejects_zero_requests_per_second() {
        assert_eq!(
            RequestRateLimitConfig::new(0, 1),
            Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(0))
        );
    }

    #[test]
    fn rejects_zero_burst() {
        assert_eq!(
            RequestRateLimitConfig::new(1, 0),
            Err(RequestRateLimitConfigError::InvalidBurst(0))
        );
    }

    #[test]
    fn allows_burst_then_rejects_until_refilled_for_local_limiter() {
        let limiter = RequestRateLimiter::new(RequestRateLimitConfig::new(2, 2).unwrap());
        let start = Instant::now();

        assert!(limiter.check_local_at("key", start).is_ok());
        assert!(limiter.check_local_at("key", start).is_ok());

        let rejection = limiter.check_local_at("key", start).unwrap_err();
        assert_eq!(rejection.retry_after_ms, 500);

        assert!(
            limiter
                .check_local_at("key", start + Duration::from_millis(500))
                .is_ok()
        );
    }

    #[tokio::test]
    async fn allows_burst_then_rejects_until_refilled_for_shared_limiter() {
        let limiter = RequestRateLimiter::new_shared(
            RequestRateLimitConfig::new(2, 2).unwrap(),
            Arc::new(InMemoryRateLimitStateStore::default()),
        );

        assert!(limiter.check_shared_at_store("key", 0).await.is_ok());
        assert!(limiter.check_shared_at_store("key", 0).await.is_ok());

        let rejection = limiter.check_shared_at_store("key", 0).await.unwrap_err();
        assert_eq!(rejection.retry_after_ms, 1);

        assert!(limiter.check_shared_at_store("key", 500).await.is_ok());
    }

    #[tokio::test]
    async fn shared_limiter_state_is_visible_across_instances() {
        let store = Arc::new(InMemoryRateLimitStateStore::default());
        let first = RequestRateLimiter::new_shared(
            RequestRateLimitConfig::new(1, 1).unwrap(),
            store.clone(),
        );
        let second =
            RequestRateLimiter::new_shared(RequestRateLimitConfig::new(1, 1).unwrap(), store);

        assert!(first.check_shared_at_store("shared", 0).await.is_ok());
        assert!(second.check_shared_at_store("shared", 0).await.is_err());
        assert!(second.check_shared_at_store("shared", 1000).await.is_ok());
    }

    impl RequestRateLimiter {
        async fn check_shared_at_store(
            &self,
            key: &str,
            now_unix_ms: i64,
        ) -> Result<(), super::RateLimitRejection> {
            let super::RequestRateLimiterBackend::Shared { store } = &self.backend else {
                panic!("check_shared_at_store called for non-shared limiter");
            };
            self.check_shared_at(store.as_ref(), key, now_unix_ms).await
        }
    }
}