use async_trait::async_trait;
use axum_session::{DatabaseError, DatabasePool, SessionError};
use chrono::Utc;
use config::session::SessionConfig;
use export::{AppRedis, RedisConnection};

use crate::config::build_session_config;

pub type Session = axum_session::Session<SessionRedisPool>;
pub type ReadOnlySession = axum_session::ReadOnlySession<SessionRedisPool>;
pub type SessionStore = axum_session::SessionStore<SessionRedisPool>;
pub type SessionLayer = axum_session::SessionLayer<SessionRedisPool>;

/// Keys deleted by a single `UNLINK` call in standalone mode. Kept at one key
/// per call in cluster mode, where multi-key commands would cross slots.
const UNLINK_BATCH: usize = 100;

/// Hint for `SCAN COUNT`, trading a few bigger replies for far fewer round trips.
const SCAN_BATCH: usize = 1000;

#[derive(Clone)]
pub struct SessionRedisPool {
    redis: AppRedis,
}

impl SessionRedisPool {
    pub fn new(redis: AppRedis) -> Self {
        Self { redis }
    }

    pub fn redis(&self) -> &AppRedis {
        &self.redis
    }
}

impl From<AppRedis> for SessionRedisPool {
    fn from(redis: AppRedis) -> Self {
        Self::new(redis)
    }
}

impl std::fmt::Debug for SessionRedisPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SessionRedisPool").finish()
    }
}

pub async fn connect(
    redis: AppRedis,
    config: &SessionConfig,
) -> Result<SessionStore, SessionError> {
    let config = build_session_config(config).map_err(|err| {
        SessionError::GenericNotSupportedError(format!("invalid session configuration: {err}"))
    })?;

    SessionStore::new(Some(SessionRedisPool::new(redis)), config).await
}

macro_rules! dispatch {
    ($conn:expr, |$c:ident| $body:expr) => {
        match &mut $conn {
            RedisConnection::Standalone($c) => $body,
            RedisConnection::Cluster($c) => $body,
        }
    };
}

fn session_key(id: &str, table_name: &str) -> String {
    if table_name.is_empty() {
        id.to_string()
    } else {
        format!("{table_name}:{id}")
    }
}

/// SCANs `pattern` across the whole store: a cursor loop on the standalone
/// connection, or the same loop on every cluster master with results merged.
async fn scan_keys(redis: &AppRedis, pattern: &str) -> Result<Vec<String>, DatabaseError> {
    let mut keys = match redis.connection_clone() {
        RedisConnection::Standalone(mut conn) => scan_conn(&mut conn, pattern).await?,
        RedisConnection::Cluster(_) => {
            let mut keys = Vec::new();

            for mut conn in redis
                .cluster_master_connections()
                .await
                .map_err(|err| DatabaseError::GenericSelectError(err.to_string()))?
            {
                keys.extend(scan_conn(&mut conn, pattern).await?);
            }

            keys
        }
    };

    keys.sort();
    keys.dedup();
    Ok(keys)
}

async fn scan_conn<C>(conn: &mut C, pattern: &str) -> Result<Vec<String>, DatabaseError>
where
    C: redis::aio::ConnectionLike,
{
    let mut keys = Vec::new();
    let mut cursor: u64 = 0;

    loop {
        let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
            .arg(cursor)
            .arg("MATCH")
            .arg(pattern)
            .arg("COUNT")
            .arg(SCAN_BATCH)
            .query_async(conn)
            .await
            .map_err(|err| DatabaseError::GenericSelectError(err.to_string()))?;

        keys.extend(batch);
        cursor = next;

        if cursor == 0 {
            break;
        }
    }

    Ok(keys)
}

#[async_trait]
impl DatabasePool for SessionRedisPool {
    async fn initiate(&self, _table_name: &str) -> Result<(), DatabaseError> {
        Ok(())
    }

    async fn delete_by_expiry(&self, _table_name: &str) -> Result<Vec<String>, DatabaseError> {
        Ok(Vec::new())
    }

    async fn count(&self, table_name: &str) -> Result<i64, DatabaseError> {
        ensure_table(table_name)?;

        let count = scan_keys(&self.redis, &format!("{table_name}:*"))
            .await?
            .len() as i64;

        Ok(count)
    }

    async fn store(
        &self,
        id: &str,
        session: &str,
        expires: i64,
        table_name: &str,
    ) -> Result<(), DatabaseError> {
        // axum_session passes `expires` as a unix timestamp. Store a relative
        // TTL instead of `EXAT` so Redis < 6.2 keeps working.
        let ttl_ms = (expires - Utc::now().timestamp()) * 1000;
        if ttl_ms <= 0 {
            return Ok(());
        }

        let key = session_key(id, table_name);
        let mut conn = self.redis.connection_clone();

        dispatch!(conn, |c| {
            redis::cmd("SET")
                .arg(&key)
                .arg(session)
                .arg("PX")
                .arg(ttl_ms)
                .exec_async(c)
                .await
        })
        .map_err(|err| DatabaseError::GenericInsertError(err.to_string()))
    }

    async fn load(&self, id: &str, table_name: &str) -> Result<Option<String>, DatabaseError> {
        let key = session_key(id, table_name);
        let mut conn = self.redis.connection_clone();

        dispatch!(conn, |c| redis::cmd("GET").arg(&key).query_async(c).await)
            .map_err(|err| DatabaseError::GenericSelectError(err.to_string()))
    }

    async fn delete_one_by_id(&self, id: &str, table_name: &str) -> Result<(), DatabaseError> {
        let key = session_key(id, table_name);
        let mut conn = self.redis.connection_clone();

        dispatch!(conn, |c| redis::cmd("UNLINK").arg(&key).exec_async(c).await)
            .map_err(|err| DatabaseError::GenericDeleteError(err.to_string()))
    }

    async fn exists(&self, id: &str, table_name: &str) -> Result<bool, DatabaseError> {
        let key = session_key(id, table_name);
        let mut conn = self.redis.connection_clone();

        dispatch!(conn, |c| redis::cmd("EXISTS")
            .arg(&key)
            .query_async(c)
            .await)
        .map_err(|err| DatabaseError::GenericSelectError(err.to_string()))
    }

    async fn delete_all(&self, table_name: &str) -> Result<(), DatabaseError> {
        ensure_table(table_name)?;

        let keys = scan_keys(&self.redis, &format!("{table_name}:*")).await?;
        if keys.is_empty() {
            return Ok(());
        }

        let mut conn = self.redis.connection_clone();
        match &mut conn {
            // A single multi-key UNLINK must not cross slots, so batching is
            // only safe on a standalone connection.
            RedisConnection::Standalone(_) => {
                for chunk in keys.chunks(UNLINK_BATCH) {
                    dispatch!(conn, |c| {
                        let mut cmd = redis::cmd("UNLINK");
                        for key in chunk {
                            cmd.arg(key);
                        }
                        cmd.exec_async(c).await
                    })
                    .map_err(|err| DatabaseError::GenericDeleteError(err.to_string()))?;
                }
            }
            RedisConnection::Cluster(_) => {
                for key in &keys {
                    dispatch!(conn, |c| redis::cmd("UNLINK").arg(key).exec_async(c).await)
                        .map_err(|err| DatabaseError::GenericDeleteError(err.to_string()))?;
                }
            }
        }

        Ok(())
    }

    async fn get_ids(&self, table_name: &str) -> Result<Vec<String>, DatabaseError> {
        ensure_table(table_name)?;

        let prefix = format!("{table_name}:");

        let ids = scan_keys(&self.redis, &format!("{table_name}:*"))
            .await?
            .into_iter()
            .filter_map(|key| key.strip_prefix(&prefix).map(str::to_string))
            .collect();

        Ok(ids)
    }

    fn auto_handles_expiry(&self) -> bool {
        true
    }
}

/// Session keys are managed under a `{table_name}:` prefix; an empty prefix
/// would make scans, counts and deletes target unrelated keys (or the whole
/// database). `connect` already rejects it, this is the defensive backstop.
fn ensure_table(table_name: &str) -> Result<(), DatabaseError> {
    if table_name.is_empty() {
        return Err(DatabaseError::GenericNotSupportedError(
            "session table_name must not be empty".into(),
        ));
    }

    Ok(())
}