gitdataai/libs/service/user/chpc.rs

use crate::AppService;
use crate::error::AppError;
use chrono::{Duration, Local, NaiveDate};
use models::repos::repo_commit;
use models::users::{user, user_email};
use redis::AsyncCommands;
use sea_orm::*;
use serde::{Deserialize, Serialize};
use session::Session;
use utoipa::{IntoParams, ToSchema};

/// Cache key prefix for contribution heatmap
const HEATMAP_CACHE_PREFIX: &str = "user:heatmap";
/// Default cache TTL in seconds (5 minutes)
const HEATMAP_CACHE_TTL: u64 = 300;

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ContributionHeatmapItem {
    pub date: String,
    pub count: i32,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ContributionHeatmapResponse {
    pub username: String,
    pub total_contributions: i64,
    pub heatmap: Vec<ContributionHeatmapItem>,
    pub start_date: String,
    pub end_date: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, IntoParams)]
pub struct ContributionHeatmapQuery {
    pub start_date: Option<String>,
    pub end_date: Option<String>,
}

impl AppService {
    pub async fn get_user_contribution_heatmap(
        &self,
        _context: Session,
        username: String,
        query: ContributionHeatmapQuery,
    ) -> Result<ContributionHeatmapResponse, AppError> {
        let user = user::Entity::find()
            .filter(user::Column::Username.eq(&username))
            .one(&self.db)
            .await?
            .ok_or(AppError::UserNotFound)?;

        let (start_date, end_date) = self.parse_date_range(query.start_date, query.end_date)?;

        let cache_key = self.build_heatmap_cache_key(&user.uid, start_date, end_date);
        if let Ok(mut conn) = self.cache.conn().await {
            if let Ok(cached) = conn.get::<_, String>(cache_key.clone()).await {
                if let Ok(cached) = serde_json::from_str::<ContributionHeatmapResponse>(&cached) {
                    return Ok(cached);
                }
            }
        }

        let emails: Vec<String> = user_email::Entity::find()
            .filter(user_email::Column::User.eq(user.uid))
            .select_only()
            .column(user_email::Column::Email)
            .into_tuple::<String>()
            .all(&self.db)
            .await?;

        if emails.is_empty() {
            let response = ContributionHeatmapResponse {
                username: user.username.clone(),
                total_contributions: 0,
                heatmap: vec![],
                start_date: start_date.format("%Y-%m-%d").to_string(),
                end_date: end_date.format("%Y-%m-%d").to_string(),
            };
            return Ok(response);
        }

        let start_dt = start_date.and_hms_opt(0, 0, 0).unwrap();
        let end_dt = end_date.and_hms_opt(23, 59, 59).unwrap();

        let commits: Vec<repo_commit::Model> = repo_commit::Entity::find()
            .filter(repo_commit::Column::AuthorEmail.is_in(emails.clone()))
            .filter(repo_commit::Column::CreatedAt.gte(start_dt))
            .filter(repo_commit::Column::CreatedAt.lte(end_dt))
            .all(&self.db)
            .await?;

        let mut heatmap_map: std::collections::HashMap<String, i64> =
            std::collections::HashMap::new();
        for commit in &commits {
            let date_str = commit.created_at.format("%Y-%m-%d").to_string();
            *heatmap_map.entry(date_str).or_insert(0) += 1;
        }

        let total_contributions = commits.len() as i64;

        let mut heatmap: Vec<ContributionHeatmapItem> = Vec::new();
        let mut current = start_date;
        while current <= end_date {
            let date_str = current.format("%Y-%m-%d").to_string();
            let count = *heatmap_map.get(&date_str).unwrap_or(&0);
            heatmap.push(ContributionHeatmapItem {
                date: date_str,
                count: count as i32,
            });
            current += Duration::days(1);
        }

        let response = ContributionHeatmapResponse {
            username: user.username,
            total_contributions,
            heatmap,
            start_date: start_date.format("%Y-%m-%d").to_string(),
            end_date: end_date.format("%Y-%m-%d").to_string(),
        };

        if let Ok(mut conn) = self.cache.conn().await {
            let _: Option<()> = conn
                .set_ex::<String, String, ()>(
                    cache_key,
                    serde_json::to_string(&response).unwrap_or_default(),
                    HEATMAP_CACHE_TTL,
                )
                .await
                .ok();
        }

        Ok(response)
    }

    pub async fn invalidate_user_heatmap_cache(
        &self,
        user_uid: uuid::Uuid,
    ) -> Result<(), AppError> {
        // Invalidate all heatmap cache entries for a user
        // Delete known date range keys (last 2 years)
        if let Ok(mut conn) = self.cache.conn().await {
            let today = Local::now().date_naive();
            let two_years_ago = today - Duration::days(730);
            let mut current = two_years_ago;
            while current <= today {
                let key = self.build_heatmap_cache_key(&user_uid, current, current);
                let _: Option<()> = conn.del::<_, ()>(key).await.ok();
                current += Duration::days(1);
            }
        }
        Ok(())
    }

    fn build_heatmap_cache_key(
        &self,
        user_uid: &uuid::Uuid,
        start_date: NaiveDate,
        end_date: NaiveDate,
    ) -> String {
        format!(
            "{}:{}:{}:{}",
            HEATMAP_CACHE_PREFIX,
            user_uid,
            start_date.format("%Y-%m-%d"),
            end_date.format("%Y-%m-%d"),
        )
    }

    fn parse_date_range(
        &self,
        start_date_str: Option<String>,
        end_date_str: Option<String>,
    ) -> Result<(NaiveDate, NaiveDate), AppError> {
        let today = Local::now().date_naive();
        let one_year_ago = today - Duration::days(365);

        let start_date = if let Some(date_str) = start_date_str {
            NaiveDate::parse_from_str(&date_str, "%Y-%m-%d").map_err(|_| {
                AppError::NotFound("Invalid start_date format, expected YYYY-MM-DD".to_string())
            })?
        } else {
            one_year_ago
        };

        let end_date = if let Some(date_str) = end_date_str {
            NaiveDate::parse_from_str(&date_str, "%Y-%m-%d").map_err(|_| {
                AppError::NotFound("Invalid end_date format, expected YYYY-MM-DD".to_string())
            })?
        } else {
            today
        };

        if start_date > end_date {
            return Err(AppError::NotFound(
                "start_date cannot be later than end_date".to_string(),
            ));
        }

        if (end_date - start_date).num_days() > 730 {
            return Err(AppError::NotFound(
                "Date range cannot exceed 2 years".to_string(),
            ));
        }

        Ok((start_date, end_date))
    }

    pub async fn get_current_user_contribution_heatmap(
        &self,
        context: Session,
        query: ContributionHeatmapQuery,
    ) -> Result<ContributionHeatmapResponse, AppError> {
        let user_uid = context.user().ok_or(AppError::Unauthorized)?;

        let user = user::Entity::find()
            .filter(user::Column::Uid.eq(user_uid))
            .one(&self.db)
            .await?
            .ok_or(AppError::UserNotFound)?;

        self.get_user_contribution_heatmap(context, user.username, query)
            .await
    }
}