nxtgauge-backend-rust/apps/users/src/ai_subscription.rs

//! Subscription lifecycle management for AI credits
//! 
//! Handles plan upgrades, downgrades, cancellations, and renewals

use db::models::ai_credits::{AiCreditsError, AiCreditsRepository, AiPlan};
use sqlx::PgPool;
use uuid::Uuid;
use chrono::{DateTime, Utc, Duration};

/// Calculate prorated credits for a mid-period upgrade
/// 
/// Formula: remaining_days / total_days * (new_monthly - old_monthly)
fn calculate_proration(
    current_period_end: DateTime<Utc>,
    old_monthly_credits: i32,
    new_monthly_credits: i32,
) -> i32 {
    if new_monthly_credits <= old_monthly_credits {
        return 0;
    }
    
    let now = Utc::now();
    let total_days = (current_period_end - now).num_days();
    let period_length = 30; // Monthly period assumption
    
    if total_days <= 0 {
        return 0;
    }
    
    let remaining_ratio = total_days as f64 / period_length as f64;
    let credit_diff = (new_monthly_credits - old_monthly_credits) as f64;
    
    (remaining_ratio * credit_diff).ceil() as i32
}

/// Upgrade a user's plan immediately
/// 
/// This applies the new plan immediately and grants prorated bonus credits.
pub async fn upgrade_plan(
    pool: &PgPool,
    user_id: Uuid,
    new_plan_id: Uuid,
    actor_id: Option<Uuid>,
) -> Result<(), AiCreditsError> {
    let mut tx = pool.begin().await?;
    
    // Get current wallet
    let wallet = sqlx::query_as::<_, db::models::ai_credits::AiWallet>(
        "SELECT * FROM user_ai_subscriptions WHERE user_id = $1 FOR UPDATE"
    )
    .bind(user_id)
    .fetch_one(&mut *tx)
    .await?;
    
    // Get new plan details
    let new_plan: Option<AiPlan> = sqlx::query_as::<_, AiPlan>(
        "SELECT id, code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features, is_active FROM ai_plans WHERE id = $1 AND is_active = TRUE"
    )
    .bind(new_plan_id)
    .fetch_optional(&mut *tx)
    .await?;
    
    let Some(new_plan) = new_plan else {
        tx.rollback().await?;
        return Err(AiCreditsError::UnknownFeature("Plan not found".to_string()));
    };
    
    // Get old plan
    let old_plan: Option<AiPlan> = sqlx::query_as::<_, AiPlan>(
        "SELECT id, code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features, is_active FROM ai_plans WHERE id = $1"
    )
    .bind(wallet.plan_id)
    .fetch_optional(&mut *tx)
    .await?;
    
    // Calculate proration
    let proration_credits = if let Some(ref old_plan) = old_plan {
        calculate_proration(
            wallet.current_period_end,
            old_plan.monthly_credits,
            new_plan.monthly_credits,
        )
    } else {
        0
    };
    
    // Update plan
    sqlx::query(
        r#"
        UPDATE user_ai_subscriptions 
        SET plan_id = $1, 
            monthly_credits_total = $2,
            bonus_credits_total = bonus_credits_total + $3,
            downgrade_scheduled_to = NULL,
            updated_at = NOW()
        WHERE id = $4
        "#
    )
    .bind(new_plan_id)
    .bind(new_plan.monthly_credits)
    .bind(proration_credits)
    .bind(wallet.id)
    .execute(&mut *tx)
    .await?;
    
    // Record in history
    sqlx::query(
        r#"
        INSERT INTO ai_subscription_history 
            (user_id, from_plan_id, to_plan_id, change_type, proration_credits, 
             proration_days_remaining, effective_at, created_by, status)
        VALUES ($1, $2, $3, 'upgrade', $4, $5, NOW(), $6, 'completed')
        "#
    )
    .bind(user_id)
    .bind(wallet.plan_id)
    .bind(new_plan_id)
    .bind(proration_credits)
    .bind((wallet.current_period_end - Utc::now()).num_days() as i32)
    .bind(actor_id)
    .execute(&mut *tx)
    .await?;
    
    // Create ledger entry for proration bonus
    if proration_credits > 0 {
        let balance_after = wallet.available_credits() + proration_credits;
        sqlx::query(
            r#"
            INSERT INTO ai_credit_ledger
                (wallet_id, entry_type, credits, balance_after, reference_type, actor_type, metadata)
            VALUES ($1, 'subscription_grant', $2, $3, 'plan_upgrade', 'system', 
                    jsonb_build_object('reason', 'upgrade_proration', 'plan_code', $4))
            "#
        )
        .bind(wallet.id)
        .bind(proration_credits)
        .bind(balance_after)
        .bind(&new_plan.code)
        .execute(&mut *tx)
        .await?;
    }
    
    tx.commit().await?;
    Ok(())
}

/// Schedule a downgrade for the end of the current period
/// 
/// The downgrade doesn't happen immediately - it will be applied
/// when the current period ends.
pub async fn schedule_downgrade(
    pool: &PgPool,
    user_id: Uuid,
    downgrade_to_plan_id: Uuid,
    actor_id: Option<Uuid>,
) -> Result<DateTime<Utc>, AiCreditsError> {
    let wallet = AiCreditsRepository::get_wallet(pool, user_id)
        .await?
        .ok_or(AiCreditsError::ReservationNotFound)?;
    
    // Get effective date (current period end)
    let effective_at = wallet.current_period_end;
    
    // Update wallet to mark scheduled downgrade
    sqlx::query(
        "UPDATE user_ai_subscriptions SET downgrade_scheduled_to = $1, updated_at = NOW() WHERE id = $2"
    )
    .bind(downgrade_to_plan_id)
    .bind(wallet.id)
    .execute(pool)
    .await?;
    
    // Record scheduled downgrade
    sqlx::query(
        r#"
        INSERT INTO ai_subscription_history 
            (user_id, from_plan_id, to_plan_id, change_type, effective_at, created_by, status)
        VALUES ($1, $2, $3, 'downgrade', $4, $5, 'scheduled')
        "#
    )
    .bind(user_id)
    .bind(wallet.plan_id)
    .bind(downgrade_to_plan_id)
    .bind(effective_at)
    .bind(actor_id)
    .execute(pool)
    .await?;
    
    Ok(effective_at)
}

/// Apply scheduled downgrades that have reached their effective date
/// 
/// This should be called periodically (e.g., by a cron job) to apply
/// downgrades that were scheduled for the current period end.
pub async fn apply_scheduled_downgrades(pool: &PgPool) -> Result<usize, AiCreditsError> {
    let now = Utc::now();
    
    // Find all wallets with scheduled downgrades that should be applied
    let to_downgrade: Vec<(Uuid, Uuid, Uuid)> = sqlx::query_as(
        r#"
        SELECT w.id as wallet_id, w.user_id, w.downgrade_scheduled_to as new_plan_id
        FROM user_ai_subscriptions w
        WHERE w.downgrade_scheduled_to IS NOT NULL 
          AND w.current_period_end <= $1
        "#
    )
    .bind(now)
    .fetch_all(pool)
    .await
    .map_err(AiCreditsError::Db)?;
    
    let mut applied_count = 0;
    
    for (wallet_id, user_id, new_plan_id) in to_downgrade {
        let mut tx = pool.begin().await?;
        
        // Get new plan credits
        let new_monthly_credits: i32 = sqlx::query_scalar(
            "SELECT monthly_credits FROM ai_plans WHERE id = $1"
        )
        .bind(new_plan_id)
        .fetch_one(&mut *tx)
        .await?;
        
        // Update wallet
        sqlx::query(
            r#"
            UPDATE user_ai_subscriptions 
            SET plan_id = downgrade_scheduled_to,
                monthly_credits_total = $1,
                monthly_credits_used = 0, -- Reset usage for new period
                downgrade_scheduled_to = NULL,
                current_period_start = NOW(),
                current_period_end = NOW() + INTERVAL '30 days',
                updated_at = NOW()
            WHERE id = $2
            "#
        )
        .bind(new_monthly_credits)
        .bind(wallet_id)
        .execute(&mut *tx)
        .await?;
        
        // Update history record
        sqlx::query(
            r#"
            UPDATE ai_subscription_history 
            SET status = 'completed'
            WHERE user_id = $1 AND change_type = 'downgrade' AND status = 'scheduled'
            "#
        )
        .bind(user_id)
        .execute(&mut *tx)
        .await?;
        
        tx.commit().await?;
        applied_count += 1;
    }
    
    Ok(applied_count)
}

/// Cancel a subscription (move to free plan)
pub async fn cancel_subscription(
    pool: &PgPool,
    user_id: Uuid,
    actor_id: Option<Uuid>,
) -> Result<(), AiCreditsError> {
    let wallet = AiCreditsRepository::get_wallet(pool, user_id)
        .await?
        .ok_or(AiCreditsError::ReservationNotFound)?;
    
    // Get free plan
    let free_plan: Option<(Uuid, i32)> = sqlx::query_as(
        "SELECT id, monthly_credits FROM ai_plans WHERE code = 'free' AND is_active = TRUE"
    )
    .fetch_optional(pool)
    .await?;
    
    let Some((free_plan_id, free_monthly)) = free_plan else {
        return Err(AiCreditsError::UnknownFeature("Free plan not found".to_string()));
    };
    
    // Update to free plan immediately
    sqlx::query(
        r#"
        UPDATE user_ai_subscriptions 
        SET plan_id = $1,
            monthly_credits_total = $2,
            monthly_credits_used = 0,
            downgrade_scheduled_to = NULL,
            status = 'cancelled',
            updated_at = NOW()
        WHERE user_id = $3
        "#
    )
    .bind(free_plan_id)
    .bind(free_monthly)
    .bind(user_id)
    .execute(pool)
    .await?;
    
    // Record cancellation
    sqlx::query(
        r#"
        INSERT INTO ai_subscription_history 
            (user_id, from_plan_id, to_plan_id, change_type, effective_at, created_by, status)
        VALUES ($1, $2, $3, 'cancel', NOW(), $4, 'completed')
        "#
    )
    .bind(user_id)
    .bind(wallet.plan_id)
    .bind(free_plan_id)
    .bind(actor_id)
    .execute(pool)
    .await?;
    
    Ok(())
}

/// Start a trial for a new user
/// 
/// Sets up a trial period with the specified plan.
pub async fn start_trial(
    pool: &PgPool,
    user_id: Uuid,
    trial_plan_id: Uuid,
    trial_days: i32,
) -> Result<DateTime<Utc>, AiCreditsError> {
    let trial_ends_at = Utc::now() + Duration::days(trial_days as i64);
    
    let plan: Option<(i32,)> = sqlx::query_as(
        "SELECT monthly_credits FROM ai_plans WHERE id = $1"
    )
    .bind(trial_plan_id)
    .fetch_optional(pool)
    .await?;
    
    let Some((monthly_credits,)) = plan else {
        return Err(AiCreditsError::UnknownFeature("Trial plan not found".to_string()));
    };
    
    sqlx::query(
        r#"
        UPDATE user_ai_subscriptions 
        SET plan_id = $1,
            monthly_credits_total = $2,
            is_trial = TRUE,
            trial_days = $3,
            trial_ends_at = $4,
            current_period_end = $4,
            updated_at = NOW()
        WHERE user_id = $5
        "#
    )
    .bind(trial_plan_id)
    .bind(monthly_credits)
    .bind(trial_days)
    .bind(trial_ends_at)
    .bind(user_id)
    .execute(pool)
    .await?;
    
    Ok(trial_ends_at)
}

/// Check and expire trials that have ended
/// 
/// Moves users whose trials have expired to the free plan.
/// Should be called periodically by a cron job.
pub async fn expire_trials(pool: &PgPool) -> Result<usize, AiCreditsError> {
    let now = Utc::now();
    
    // Get free plan
    let free_plan: Option<(Uuid, i32)> = sqlx::query_as(
        "SELECT id, monthly_credits FROM ai_plans WHERE code = 'free' AND is_active = TRUE"
    )
    .fetch_optional(pool)
    .await?;
    
    let Some((free_plan_id, free_monthly)) = free_plan else {
        return Err(AiCreditsError::UnknownFeature("Free plan not found".to_string()));
    };
    
    // Find expired trials
    let expired: Vec<(Uuid, Uuid)> = sqlx::query_as(
        r#"
        SELECT user_id, plan_id
        FROM user_ai_subscriptions
        WHERE is_trial = TRUE AND trial_ends_at <= $1
        "#
    )
    .bind(now)
    .fetch_all(pool)
    .await
    .map_err(AiCreditsError::Db)?;
    
    let mut expired_count = 0;
    
    for (user_id, old_plan_id) in expired {
        sqlx::query(
            r#"
            UPDATE user_ai_subscriptions 
            SET plan_id = $1,
                monthly_credits_total = $2,
                is_trial = FALSE,
                trial_days = NULL,
                trial_ends_at = NULL,
                current_period_end = NOW() + INTERVAL '30 days',
                updated_at = NOW()
            WHERE user_id = $3
            "#
        )
        .bind(free_plan_id)
        .bind(free_monthly)
        .bind(user_id)
        .execute(pool)
        .await?;
        
        // Record trial expiration
        sqlx::query(
            r#"
            INSERT INTO ai_subscription_history 
                (user_id, from_plan_id, to_plan_id, change_type, effective_at, status)
            VALUES ($1, $2, $3, 'trial_expired', NOW(), 'completed')
            "#
        )
        .bind(user_id)
        .bind(old_plan_id)
        .bind(free_plan_id)
        .execute(pool)
        .await?;
        
        expired_count += 1;
    }
    
    Ok(expired_count)
}

/// Get subscription history for a user
pub async fn get_subscription_history(
    pool: &PgPool,
    user_id: Uuid,
    limit: i64,
) -> Result<Vec<SubscriptionHistoryRow>, sqlx::Error> {
    sqlx::query_as(
        r#"
        SELECT 
            h.id,
            h.user_id,
            fp.name as from_plan_name,
            tp.name as to_plan_name,
            h.change_type,
            h.proration_credits,
            h.effective_at,
            h.status,
            h.created_at
        FROM ai_subscription_history h
        LEFT JOIN ai_plans fp ON h.from_plan_id = fp.id
        JOIN ai_plans tp ON h.to_plan_id = tp.id
        WHERE h.user_id = $1
        ORDER BY h.created_at DESC
        LIMIT $2
        "#
    )
    .bind(user_id)
    .bind(limit)
    .fetch_all(pool)
    .await
}

#[derive(Debug, sqlx::FromRow)]
pub struct SubscriptionHistoryRow {
    pub id: Uuid,
    pub user_id: Uuid,
    pub from_plan_name: Option<String>,
    pub to_plan_name: String,
    pub change_type: String,
    pub proration_credits: i32,
    pub effective_at: DateTime<Utc>,
    pub status: String,
    pub created_at: DateTime<Utc>,
}