quoter/src/app_state.rs

use crate::s3_utils::get_s3_filelist;
use actix_web::error::ErrorInternalServerError;
use aws_config::BehaviorVersion;
use aws_sdk_s3::{config::Credentials, Client as S3Client};
use log::warn;
use redis::{aio::MultiplexedConnection, AsyncCommands, Client as RedisClient};
use std::env;

#[derive(Clone)]
pub struct AppState {
    pub redis: MultiplexedConnection,
    pub storj_client: S3Client,
    pub aws_client: S3Client,
    pub bucket: String,
}

const PATH_MAPPING_KEY: &str = "filepath_mapping"; // Ключ для хранения маппинга путей
                                                   // Убираем TTL для квоты - она должна быть постоянной на пользователя

impl AppState {
    /// Инициализация нового состояния приложения.
    pub async fn new() -> Self {
        // Получаем конфигурацию для Redis
        let redis_url = env::var("REDIS_URL").expect("REDIS_URL must be set");
        let redis_client = RedisClient::open(redis_url).expect("Invalid Redis URL");
        let redis_connection = redis_client
            .get_multiplexed_async_connection()
            .await
            .unwrap();

        // Получаем конфигурацию для S3 (Storj)
        let s3_access_key = env::var("STORJ_ACCESS_KEY").expect("STORJ_ACCESS_KEY must be set");
        let s3_secret_key = env::var("STORJ_SECRET_KEY").expect("STORJ_SECRET_KEY must be set");
        let s3_endpoint = env::var("STORJ_END_POINT")
            .unwrap_or_else(|_| "https://gateway.storjshare.io".to_string());
        let bucket = env::var("STORJ_BUCKET_NAME").unwrap_or_else(|_| "discours-io".to_string());

        // Получаем конфигурацию для AWS S3
        let aws_access_key = env::var("AWS_ACCESS_KEY").expect("AWS_ACCESS_KEY must be set");
        let aws_secret_key = env::var("AWS_SECRET_KEY").expect("AWS_SECRET_KEY must be set");
        let aws_endpoint =
            env::var("AWS_END_POINT").unwrap_or_else(|_| "https://s3.amazonaws.com".to_string());

        // Конфигурируем клиент S3 для Storj
        let storj_config = aws_config::defaults(BehaviorVersion::latest())
            .region("eu-west-1")
            .endpoint_url(s3_endpoint)
            .credentials_provider(Credentials::new(
                s3_access_key,
                s3_secret_key,
                None,
                None,
                "rust-storj-client",
            ))
            .load()
            .await;

        let storj_client = S3Client::new(&storj_config);

        // Конфигурируем клиент S3 для AWS
        let aws_config = aws_config::defaults(BehaviorVersion::latest())
            .region("eu-west-1")
            .endpoint_url(aws_endpoint)
            .credentials_provider(Credentials::new(
                aws_access_key,
                aws_secret_key,
                None,
                None,
                "rust-aws-client",
            ))
            .load()
            .await;

        let aws_client = S3Client::new(&aws_config);

        let app_state = AppState {
            redis: redis_connection,
            storj_client,
            aws_client,
            bucket,
        };

        // Кэшируем список файлов из AWS при старте приложения
        app_state.cache_filelist().await;

        app_state
    }

    /// Кэширует список файлов из Storj S3 в Redis.
    pub async fn cache_filelist(&self) {
        warn!("caching AWS filelist...");
        let mut redis = self.redis.clone();

        // Запрашиваем список файлов из Storj S3
        let filelist = get_s3_filelist(&self.aws_client, &self.bucket).await;

        for [filename, filepath] in filelist.clone() {
            // Сохраняем список файлов в Redis, используя HSET для каждого файла
            let _: () = redis
                .hset(PATH_MAPPING_KEY, filename.clone(), filepath)
                .await
                .unwrap();
        }

        warn!("cached {} files", filelist.len());
    }

    /// Получает путь из ключа (имени файла) в Redis.
    pub async fn get_path(&self, filename: &str) -> Result<Option<String>, actix_web::Error> {
        let mut redis = self.redis.clone();
        let new_path: Option<String> = redis
            .hget(PATH_MAPPING_KEY, filename)
            .await
            .map_err(|_| ErrorInternalServerError("Failed to get path mapping from Redis"))?;
        Ok(new_path)
    }

    pub async fn set_path(&self, filename: &str, filepath: &str) {
        let mut redis = self.redis.clone();
        let _: () = redis
            .hset(PATH_MAPPING_KEY, filename, filepath)
            .await
            .unwrap_or_else(|_| panic!("Failed to cache file {} in Redis", filename));
    }

    /// создает или получает текущее значение квоты пользователя
    pub async fn get_or_create_quota(&self, user_id: &str) -> Result<u64, actix_web::Error> {
        let mut redis = self.redis.clone();
        let quota_key = format!("quota:{}", user_id);

        // Попытка получить квоту из Redis
        let quota: u64 = redis.get(&quota_key).await.unwrap_or(0);

        if quota == 0 {
            // Если квота не найдена, устанавливаем её в 0 байт без TTL (постоянная квота)
            redis
                .set::<&str, u64, ()>(&quota_key, 0)
                .await
                .map_err(|_| {
                    ErrorInternalServerError("Failed to set initial user quota in Redis")
                })?;

            Ok(0) // Возвращаем 0 как начальную квоту
        } else {
            Ok(quota)
        }
    }

    /// инкрементирует значение квоты пользователя в байтах
    pub async fn increment_uploaded_bytes(
        &self,
        user_id: &str,
        bytes: u64,
    ) -> Result<u64, actix_web::Error> {
        let mut redis = self.redis.clone();
        let quota_key = format!("quota:{}", user_id);

        // Проверяем, существует ли ключ в Redis
        let exists: bool = redis.exists::<_, bool>(&quota_key).await.map_err(|_| {
            ErrorInternalServerError("Failed to check if user quota exists in Redis")
        })?;

        // Если ключ не существует, создаем его с начальным значением без TTL
        if !exists {
            redis
                .set::<_, u64, ()>(&quota_key, bytes)
                .await
                .map_err(|_| {
                    ErrorInternalServerError("Failed to set initial user quota in Redis")
                })?;
            return Ok(bytes);
        }

        // Если ключ существует, инкрементируем его значение на заданное количество байт
        let new_quota: u64 = redis
            .incr::<_, u64, u64>(&quota_key, bytes)
            .await
            .map_err(|_| ErrorInternalServerError("Failed to increment user quota in Redis"))?;

        Ok(new_quota)
    }

    /// Устанавливает квоту пользователя в байтах (позволяет увеличить или уменьшить)
    pub async fn set_user_quota(&self, user_id: &str, bytes: u64) -> Result<u64, actix_web::Error> {
        let mut redis = self.redis.clone();
        let quota_key = format!("quota:{}", user_id);

        // Устанавливаем новое значение квоты
        redis
            .set::<_, u64, ()>(&quota_key, bytes)
            .await
            .map_err(|_| ErrorInternalServerError("Failed to set user quota in Redis"))?;

        Ok(bytes)
    }

    /// Увеличивает квоту пользователя на указанное количество байт
    pub async fn increase_user_quota(
        &self,
        user_id: &str,
        additional_bytes: u64,
    ) -> Result<u64, actix_web::Error> {
        let mut redis = self.redis.clone();
        let quota_key = format!("quota:{}", user_id);

        // Получаем текущую квоту
        let current_quota: u64 = redis.get(&quota_key).await.unwrap_or(0);

        // Вычисляем новую квоту
        let new_quota = current_quota + additional_bytes;

        // Устанавливаем новое значение
        redis
            .set::<_, u64, ()>(&quota_key, new_quota)
            .await
            .map_err(|_| ErrorInternalServerError("Failed to increase user quota in Redis"))?;

        Ok(new_quota)
    }
}