MycLib/GeminiTest/Myc.System.Quota.pas

unit Myc.System.Quota;

interface

uses
    System.SysUtils,
    System.Win.Registry,
    Winapi.Windows,
    System.Classes;

type
    // Handles global quota tracking via Windows Registry
    TGlobalQuota = class
    private
        const
            REG_KEY = 'Software\Myc\GeminiTools\Quota';

        // Value names
        const
            VAL_TOTAL_TOKENS = 'TotalTokens';
        const
            VAL_DAILY_TOKENS = 'DailyTokens';
        const
            VAL_LAST_RESET = 'LastReset'; // Manual reset timestamp
        const
            VAL_LAST_UPDATE = 'LastUpdate'; // Last modification timestamp
        const
            VAL_DAILY_DATE = 'DailyDate'; // Date for the daily counter reference

    public
        // Adds tokens to global and daily counters thread-safely
        // Automatically resets the daily counter if the day has changed
        class procedure AddTokens(Count: Integer);

        // Reads the all-time total from registry
        class function GetTotal: Int64;

        // Reads the accumulated tokens for the current day
        class function GetDailyTotal: Int64;

        // Returns the date and time of the last modification (AddTokens)
        class function GetLastUpdate: TDateTime;

        // Returns the date and time of the last manual reset
        class function GetLastReset: TDateTime;

        // Resets all counters (Total and Daily) to zero
        class procedure Reset;
    end;

implementation

{ TGlobalQuota }

class procedure TGlobalQuota.AddTokens(Count: Integer);
var
    reg: TRegistry;
    mutex: THandle;
    totalTokens, dailyTokens: Int64;
    strValue: string;
    lastDailyDate, today: TDateTime;
begin
    if Count <= 0 then
        exit;

    // Use a named mutex to prevent race conditions between different processes
    mutex := CreateMutex(nil, False, 'GlobalGeminiQuotaMutex');
    if WaitForSingleObject(mutex, 2000) = WAIT_OBJECT_0 then
        try
            reg := TRegistry.Create(KEY_READ or KEY_WRITE);
            try
                reg.RootKey := HKEY_CURRENT_USER;
                if reg.OpenKey(REG_KEY, True) then
                begin
                    // 1. Handle Total Tokens
                    totalTokens := 0;
                    if reg.ValueExists(VAL_TOTAL_TOKENS) then
                    begin
                        strValue := reg.ReadString(VAL_TOTAL_TOKENS);
                        TryStrToInt64(strValue, totalTokens);
                    end;
                    totalTokens := totalTokens + Count;

                    // 2. Handle Daily Tokens
                    dailyTokens := 0;
                    lastDailyDate := 0;
                    today := Trunc(Now); // Date part only

                    if reg.ValueExists(VAL_DAILY_DATE) then
                        lastDailyDate := reg.ReadDate(VAL_DAILY_DATE);

                    // Check if we are on a new day
                    if lastDailyDate = today then
                    begin
                        // Same day, read existing daily count
                        if reg.ValueExists(VAL_DAILY_TOKENS) then
                        begin
                            strValue := reg.ReadString(VAL_DAILY_TOKENS);
                            TryStrToInt64(strValue, dailyTokens);
                        end;
                    end
                    else
                    begin
                        // New day (or first run), dailyTokens remains 0
                        // Update the reference date
                        reg.WriteDate(VAL_DAILY_DATE, today);
                    end;

                    dailyTokens := dailyTokens + Count;

                    // 3. Write updates
                    // Store numbers as String to support Int64 (Registry WriteInteger is 32-bit in older API wrappers)
                    reg.WriteString(VAL_TOTAL_TOKENS, totalTokens.ToString);
                    reg.WriteString(VAL_DAILY_TOKENS, dailyTokens.ToString);
                    reg.WriteDateTime(VAL_LAST_UPDATE, Now);

                    reg.CloseKey;
                end;
            finally
                reg.Free;
            end;
        finally
            ReleaseMutex(mutex);
            CloseHandle(mutex);
        end;
end;

class function TGlobalQuota.GetTotal: Int64;
var
    reg: TRegistry;
    strValue: string;
begin
    Result := 0;
    reg := TRegistry.Create(KEY_READ);
    try
        reg.RootKey := HKEY_CURRENT_USER;
        if reg.OpenKey(REG_KEY, False) then
        begin
            if reg.ValueExists(VAL_TOTAL_TOKENS) then
            begin
                strValue := reg.ReadString(VAL_TOTAL_TOKENS);
                TryStrToInt64(strValue, Result);
            end;
            reg.CloseKey;
        end;
    finally
        reg.Free;
    end;
end;

class function TGlobalQuota.GetDailyTotal: Int64;
var
    reg: TRegistry;
    strValue: string;
    lastDailyDate: TDateTime;
begin
    Result := 0;
    reg := TRegistry.Create(KEY_READ);
    try
        reg.RootKey := HKEY_CURRENT_USER;
        if reg.OpenKey(REG_KEY, False) then
        begin
            // Check if the stored daily count belongs to today
            if reg.ValueExists(VAL_DAILY_DATE) then
            begin
                lastDailyDate := reg.ReadDate(VAL_DAILY_DATE);
                if lastDailyDate = Trunc(Now) then
                begin
                    if reg.ValueExists(VAL_DAILY_TOKENS) then
                    begin
                        strValue := reg.ReadString(VAL_DAILY_TOKENS);
                        TryStrToInt64(strValue, Result);
                    end;
                end;
                // If date differs, Result remains 0 (logic effectively resets on read, though write happens in AddTokens)
            end;
            reg.CloseKey;
        end;
    finally
        reg.Free;
    end;
end;

class function TGlobalQuota.GetLastUpdate: TDateTime;
var
    reg: TRegistry;
begin
    Result := 0;
    reg := TRegistry.Create(KEY_READ);
    try
        reg.RootKey := HKEY_CURRENT_USER;
        if reg.OpenKey(REG_KEY, False) then
        begin
            if reg.ValueExists(VAL_LAST_UPDATE) then
                Result := reg.ReadDateTime(VAL_LAST_UPDATE);
            reg.CloseKey;
        end;
    finally
        reg.Free;
    end;
end;

class function TGlobalQuota.GetLastReset: TDateTime;
var
    reg: TRegistry;
begin
    Result := 0;
    reg := TRegistry.Create(KEY_READ);
    try
        reg.RootKey := HKEY_CURRENT_USER;
        if reg.OpenKey(REG_KEY, False) then
        begin
            if reg.ValueExists(VAL_LAST_RESET) then
                Result := reg.ReadDateTime(VAL_LAST_RESET);
            reg.CloseKey;
        end;
    finally
        reg.Free;
    end;
end;

class procedure TGlobalQuota.Reset;
var
    reg: TRegistry;
    mutex: THandle;
begin
    // Reset also needs Mutex protection to avoid conflicting writes with AddTokens
    mutex := CreateMutex(nil, False, 'GlobalGeminiQuotaMutex');
    if WaitForSingleObject(mutex, 2000) = WAIT_OBJECT_0 then
        try
            reg := TRegistry.Create(KEY_WRITE);
            try
                reg.RootKey := HKEY_CURRENT_USER;
                if reg.OpenKey(REG_KEY, True) then
                begin
                    reg.WriteString(VAL_TOTAL_TOKENS, '0');
                    reg.WriteString(VAL_DAILY_TOKENS, '0');
                    reg.WriteDate(VAL_DAILY_DATE, Trunc(Now));
                    reg.WriteDateTime(VAL_LAST_RESET, Now);
                    // We typically update LastUpdate on reset as well, or leave it as the last "Add" action.
                    // Here we update it to indicate change.
                    reg.WriteDateTime(VAL_LAST_UPDATE, Now);
                    reg.CloseKey;
                end;
            finally
                reg.Free;
            end;
        finally
            ReleaseMutex(mutex);
            CloseHandle(mutex);
        end;
end;

end.