neiam /dms

use uuid::Uuid;
use chrono::{Duration, Utc};
use sea_orm::*;
use anyhow::{Context, Result};
use crate::entities::{users, user_sessions, registration_tokens};

/// Generate a new session token
pub fn generate_session_token() -> String {
    Uuid::new_v4().to_string()
}

/// Generate a new registration token
pub fn generate_registration_token() -> String {
    Uuid::new_v4().to_string()
}

/// Create a new user session
pub async fn create_session(
    db: &DatabaseConnection,
    user_id: Uuid,
    session_duration_hours: i64,
) -> Result<user_sessions::Model> {
    let session_token = generate_session_token();
    let expires_at = Utc::now() + Duration::hours(session_duration_hours);

    let session = user_sessions::ActiveModel {
        id: Set(Uuid::now_v7()),
        user_id: Set(user_id),
        session_token: Set(session_token),
        expires_at: Set(expires_at.into()),
        created_at: Set(Utc::now().into()),
        last_used_at: Set(Utc::now().into()),
    };

    let session = session.insert(db).await.context("Failed to create session")?;
    Ok(session)
}

/// Validate and refresh a session
pub async fn validate_session(
    db: &DatabaseConnection,
    token: &str,
) -> Result<Option<(users::Model, user_sessions::Model)>> {
    let session = user_sessions::Entity::find()
        .filter(user_sessions::Column::SessionToken.eq(token))
        .filter(user_sessions::Column::ExpiresAt.gt(Utc::now()))
        .find_also_related(users::Entity)
        .one(db)
        .await
        .context("Failed to query session")?;

    if let Some((session_model, Some(user_model))) = session {
        // Update last_used_at
        let mut session_update: user_sessions::ActiveModel = session_model.clone().into();
        session_update.last_used_at = Set(Utc::now().into());
        session_update.update(db).await.context("Failed to update session")?;

        Ok(Some((user_model, session_model)))
    } else {
        Ok(None)
    }
}

/// Delete a session (logout)
pub async fn delete_session(db: &DatabaseConnection, token: &str) -> Result<bool> {
    let result = user_sessions::Entity::delete_many()
        .filter(user_sessions::Column::SessionToken.eq(token))
        .exec(db)
        .await
        .context("Failed to delete session")?;

    Ok(result.rows_affected > 0)
}

/// Clean up expired sessions
pub async fn cleanup_expired_sessions(db: &DatabaseConnection) -> Result<u64> {
    let result = user_sessions::Entity::delete_many()
        .filter(user_sessions::Column::ExpiresAt.lt(Utc::now()))
        .exec(db)
        .await
        .context("Failed to cleanup expired sessions")?;

    Ok(result.rows_affected)
}

/// Create a registration token
pub async fn create_registration_token(
    db: &DatabaseConnection,
    created_by_user_id: Option<Uuid>,
    email: Option<String>,
    expires_in_days: i64,
) -> Result<registration_tokens::Model> {
    let token = generate_registration_token();
    let expires_at = Utc::now() + Duration::days(expires_in_days);

    let registration_token = registration_tokens::ActiveModel {
        id: Set(Uuid::now_v7()),
        token: Set(token),
        email: Set(email),
        created_by_user_id: Set(created_by_user_id),
        expires_at: Set(expires_at.into()),
        is_used: Set(false),
        used_by_user_id: Set(None),
        used_at: Set(None),
        created_at: Set(Utc::now().into()),
    };

    let token_model = registration_token
        .insert(db)
        .await
        .context("Failed to create registration token")?;

    Ok(token_model)
}

/// Validate a registration token
pub async fn validate_registration_token(
    db: &DatabaseConnection,
    token: &str,
    email: &str,
) -> Result<Option<registration_tokens::Model>> {
    let reg_token = registration_tokens::Entity::find()
        .filter(registration_tokens::Column::Token.eq(token))
        .filter(registration_tokens::Column::IsUsed.eq(false))
        .filter(registration_tokens::Column::ExpiresAt.gt(Utc::now()))
        .one(db)
        .await
        .context("Failed to query registration token")?;

    match reg_token {
        Some(token_model) => {
            // If token has a specific email, check it matches
            if let Some(token_email) = &token_model.email
                && token_email != email
            {
                return Ok(None);
            }
            Ok(Some(token_model))
        }
        None => Ok(None),
    }
}

/// Mark a registration token as used
pub async fn mark_registration_token_used(
    db: &DatabaseConnection,
    token_id: Uuid,
    used_by_user_id: Uuid,
) -> Result<()> {
    let mut token: registration_tokens::ActiveModel = registration_tokens::Entity::find_by_id(token_id)
        .one(db)
        .await
        .context("Failed to find registration token")?
        .context("Registration token not found")?
        .into();

    token.is_used = Set(true);
    token.used_by_user_id = Set(Some(used_by_user_id));
    token.used_at = Set(Some(Utc::now().into()));

    token.update(db).await.context("Failed to mark token as used")?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_token_generation() {
        let session_token = generate_session_token();
        let reg_token = generate_registration_token();
        
        assert_ne!(session_token, reg_token);
        assert!(session_token.len() > 20); // UUID should be longer than 20 chars
        assert!(reg_token.len() > 20);
        
        // Ensure they're valid UUIDs
        assert!(Uuid::parse_str(&session_token).is_ok());
        assert!(Uuid::parse_str(&reg_token).is_ok());
    }
}