src/application/services/auth.rs

use crate::domain::error::{AppError, AppResult};
use crate::domain::models::Merchant;
use crate::infrastructure::repositories::MerchantRepository;
use argon2::{
    password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString},
    Argon2,
};
use async_trait::async_trait;
use dashmap::DashMap;
use once_cell::sync::Lazy;
use std::sync::Arc;
use std::time::{Duration, Instant};
use uuid::Uuid;

// ─── Session Cache ────────────────────────────────────────────────────────────

struct CachedSession {
    version: i64,
    expires_at: Instant,
}

static SESSION_CACHE: Lazy<DashMap<String, CachedSession>> = Lazy::new(DashMap::new);

// ─── Login Attempt Tracker — Brute-Force Protection ──────────────────────────
//
// Keyed by lowercase email address. Tracks consecutive failed attempts within
// a rolling 15-minute window. After MAX_FAILURES attempts the account is locked
// until the window expires. A successful login resets the counter.
//
// This is intentionally in-memory (not DB) so it:
//   a) adds zero latency to the happy path
//   b) cannot be bypassed by hitting a different DB replica
//   c) auto-recovers on restart (acceptable for in-memory rate limiting)
//
// For persistent cross-replica locking, wire this to Redis or a Postgres
// advisory lock keyed on the email hash.

const MAX_FAILURES: u32 = 5;
const LOCKOUT_WINDOW: Duration = Duration::from_secs(15 * 60); // 15 minutes

struct LoginAttempt {
    failures: u32,
    first_failure: Instant,
    locked_until: Option<Instant>,
}

static LOGIN_ATTEMPTS: Lazy<DashMap<String, LoginAttempt>> = Lazy::new(DashMap::new);

/// Check if an email is currently locked out. Returns `Err(AppError::RateLimited)`
/// if the account should be refused, `Ok(())` if the attempt may proceed.
fn check_lockout(email: &str) -> AppResult<()> {
    let key = email.to_lowercase();
    if let Some(attempt) = LOGIN_ATTEMPTS.get(&key) {
        // If a hard lockout timestamp is set and still in the future, refuse.
        if let Some(locked_until) = attempt.locked_until {
            if Instant::now() < locked_until {
                tracing::warn!(
                    email = %email,
                    "Login blocked: account locked out after {} failed attempts",
                    attempt.failures
                );
                return Err(AppError::RateLimited);
            }
        }
        // Rolling window: if the first failure was > 15 minutes ago, the window
        // has expired and the counter will be reset on the next record_failure call.
    }
    Ok(())
}

/// Record a failed login. Increments the counter; locks the account after
/// MAX_FAILURES attempts within LOCKOUT_WINDOW.
fn record_failure(email: &str) {
    let key = email.to_lowercase();
    let now = Instant::now();

    let mut entry = LOGIN_ATTEMPTS.entry(key).or_insert_with(|| LoginAttempt {
        failures: 0,
        first_failure: now,
        locked_until: None,
    });

    // Reset window if the previous failure was outside the rolling window
    if now.duration_since(entry.first_failure) > LOCKOUT_WINDOW {
        entry.failures = 0;
        entry.first_failure = now;
        entry.locked_until = None;
    }

    entry.failures += 1;

    if entry.failures >= MAX_FAILURES {
        entry.locked_until = Some(now + LOCKOUT_WINDOW);
        tracing::warn!(
            email = %email,
            failures = entry.failures,
            "Login lockout triggered: too many failed attempts"
        );
    }
}

/// Reset the login attempt counter on successful authentication.
fn record_success(email: &str) {
    LOGIN_ATTEMPTS.remove(&email.to_lowercase());
}

// ─── Auth Service Trait ───────────────────────────────────────────────────────

#[async_trait]
pub trait AuthService: Send + Sync {
    async fn register(
        &self,
        email: &str,
        password: &str,
        brand_name: &str,
        slug: Option<&str>,
        social_url: Option<&str>,
        upi_id: Option<&str>,
    ) -> AppResult<(Merchant, String, String)>;
    async fn login(&self, email: &str, password: &str) -> AppResult<(Merchant, String)>;
    async fn reset_password(
        &self,
        email: &str,
        recovery_key: &str,
        new_password: &str,
    ) -> AppResult<()>;
    async fn refresh_token(&self, merchant_id: &str) -> AppResult<String>;
    async fn verify_session(&self, merchant_id: &str, version: i64) -> AppResult<()>;
}

pub struct RtixAuthService {
    merchant_repo: Arc<dyn MerchantRepository>,
    jwt_secret: Vec<u8>,
}

impl RtixAuthService {
    pub fn new(merchant_repo: Arc<dyn MerchantRepository>, jwt_secret: Vec<u8>) -> Self {
        Self {
            merchant_repo,
            jwt_secret,
        }
    }

    fn hash_password(&self, password: &str) -> AppResult<String> {
        let mut rng = rand::thread_rng();
        let salt = SaltString::generate(&mut rng);
        let argon2 = Argon2::default();
        argon2
            .hash_password(password.as_bytes(), &salt)
            .map(|h| h.to_string())
            .map_err(|e| AppError::Internal(format!("Hashing failed: {}", e)))
    }

    fn verify_password(&self, password: &str, hash: &str) -> AppResult<()> {
        let parsed_hash = PasswordHash::new(hash)
            .map_err(|_| AppError::Internal("Invalid hash stored".to_string()))?;
        Argon2::default()
            .verify_password(password.as_bytes(), &parsed_hash)
            .map_err(|_| AppError::Auth("Invalid credentials".to_string()))
    }

    fn issue_token(&self, merchant: &crate::domain::models::Merchant) -> AppResult<String> {
        let claims = crate::interfaces::http::routes::auth::Claims {
            sub: merchant.merchant_id.clone(),
            email: merchant.email.clone(),
            brand_name: merchant.brand_name.clone(),
            slug: merchant.slug.clone(),
            role: Some(merchant.role.clone()),
            version: merchant.session_version,
            exp: crate::core::session::access_token_expiry(),
        };

        jsonwebtoken::encode(
            &jsonwebtoken::Header::default(),
            &claims,
            &jsonwebtoken::EncodingKey::from_secret(&self.jwt_secret),
        )
        .map_err(|e| AppError::Internal(format!("Token issuance failed: {}", e)))
    }
}

#[async_trait]
impl AuthService for RtixAuthService {
    async fn register(
        &self,
        email: &str,
        password: &str,
        brand_name: &str,
        slug: Option<&str>,
        social_url: Option<&str>,
        upi_id: Option<&str>,
    ) -> AppResult<(Merchant, String, String)> {
        let existing: Option<Merchant> = self.merchant_repo.find_by_email(email).await?;
        if existing.is_some() {
            return Err(AppError::BadRequest("Email already exists".to_string()));
        }

        let merchant_id = Uuid::new_v4().to_string();
        let password_hash = self.hash_password(password)?;
        let recovery_key = format!(
            "VTX-{}",
            &Uuid::new_v4().to_string().replace('-', "")[..16].to_uppercase()
        );
        let recovery_key_hash = self.hash_password(&recovery_key)?;

        let mut final_slug = slug.unwrap_or(brand_name).to_lowercase().replace(' ', "-");

        // Pillar IV: Identity Resilience - Automatic collision resolution
        let mut retry_count = 0;
        while self
            .merchant_repo
            .find_by_slug(&final_slug)
            .await?
            .is_some()
        {
            retry_count += 1;
            if retry_count > 5 {
                return Err(AppError::Internal(
                    "Could not generate a unique slug".to_string(),
                ));
            }
            let suffix = &Uuid::new_v4().to_string()[..4].to_lowercase();
            final_slug = format!(
                "{}-{}",
                slug.unwrap_or(brand_name).to_lowercase().replace(' ', "-"),
                suffix
            );
        }

        let merchant = Merchant {
            merchant_id: merchant_id.clone(),
            email: email.to_string(),
            password_hash,
            brand_name: brand_name.to_string(),
            slug: final_slug,
            social_url: social_url.map(ToString::to_string),
            upi_id: upi_id.map(ToString::to_string),
            business_address: None,
            recovery_key: Some(recovery_key_hash),
            session_version: 1,
            delivery_rate_per_km: 10.0,
            delivery_base_fee: 20.0,
            logistics_config: serde_json::json!({
                "complexity_bias": 1.0,
                "weight_coefficient": 0.02,
                "distance_coefficient": 1.0
            }),
            base_pincode: "560001".to_string(),
            auto_settle_threshold: 50.0,
            trust_score: 100.0,
            verification_level: "UNVERIFIED".to_string(),
            max_order_value_inr: 10000.0,
            created_at: None,
            state_code: Some(29),
            gstin: None,
            announcement_banner: None,
            plan: "FREE".to_string(),
            role: "MERCHANT".to_string(),
            is_frozen: false,
            billing_cycle_start: None,
        };

        self.merchant_repo.create(&merchant).await?;

        let token = self.issue_token(&merchant)?;

        Ok((merchant, token, recovery_key))
    }

    /// Authenticate a merchant with brute-force protection.
    ///
    /// # Brute-Force Guard
    /// - After **5 consecutive failed attempts** within a 15-minute window the
    ///   email is locked and every further attempt returns HTTP 429 immediately
    ///   (before any DB query or Argon2 verification, preventing timing oracles).
    /// - A successful login resets the attempt counter.
    /// - The lockout is in-memory and resets on server restart (acceptable for
    ///   single-replica deployments; use Redis for multi-replica hardening).
    async fn login(&self, email: &str, password: &str) -> AppResult<(Merchant, String)> {
        // ── Brute-force check BEFORE any DB work ──────────────────────────────
        check_lockout(email)?;

        let merchant: Option<Merchant> = self.merchant_repo.find_by_email(email).await?;
        let merchant = merchant.ok_or_else(|| {
            // Record failure even on unknown email to prevent user enumeration
            // via timing differences (non-existent vs. wrong-password paths).
            record_failure(email);
            AppError::Auth("Invalid email or password".to_string())
        })?;

        // ── Argon2 password verification ──────────────────────────────────────
        if let Err(e) = self.verify_password(password, &merchant.password_hash) {
            record_failure(email);
            tracing::warn!(
                email = %email,
                "Failed login attempt recorded"
            );
            return Err(e);
        }

        // ── Success: clear attempt counter ────────────────────────────────────
        record_success(email);

        let token = self.issue_token(&merchant)?;
        Ok((merchant, token))
    }

    async fn reset_password(
        &self,
        email: &str,
        recovery_key: &str,
        new_password: &str,
    ) -> AppResult<()> {
        let merchant: Option<Merchant> = self.merchant_repo.find_by_email(email).await?;
        let merchant = merchant.ok_or(AppError::NotFound("Merchant not found".to_string()))?;

        if let Some(hash) = &merchant.recovery_key {
            self.verify_password(recovery_key, hash)?;
        } else {
            return Err(AppError::Auth("No recovery key set".to_string()));
        }

        let new_hash = self.hash_password(new_password)?;
        self.merchant_repo.update_password(email, &new_hash).await?;

        Ok(())
    }

    async fn refresh_token(&self, merchant_id: &str) -> AppResult<String> {
        let merchant = self.merchant_repo.find_by_id(merchant_id).await?;
        let merchant = merchant.ok_or(AppError::NotFound("Merchant not found".to_string()))?;

        let token = self.issue_token(&merchant)?;
        Ok(token)
    }

    async fn verify_session(&self, merchant_id: &str, version: i64) -> AppResult<()> {
        let now = Instant::now();
        if let Some(cached) = SESSION_CACHE.get(merchant_id) {
            if cached.version == version && cached.expires_at > now {
                return Ok(());
            }
        }

        let merchant = self.merchant_repo.find_by_id(merchant_id).await?;
        let merchant = merchant.ok_or(AppError::Auth("Session invalid".to_string()))?;

        if merchant.session_version != version {
            return Err(AppError::Auth("Session version mismatch".to_string()));
        }

        if !SESSION_CACHE.contains_key(merchant_id) && SESSION_CACHE.len() > 10_000 {
            SESSION_CACHE.retain(|_, v| v.expires_at > now);
        }

        SESSION_CACHE.insert(
            merchant_id.to_string(),
            CachedSession {
                version,
                expires_at: now + Duration::from_secs(60),
            },
        );

        Ok(())
    }
}