gitdataai/lib/service/user/chpc.rs

use std::collections::HashMap;

use chrono::{Duration, NaiveDate, Utc};
use db::sqlx;
use serde::{Deserialize, Serialize};
use session::Session;
use utoipa::{IntoParams, ToSchema};

use crate::{AppService, error::AppError};

const HEATMAP_CACHE_PREFIX: &str = "user:heatmap";

#[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 user_chpc(
        &self,
        ctx: &Session,
        query: ContributionHeatmapQuery,
    ) -> Result<ContributionHeatmapResponse, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        let user = self.auth_find_user_by_uid(user_uid).await?;
        self.user_contribution_heatmap_for_user(user.id, user.username, query)
            .await
    }

    pub async fn user_contribution_heatmap_for_user(
        &self,
        user_uid: uuid::Uuid,
        username: String,
        query: ContributionHeatmapQuery,
    ) -> Result<ContributionHeatmapResponse, AppError> {
        let (start_date, end_date) =
            parse_date_range(query.start_date, query.end_date)?;
        let cache_key = build_heatmap_cache_key(user_uid, start_date, end_date);

        if let Ok(Some(cached)) = self
            .cache
            .get::<ContributionHeatmapResponse>(&cache_key)
            .await
        {
            return Ok(cached);
        }

        let start_dt = start_date
            .and_hms_opt(0, 0, 0)
            .ok_or(AppError::InternalError)?
            .and_utc();
        let end_dt = end_date
            .and_hms_opt(23, 59, 59)
            .ok_or(AppError::InternalError)?
            .and_utc();

        let rows = sqlx::query_as::<_, (chrono::DateTime<Utc>,)>(
            "SELECT created_at FROM repo_commit \
             WHERE (author = $1 OR committer = $1) AND created_at >= $2 AND created_at <= $3",
        )
        .bind(user_uid)
        .bind(start_dt)
        .bind(end_dt)
        .fetch_all(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        let mut sparse = HashMap::<String, i64>::new();
        for (created_at,) in &rows {
            let date = created_at.format("%Y-%m-%d").to_string();
            *sparse.entry(date).or_insert(0) += 1;
        }

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

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

        self.cache
            .set(&cache_key, &response)
            .await
            .map_err(|e| AppError::InternalServerError(e.to_string()))?;

        Ok(response)
    }

    pub async fn user_invalidate_chpc_cache(
        &self,
        ctx: &Session,
    ) -> Result<(), AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        self.invalidate_user_heatmap_cache(user_uid).await
    }

    pub async fn invalidate_user_heatmap_cache(
        &self,
        user_uid: uuid::Uuid,
    ) -> Result<(), AppError> {
        let pattern = format!("{}:{}:*", HEATMAP_CACHE_PREFIX, user_uid);
        self.cache
            .delete_pattern(&pattern)
            .await
            .map_err(|e| AppError::InternalServerError(e.to_string()))?;
        Ok(())
    }
}

fn build_heatmap_cache_key(
    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(
    start_date_str: Option<String>,
    end_date_str: Option<String>,
) -> Result<(NaiveDate, NaiveDate), AppError> {
    let today = Utc::now().date_naive();
    let one_year_ago = today - Duration::days(365);

    let start_date = match start_date_str {
        Some(date) => parse_date(&date, "start_date")?,
        None => one_year_ago,
    };
    let end_date = match end_date_str {
        Some(date) => parse_date(&date, "end_date")?,
        None => today,
    };

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

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

    Ok((start_date, end_date))
}

fn parse_date(value: &str, field: &str) -> Result<NaiveDate, AppError> {
    NaiveDate::parse_from_str(value, "%Y-%m-%d").map_err(|_| {
        AppError::BadRequest(format!(
            "invalid {field} format, expected YYYY-MM-DD"
        ))
    })
}