98 lines
3.0 KiB
C#
98 lines
3.0 KiB
C#
using System.Data;
|
|
using MiaoJiZhang.Domain.Entities;
|
|
using MiaoJiZhang.Domain.Enums;
|
|
using MiaoJiZhang.Infrastructure.Persistence;
|
|
using Microsoft.EntityFrameworkCore;
|
|
|
|
namespace MiaoJiZhang.Api.Services;
|
|
|
|
public sealed record AiChatQuotaStatus(
|
|
int Limit,
|
|
int Used,
|
|
int Remaining,
|
|
AiQuotaPeriod Period,
|
|
DateTime WindowStartedAt,
|
|
DateTime ResetAt,
|
|
bool Allowed);
|
|
|
|
public sealed class AiChatQuotaService(AppDbContext db)
|
|
{
|
|
public async Task<AiChatQuotaStatus> TryConsumeAsync(
|
|
long userId,
|
|
CancellationToken ct = default)
|
|
{
|
|
await using var transaction = await db.Database.BeginTransactionAsync(
|
|
IsolationLevel.Serializable,
|
|
ct);
|
|
var user = await db.Users.SingleAsync(item => item.Id == userId, ct);
|
|
var status = GetStatus(user);
|
|
if (!status.Allowed)
|
|
{
|
|
await transaction.CommitAsync(ct);
|
|
return status;
|
|
}
|
|
|
|
if (user.AiChatWindowStartedAt != status.WindowStartedAt)
|
|
{
|
|
user.AiChatWindowStartedAt = status.WindowStartedAt;
|
|
user.AiChatUsed = 0;
|
|
}
|
|
user.AiChatUsed++;
|
|
await db.SaveChangesAsync(ct);
|
|
await transaction.CommitAsync(ct);
|
|
return GetStatus(user) with { Allowed = true };
|
|
}
|
|
|
|
public static AiChatQuotaStatus GetStatus(User user)
|
|
{
|
|
var (start, end) = CurrentWindow(user.AiChatPeriod);
|
|
var used = user.AiChatWindowStartedAt == start
|
|
? Math.Max(0, user.AiChatUsed)
|
|
: 0;
|
|
var limit = Math.Max(0, user.AiChatLimit);
|
|
var remaining = limit == 0 ? -1 : Math.Max(0, limit - used);
|
|
return new AiChatQuotaStatus(
|
|
limit,
|
|
used,
|
|
remaining,
|
|
user.AiChatPeriod,
|
|
start,
|
|
end,
|
|
limit == 0 || used < limit);
|
|
}
|
|
|
|
public static (DateTime Start, DateTime End) CurrentWindow(
|
|
AiQuotaPeriod period)
|
|
{
|
|
var now = ChinaClock.Now;
|
|
var start = period switch
|
|
{
|
|
AiQuotaPeriod.Week => now.Date.AddDays(-(now.DayOfWeek == DayOfWeek.Sunday
|
|
? 6
|
|
: (int)now.DayOfWeek - 1)),
|
|
AiQuotaPeriod.Month => new DateTime(now.Year, now.Month, 1),
|
|
_ => now.Date,
|
|
};
|
|
var end = period switch
|
|
{
|
|
AiQuotaPeriod.Week => start.AddDays(7),
|
|
AiQuotaPeriod.Month => start.AddMonths(1),
|
|
_ => start.AddDays(1),
|
|
};
|
|
return (ChinaClock.ToUtc(start), ChinaClock.ToUtc(end));
|
|
}
|
|
|
|
public static bool TryParsePeriod(string? value, out AiQuotaPeriod period)
|
|
{
|
|
period = value?.Trim().ToLowerInvariant() switch
|
|
{
|
|
"week" => AiQuotaPeriod.Week,
|
|
"month" => AiQuotaPeriod.Month,
|
|
_ => AiQuotaPeriod.Day,
|
|
};
|
|
return value?.Trim().ToLowerInvariant() is "day" or "week" or "month";
|
|
}
|
|
|
|
public static string PeriodKey(AiQuotaPeriod period) =>
|
|
period.ToString().ToLowerInvariant();
|
|
} |