use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};

use dashmap::DashMap;
use tokio::sync::Mutex;
use uuid::Uuid;

use crate::cache::redis::AppRedis;
use crate::error::{AppError, AppResult};
use ::redis::Cmd;

use super::redis_keys::*;

struct Segment {
    end: i64,
    next: AtomicI64,
}

pub struct SeqAllocator {
    redis: AppRedis,
    segments: DashMap<Uuid, Arc<Segment>>,
    locks: DashMap<Uuid, Arc<Mutex<()>>>,
    segment_size: u64,
}

const MAX_RETRIES: u32 = 3;

impl SeqAllocator {
    pub fn new(redis: AppRedis) -> Self {
        Self {
            redis,
            segments: DashMap::new(),
            locks: DashMap::new(),
            segment_size: WS_SEQ_SEGMENT_SIZE,
        }
    }

    pub async fn next(&self, channel_id: Uuid) -> AppResult<i64> {
        for _ in 0..MAX_RETRIES {
            if let Some(seq) = self.try_allocate(&channel_id) {
                return Ok(seq);
            }
            let lock = self
                .locks
                .entry(channel_id)
                .or_insert_with(|| Arc::new(Mutex::new(())))
                .clone();
            let _guard = lock.lock().await;
            if let Some(seq) = self.try_allocate(&channel_id) {
                return Ok(seq);
            }
            self.refresh(channel_id).await?;
        }
        Err(AppError::InternalServerError(
            "seq allocation exhausted retries".into(),
        ))
    }

    pub async fn bootstrap(&self, channel_id: Uuid, db_max: i64) -> AppResult<i64> {
        let key = format!("{WS_SEQ_PREFIX}{channel_id}");
        let mut conn = self.redis.get_connection()?;
        let current: i64 = Cmd::new()
            .arg("SET")
            .arg(&key)
            .arg(db_max)
            .arg("NX")
            .arg("EX")
            .arg(86400)
            .query::<Option<String>>(&mut *conn.inner_mut())
            .map_err(AppError::Redis)?
            .and_then(|v| v.parse().ok())
            .unwrap_or_else(|| {
                let existing: i64 = Cmd::new()
                    .arg("GET")
                    .arg(&key)
                    .query(&mut *conn.inner_mut())
                    .map_err(AppError::Redis)
                    .unwrap_or(db_max);
                if existing < db_max { db_max } else { existing }
            });
        self.segments.remove(&channel_id);
        Ok(current)
    }

    fn try_allocate(&self, channel_id: &Uuid) -> Option<i64> {
        let state = self.segments.get(channel_id)?;
        let next = state.next.fetch_add(1, Ordering::Relaxed);
        if next < state.end { Some(next) } else { None }
    }

    async fn refresh(&self, channel_id: Uuid) -> AppResult<()> {
        let key = format!("{WS_SEQ_PREFIX}{channel_id}");
        let mut conn = self.redis.get_connection()?;
        let counter: i64 = Cmd::new()
            .arg("INCRBY")
            .arg(&key)
            .arg(self.segment_size as i64)
            .query(&mut *conn.inner_mut())
            .map_err(AppError::Redis)?;

        let start = counter - self.segment_size as i64 + 1;
        let end = counter + 1;
        self.segments.insert(
            channel_id,
            Arc::new(Segment {
                end,
                next: AtomicI64::new(start),
            }),
        );
        let _ = Cmd::new()
            .arg("EXPIRE")
            .arg(&key)
            .arg(86400_u64)
            .query::<()>(&mut *conn.inner_mut());
        Ok(())
    }
}