Genarrative/server-rs/crates/spacetime-module/src/admin_dashboard.rs

use std::collections::{BTreeMap, BTreeSet, HashSet};

use crate::*;

const ADMIN_DASHBOARD_MAX_RANGE_DAYS: i64 = 366;
const ADMIN_DASHBOARD_ACTIVE_WINDOW_MICROS: i64 = 5 * 60 * 1_000_000;

#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardStatsInput {
    pub start_day_key: i64,
    pub end_day_key: i64,
    pub now_micros: i64,
}

#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardDailyCountSnapshot {
    pub day_key: i64,
    pub count: u64,
}

#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardBreakdownSnapshot {
    pub key: String,
    pub count: u64,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardRetentionSnapshot {
    pub eligible_users: u64,
    pub retained_users: u64,
    pub rate_basis_points: u32,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardPaymentConversionSnapshot {
    pub paid_users: u64,
    pub new_registered_users: u64,
    pub rate_basis_points: u32,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardStatsSnapshot {
    pub total_registered_users: u64,
    pub new_registered_users: Vec<AdminDashboardDailyCountSnapshot>,
    pub visit_count: Vec<AdminDashboardDailyCountSnapshot>,
    pub visit_users_by_day: Vec<AdminDashboardDailyCountSnapshot>,
    pub visit_users: u64,
    pub total_visit_count: u64,
    pub total_visit_users: u64,
    pub current_users: u64,
    pub day1_retention: AdminDashboardRetentionSnapshot,
    pub day7_retention: AdminDashboardRetentionSnapshot,
    pub module_visit_breakdown: Vec<AdminDashboardBreakdownSnapshot>,
    pub new_user_payment_conversion: AdminDashboardPaymentConversionSnapshot,
}

#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct AdminDashboardStatsProcedureResult {
    pub ok: bool,
    pub snapshot: Option<AdminDashboardStatsSnapshot>,
    pub error_message: Option<String>,
}

struct AdminDashboardDailyStatsAccumulator<'a> {
    start_day_key: i64,
    end_day_key: i64,
    selected_registrations: &'a BTreeMap<String, i64>,
    visit_count: BTreeMap<i64, u64>,
    visit_users_by_day: BTreeMap<i64, BTreeSet<String>>,
    period_visit_users: BTreeSet<String>,
    total_visit_users: BTreeSet<String>,
    total_visit_count: u64,
    retained_activity_days: HashSet<(String, i64)>,
}

impl<'a> AdminDashboardDailyStatsAccumulator<'a> {
    fn new(
        start_day_key: i64,
        end_day_key: i64,
        selected_registrations: &'a BTreeMap<String, i64>,
    ) -> Self {
        Self {
            start_day_key,
            end_day_key,
            selected_registrations,
            visit_count: BTreeMap::new(),
            visit_users_by_day: BTreeMap::new(),
            period_visit_users: BTreeSet::new(),
            total_visit_users: BTreeSet::new(),
            total_visit_count: 0,
            retained_activity_days: HashSet::new(),
        }
    }

    fn apply(
        &mut self,
        scope_kind: RuntimeTrackingScopeKind,
        scope_id: &str,
        day_key: i64,
        count: u32,
    ) {
        if count == 0 {
            return;
        }
        match scope_kind {
            RuntimeTrackingScopeKind::Site => {
                let count = u64::from(count);
                self.total_visit_count = self.total_visit_count.saturating_add(count);
                if day_key_is_in_range(day_key, self.start_day_key, self.end_day_key) {
                    increment_count(&mut self.visit_count, day_key, count);
                }
            }
            RuntimeTrackingScopeKind::User => {
                let user_id = scope_id.trim();
                if user_id.is_empty() || user_id == "anonymous" {
                    return;
                }
                self.total_visit_users.insert(user_id.to_string());
                if day_key_is_in_range(day_key, self.start_day_key, self.end_day_key) {
                    self.period_visit_users.insert(user_id.to_string());
                    self.visit_users_by_day
                        .entry(day_key)
                        .or_default()
                        .insert(user_id.to_string());
                }
                if self
                    .selected_registrations
                    .get(user_id)
                    .is_some_and(|registration_day_key| {
                        day_key == registration_day_key.saturating_add(1)
                            || day_key == registration_day_key.saturating_add(7)
                    })
                {
                    self.retained_activity_days
                        .insert((user_id.to_string(), day_key));
                }
            }
            RuntimeTrackingScopeKind::Work | RuntimeTrackingScopeKind::Module => {}
        }
    }
}

struct AdminDashboardEventAccumulator {
    start_day_key: i64,
    end_day_key: i64,
    active_since_micros: i64,
    now_micros: i64,
    current_users: BTreeSet<String>,
    module_visit_breakdown: BTreeMap<String, u64>,
}

struct AdminDashboardPaymentConversionAccumulator<'a> {
    selected_registrations: &'a BTreeMap<String, i64>,
    now_micros: i64,
    paid_users: BTreeSet<String>,
}

impl<'a> AdminDashboardPaymentConversionAccumulator<'a> {
    fn new(selected_registrations: &'a BTreeMap<String, i64>, now_micros: i64) -> Self {
        Self {
            selected_registrations,
            now_micros,
            paid_users: BTreeSet::new(),
        }
    }

    fn apply(&mut self, user_id: &str, paid_at_micros: Option<i64>) {
        let Some(paid_at_micros) = paid_at_micros else {
            return;
        };
        if !self.selected_registrations.contains_key(user_id) {
            return;
        }
        if paid_at_micros > self.now_micros {
            return;
        }
        self.paid_users.insert(user_id.to_string());
    }

    fn snapshot(self) -> AdminDashboardPaymentConversionSnapshot {
        let paid_users = self.paid_users.len() as u64;
        let new_registered_users = self.selected_registrations.len() as u64;
        AdminDashboardPaymentConversionSnapshot {
            paid_users,
            new_registered_users,
            rate_basis_points: payment_conversion_rate_basis_points(
                paid_users,
                new_registered_users,
            ),
        }
    }
}

impl AdminDashboardEventAccumulator {
    fn new(start_day_key: i64, end_day_key: i64, now_micros: i64) -> Self {
        Self {
            start_day_key,
            end_day_key,
            active_since_micros: now_micros.saturating_sub(ADMIN_DASHBOARD_ACTIVE_WINDOW_MICROS),
            now_micros,
            current_users: BTreeSet::new(),
            module_visit_breakdown: BTreeMap::new(),
        }
    }

    fn apply(
        &mut self,
        user_id: Option<&str>,
        module_key: Option<&str>,
        day_key: i64,
        occurred_at_micros: i64,
    ) {
        if day_key_is_in_range(day_key, self.start_day_key, self.end_day_key) {
            let module_key = module_key
                .map(str::trim)
                .filter(|value| !value.is_empty())
                .unwrap_or("unknown");
            increment_string_count(&mut self.module_visit_breakdown, module_key, 1);
        }

        if occurred_at_micros < self.active_since_micros || occurred_at_micros > self.now_micros {
            return;
        }
        if let Some(user_id) = user_id
            .map(str::trim)
            .filter(|value| !value.is_empty() && *value != "anonymous")
        {
            self.current_users.insert(user_id.to_string());
        }
    }
}

#[spacetimedb::procedure]
pub fn get_admin_dashboard_stats_and_return(
    ctx: &mut ProcedureContext,
    input: AdminDashboardStatsInput,
) -> AdminDashboardStatsProcedureResult {
    let caller = ctx.sender();
    match ctx.try_with_tx(|tx| {
        crate::editor_project_storage::require_editor_generation_runtime_service_identity(
            tx, caller,
        )?;
        build_admin_dashboard_stats(tx, &input)
    }) {
        Ok(snapshot) => AdminDashboardStatsProcedureResult {
            ok: true,
            snapshot: Some(snapshot),
            error_message: None,
        },
        Err(message) => AdminDashboardStatsProcedureResult {
            ok: false,
            snapshot: None,
            error_message: Some(message),
        },
    }
}

fn build_admin_dashboard_stats(
    ctx: &ReducerContext,
    input: &AdminDashboardStatsInput,
) -> Result<AdminDashboardStatsSnapshot, String> {
    validate_admin_dashboard_stats_input(input)?;

    let current_day_key = runtime_profile_beijing_day_key(input.now_micros);
    let mut new_registered_users = BTreeMap::new();
    let mut selected_registrations = BTreeMap::new();
    let mut total_registered_users = 0u64;

    for row in ctx.db.profile_dashboard_state().iter() {
        total_registered_users = total_registered_users.saturating_add(1);
        let registration_micros = row.created_at.to_micros_since_unix_epoch();
        let registration_day_key = runtime_profile_beijing_day_key(registration_micros);
        if day_key_is_in_range(registration_day_key, input.start_day_key, input.end_day_key) {
            increment_count(&mut new_registered_users, registration_day_key, 1);
            selected_registrations.insert(row.user_id, registration_day_key);
        }
    }

    let mut payment_conversion =
        AdminDashboardPaymentConversionAccumulator::new(&selected_registrations, input.now_micros);
    for user_id in selected_registrations.keys() {
        for row in ctx
            .db
            .profile_recharge_order()
            .by_profile_recharge_order_user_id()
            .filter(user_id.as_str())
        {
            payment_conversion.apply(
                &row.user_id,
                row.paid_at.map(|value| value.to_micros_since_unix_epoch()),
            );
        }
    }

    let mut daily_stats = AdminDashboardDailyStatsAccumulator::new(
        input.start_day_key,
        input.end_day_key,
        &selected_registrations,
    );
    for row in ctx.db.tracking_daily_stat().iter() {
        daily_stats.apply(row.scope_kind, &row.scope_id, row.day_key, row.count);
    }

    let mut event_stats = AdminDashboardEventAccumulator::new(
        input.start_day_key,
        input.end_day_key,
        input.now_micros,
    );
    for row in ctx.db.tracking_event().iter() {
        event_stats.apply(
            row.user_id.as_deref(),
            row.module_key.as_deref(),
            row.day_key,
            row.occurred_at.to_micros_since_unix_epoch(),
        );
    }

    Ok(AdminDashboardStatsSnapshot {
        total_registered_users,
        new_registered_users: daily_counts(new_registered_users),
        visit_count: daily_counts(daily_stats.visit_count),
        visit_users_by_day: daily_unique_counts(daily_stats.visit_users_by_day),
        visit_users: daily_stats.period_visit_users.len() as u64,
        total_visit_count: daily_stats.total_visit_count,
        total_visit_users: daily_stats.total_visit_users.len() as u64,
        current_users: event_stats.current_users.len() as u64,
        day1_retention: calculate_admin_dashboard_retention(
            &selected_registrations,
            &daily_stats.retained_activity_days,
            current_day_key,
            1,
        ),
        day7_retention: calculate_admin_dashboard_retention(
            &selected_registrations,
            &daily_stats.retained_activity_days,
            current_day_key,
            7,
        ),
        module_visit_breakdown: event_stats
            .module_visit_breakdown
            .into_iter()
            .map(|(key, count)| AdminDashboardBreakdownSnapshot { key, count })
            .collect(),
        new_user_payment_conversion: payment_conversion.snapshot(),
    })
}

fn validate_admin_dashboard_stats_input(input: &AdminDashboardStatsInput) -> Result<(), String> {
    if input.start_day_key > input.end_day_key {
        return Err("dashboard start_day_key 不能晚于 end_day_key".to_string());
    }
    if input.end_day_key.saturating_sub(input.start_day_key) + 1 > ADMIN_DASHBOARD_MAX_RANGE_DAYS {
        return Err(format!(
            "dashboard 查询范围最多支持 {ADMIN_DASHBOARD_MAX_RANGE_DAYS} 天"
        ));
    }
    if input.now_micros <= 0 {
        return Err("dashboard now_micros 必须大于 0".to_string());
    }
    let current_day_key = runtime_profile_beijing_day_key(input.now_micros);
    if input.end_day_key > current_day_key {
        return Err("dashboard end_day_key 不能晚于今天".to_string());
    }
    Ok(())
}

fn calculate_admin_dashboard_retention(
    registrations: &BTreeMap<String, i64>,
    active_user_days: &HashSet<(String, i64)>,
    current_day_key: i64,
    retention_days: i64,
) -> AdminDashboardRetentionSnapshot {
    let mut eligible_users = 0u64;
    let mut retained_users = 0u64;
    for (user_id, registration_day_key) in registrations {
        let retention_day_key = registration_day_key.saturating_add(retention_days);
        // 今天仍是进行中的业务日，只把目标日已完整结束的 cohort 纳入分母。
        if retention_day_key >= current_day_key {
            continue;
        }
        eligible_users = eligible_users.saturating_add(1);
        if active_user_days.contains(&(user_id.clone(), retention_day_key)) {
            retained_users = retained_users.saturating_add(1);
        }
    }
    AdminDashboardRetentionSnapshot {
        eligible_users,
        retained_users,
        rate_basis_points: retention_rate_basis_points(retained_users, eligible_users),
    }
}

fn retention_rate_basis_points(retained_users: u64, eligible_users: u64) -> u32 {
    if eligible_users == 0 {
        return 0;
    }
    let numerator = u128::from(retained_users)
        .saturating_mul(10_000)
        .saturating_add(u128::from(eligible_users) / 2);
    (numerator / u128::from(eligible_users)).min(10_000) as u32
}

fn payment_conversion_rate_basis_points(paid_users: u64, new_registered_users: u64) -> u32 {
    if new_registered_users == 0 {
        return 0;
    }
    ((u128::from(paid_users) * 10_000 + u128::from(new_registered_users) / 2)
        / u128::from(new_registered_users))
    .min(10_000) as u32
}

fn day_key_is_in_range(day_key: i64, start_day_key: i64, end_day_key: i64) -> bool {
    day_key >= start_day_key && day_key <= end_day_key
}

fn increment_count(counts: &mut BTreeMap<i64, u64>, key: i64, amount: u64) {
    let next = counts
        .get(&key)
        .copied()
        .unwrap_or(0)
        .saturating_add(amount);
    counts.insert(key, next);
}

fn increment_string_count(counts: &mut BTreeMap<String, u64>, key: &str, amount: u64) {
    let next = counts.get(key).copied().unwrap_or(0).saturating_add(amount);
    counts.insert(key.to_string(), next);
}

fn daily_counts(counts: BTreeMap<i64, u64>) -> Vec<AdminDashboardDailyCountSnapshot> {
    counts
        .into_iter()
        .map(|(day_key, count)| AdminDashboardDailyCountSnapshot { day_key, count })
        .collect()
}

fn daily_unique_counts(
    users_by_day: BTreeMap<i64, BTreeSet<String>>,
) -> Vec<AdminDashboardDailyCountSnapshot> {
    users_by_day
        .into_iter()
        .map(|(day_key, users)| AdminDashboardDailyCountSnapshot {
            day_key,
            count: users.len() as u64,
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::{
        AdminDashboardDailyStatsAccumulator, AdminDashboardEventAccumulator,
        AdminDashboardPaymentConversionAccumulator, AdminDashboardStatsInput,
        RuntimeTrackingScopeKind, calculate_admin_dashboard_retention, daily_unique_counts,
        payment_conversion_rate_basis_points, retention_rate_basis_points,
        validate_admin_dashboard_stats_input,
    };
    use std::collections::{BTreeMap, BTreeSet, HashSet};

    #[test]
    fn retention_only_counts_mature_cohorts_and_exact_return_day() {
        let registrations = BTreeMap::from([
            ("day1-returned".to_string(), 100),
            ("day1-missed".to_string(), 100),
            ("day7-returned".to_string(), 94),
            ("not-mature".to_string(), 101),
        ]);
        let active_user_days = HashSet::from([
            ("day1-returned".to_string(), 101),
            ("day1-missed".to_string(), 102),
            ("day7-returned".to_string(), 101),
            ("not-mature".to_string(), 101),
        ]);

        let day1 = calculate_admin_dashboard_retention(&registrations, &active_user_days, 102, 1);
        let day7 = calculate_admin_dashboard_retention(&registrations, &active_user_days, 102, 7);

        assert_eq!(day1.eligible_users, 3);
        assert_eq!(day1.retained_users, 1);
        assert_eq!(day1.rate_basis_points, 3_333);
        assert_eq!(day7.eligible_users, 1);
        assert_eq!(day7.retained_users, 1);
        assert_eq!(day7.rate_basis_points, 10_000);
    }

    #[test]
    fn retention_rate_is_zero_without_eligible_users_and_rounds_to_basis_points() {
        assert_eq!(retention_rate_basis_points(0, 0), 0);
        assert_eq!(retention_rate_basis_points(2, 3), 6_667);
        assert_eq!(retention_rate_basis_points(4, 3), 10_000);
    }

    #[test]
    fn new_user_payment_conversion_deduplicates_paid_users_through_query_time() {
        let registrations = BTreeMap::from([
            ("paid".to_string(), 100),
            ("refunded-with-paid-at".to_string(), 100),
            ("pending".to_string(), 100),
            ("paid-late".to_string(), 200),
        ]);
        let mut conversion = AdminDashboardPaymentConversionAccumulator::new(&registrations, 500);

        conversion.apply("paid", Some(99));
        conversion.apply("paid", Some(150));
        conversion.apply("paid", Some(160));
        conversion.apply("refunded-with-paid-at", Some(180));
        conversion.apply("pending", None);
        conversion.apply("paid-late", Some(501));
        conversion.apply("outside-cohort", Some(150));

        let snapshot = conversion.snapshot();
        assert_eq!(snapshot.paid_users, 2);
        assert_eq!(snapshot.new_registered_users, 4);
        assert_eq!(snapshot.rate_basis_points, 5_000);
    }

    #[test]
    fn new_user_payment_conversion_rate_is_zero_without_new_users_and_rounds() {
        assert_eq!(payment_conversion_rate_basis_points(0, 0), 0);
        assert_eq!(payment_conversion_rate_basis_points(2, 3), 6_667);
        assert_eq!(payment_conversion_rate_basis_points(4, 3), 10_000);
    }

    #[test]
    fn daily_visit_users_are_deduplicated_per_day_without_cross_day_collapsing() {
        let counts = daily_unique_counts(BTreeMap::from([
            (
                100,
                BTreeSet::from(["user-a".to_string(), "user-b".to_string()]),
            ),
            (101, BTreeSet::from(["user-a".to_string()])),
        ]));

        assert_eq!(counts.len(), 2);
        assert_eq!(counts[0].day_key, 100);
        assert_eq!(counts[0].count, 2);
        assert_eq!(counts[1].day_key, 101);
        assert_eq!(counts[1].count, 1);
    }

    #[test]
    fn dashboard_tracking_aggregation_stays_exact_past_legacy_row_limit() {
        let registrations = BTreeMap::new();
        let mut daily = AdminDashboardDailyStatsAccumulator::new(100, 100, &registrations);
        let mut events = AdminDashboardEventAccumulator::new(100, 100, 1_000_000);

        for index in 0..50_001 {
            daily.apply(RuntimeTrackingScopeKind::Site, "site", 100, 1);
            daily.apply(
                RuntimeTrackingScopeKind::User,
                &format!("user-{index}"),
                100,
                1,
            );
            events.apply(Some("current-user"), Some("profile"), 100, 999_999);
        }

        assert_eq!(daily.total_visit_count, 50_001);
        assert_eq!(daily.total_visit_users.len(), 50_001);
        assert_eq!(daily.period_visit_users.len(), 50_001);
        assert_eq!(daily.visit_count[&100], 50_001);
        assert_eq!(events.module_visit_breakdown["profile"], 50_001);
        assert_eq!(events.current_users.len(), 1);
    }

    #[test]
    fn dashboard_stats_input_rejects_future_and_oversized_ranges() {
        const DAY_MICROS: i64 = 86_400_000_000;
        let now_micros = 200 * DAY_MICROS;
        let future = validate_admin_dashboard_stats_input(&AdminDashboardStatsInput {
            start_day_key: 200,
            end_day_key: 201,
            now_micros,
        });
        let oversized = validate_admin_dashboard_stats_input(&AdminDashboardStatsInput {
            start_day_key: 1,
            end_day_key: 367,
            now_micros: 500 * DAY_MICROS,
        });

        assert!(future.is_err());
        assert!(oversized.is_err());
    }
}