use anyhow::Result;
use mongodb::bson::{doc, DateTime};
use mongodb::Collection;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Clone)]
pub struct AccountLockout {
    user_collection: Arc<RwLock<Collection<mongodb::bson::Document>>>,
    max_attempts: u32,
    base_duration_minutes: u32,
    max_duration_minutes: u32,
}

impl AccountLockout {
    pub fn new(
        user_collection: Collection<mongodb::bson::Document>,
        max_attempts: u32,
        base_duration_minutes: u32,
        max_duration_minutes: u32,
    ) -> Self {
        Self {
            user_collection: Arc::new(RwLock::new(user_collection)),
            max_attempts,
            base_duration_minutes,
            max_duration_minutes,
        }
    }

    pub async fn check_lockout(&self, email: &str) -> Result<bool> {
        let collection = self.user_collection.read().await;
        let user = collection.find_one(doc! { "email": email }, None).await?;

        if let Some(user_doc) = user {
            if let Some(locked_until_val) = user_doc.get("locked_until") {
                if let Some(dt) = locked_until_val.as_datetime() {
                    let now = DateTime::now();
                    if dt.timestamp_millis() > now.timestamp_millis() {
                        return Ok(true); // Account is locked
                    }
                }
            }
        }

        Ok(false) // Account is not locked
    }

    pub async fn record_failed_attempt(&self, email: &str) -> Result<bool> {
        let collection = self.user_collection.write().await;

        // Get current failed attempts
        let user = collection.find_one(doc! { "email": email }, None).await?;

        let current_attempts = if let Some(user_doc) = user {
            user_doc
                .get("failed_login_attempts")
                .and_then(|v| v.as_i64())
                .unwrap_or(0) as u32
        } else {
            0
        };

        let new_attempts = current_attempts + 1;
        let should_lock = new_attempts >= self.max_attempts;

        // Calculate lockout duration
        let lock_duration = if should_lock {
            let multiplier = new_attempts.saturating_sub(self.max_attempts) + 1;
            let duration = self.base_duration_minutes * multiplier;
            std::cmp::min(duration, self.max_duration_minutes)
        } else {
            0
        };

        let locked_until = if lock_duration > 0 {
            let now = DateTime::now();
            let duration_millis = lock_duration as u64 * 60 * 1000;
            DateTime::from_millis(now.timestamp_millis() + duration_millis as i64)
        } else {
            DateTime::now()
        };

        // Update user
        collection
            .update_one(
                doc! { "email": email },
                doc! {
                    "$set": {
                        "failed_login_attempts": new_attempts as i32,
                        "last_failed_login": DateTime::now(),
                        "locked_until": locked_until,
                    }
                },
                None,
            )
            .await?;

        Ok(should_lock)
    }

    pub async fn reset_attempts(&self, email: &str) -> Result<()> {
        let collection = self.user_collection.write().await;

        collection
            .update_one(
                doc! { "email": email },
                doc! {
                    "$set": {
                        "failed_login_attempts": 0,
                        "locked_until": null,
                    }
                },
                None,
            )
            .await?;

        Ok(())
    }
}
