From b34e3aa8861c852ab2707c88d2e80234dc11a307 Mon Sep 17 00:00:00 2001
From: cryptooda <crypto.oda@gmail.com>
Date: Fri, 16 May 2025 00:27:07 +0700
Subject: [PATCH] Add balance tracking worker

---
 src/Managing.Api/Program.cs                   |   3 +
 .../Workers/BalanceTrackingWorker.cs          | 135 ++++++++++++++++++
 2 files changed, 138 insertions(+)
 create mode 100644 src/Managing.Application/Workers/BalanceTrackingWorker.cs

diff --git a/src/Managing.Api/Program.cs b/src/Managing.Api/Program.cs
index c23572c..06df818 100644
--- a/src/Managing.Api/Program.cs
+++ b/src/Managing.Api/Program.cs
@@ -6,6 +6,7 @@ using Managing.Api.Filters;
 using Managing.Api.HealthChecks;
 using Managing.Api.Workers;
 using Managing.Application.Hubs;
+using Managing.Application.Workers;
 using Managing.Bootstrap;
 using Managing.Common;
 using Managing.Core.Middleawares;
@@ -203,6 +204,8 @@ if (builder.Configuration.GetValue<bool>("EnableBotManager", false))
     builder.Services.AddHostedService<BotManagerWorker>();
 }
 
+builder.Services.AddHostedService<BalanceTrackingWorker>();
+
 // App
 var app = builder.Build();
 app.UseSerilogRequestLogging();
diff --git a/src/Managing.Application/Workers/BalanceTrackingWorker.cs b/src/Managing.Application/Workers/BalanceTrackingWorker.cs
new file mode 100644
index 0000000..23771f8
--- /dev/null
+++ b/src/Managing.Application/Workers/BalanceTrackingWorker.cs
@@ -0,0 +1,135 @@
+using Managing.Application.Abstractions.Services;
+using Managing.Application.ManageBot.Commands;
+using MediatR;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
+
+namespace Managing.Application.Workers;
+
+public class BalanceTrackingWorker : BackgroundService
+{
+    private readonly ILogger<BalanceTrackingWorker> _logger;
+    private readonly IMediator _mediator;
+    private readonly IAccountService _accountService;
+    private readonly TimeSpan _interval = TimeSpan.FromMinutes(1);
+
+    public BalanceTrackingWorker(
+        ILogger<BalanceTrackingWorker> logger,
+        IMediator mediator,
+        IAccountService accountService)
+    {
+        _logger = logger;
+        _mediator = mediator;
+        _accountService = accountService;
+    }
+
+    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+    {
+        while (!stoppingToken.IsCancellationRequested)
+        {
+            try
+            {
+                _logger.LogInformation("Starting balance tracking...");
+                await TrackBalances();
+                _logger.LogInformation("Completed balance tracking");
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "Error occurred while tracking balances");
+            }
+
+            await Task.Delay(_interval, stoppingToken);
+        }
+    }
+
+    private async Task TrackBalances()
+    {
+        // Get all active bots
+        var bots = await _mediator.Send(new GetActiveBotsCommand());
+        
+        // Group bots by agent/user
+        var botsByAgent = bots
+            .Where(b => b.User != null)
+            .GroupBy(b => b.User.AgentName)
+            .ToDictionary(g => g.Key, g => g.ToList());
+
+        foreach (var agentEntry in botsByAgent)
+        {
+            try
+            {
+                var agentName = agentEntry.Key;
+                var agentBots = agentEntry.Value;
+                decimal totalAgentValue = 0;
+                decimal totalBotAllocatedBalance = 0;
+
+                _logger.LogInformation($"Processing agent: {agentName} with {agentBots.Count} bots");
+
+                // Calculate total allocated balance for all bots
+                foreach (var bot in agentBots)
+                {
+                    totalBotAllocatedBalance += bot.Config.BotTradingBalance;
+                    _logger.LogInformation(
+                        $"Bot {bot.Name} allocated balance: {bot.Config.BotTradingBalance} USD");
+                }
+
+                // Get account balances for this agent (only once per agent)
+                var agent = agentBots.First().User; // Get the user object from the first bot
+                var accountBalances = _accountService.GetAccountsBalancesByUser(agent, true);
+                foreach (var accountBalance in accountBalances)
+                {
+                    if (accountBalance.Balances != null)
+                    {
+                        var accountTotalValue = accountBalance.Balances.Sum(b => b.Value);
+                        
+                        // If this is the account that holds the bot balances (USDC), subtract the allocated amounts
+                        var usdcBalance = accountBalance.Balances.FirstOrDefault(b => b.TokenName == "USDC");
+                        if (usdcBalance != null)
+                        {
+                            _logger.LogInformation(
+                                $"Account {accountBalance.Name} USDC balance before bot allocation: {usdcBalance.Value} USD");
+                            usdcBalance.Value -= totalBotAllocatedBalance;
+                            _logger.LogInformation(
+                                $"Account {accountBalance.Name} USDC balance after bot allocation: {usdcBalance.Value} USD");
+                        }
+
+                        totalAgentValue += accountTotalValue;
+
+                        _logger.LogInformation(
+                            $"Account {accountBalance.Name} total value: {accountTotalValue} USD");
+                        
+                        // Log individual token balances for debugging
+                        foreach (var balance in accountBalance.Balances)
+                        {
+                            _logger.LogInformation(
+                                $"  - {balance.TokenName}: {balance.Amount} (Value: {balance.Value} USD)");
+                        }
+                    }
+                }
+
+                // Add up all bot wallet balances for this agent
+                foreach (var bot in agentBots)
+                {
+                    var latestBotBalance = bot.WalletBalances
+                        .OrderByDescending(x => x.Key)
+                        .FirstOrDefault();
+
+                    if (latestBotBalance.Key != default)
+                    {
+                        totalAgentValue += latestBotBalance.Value;
+                        _logger.LogInformation(
+                            $"Bot {bot.Name} wallet balance: {latestBotBalance.Value} USD at {latestBotBalance.Key}");
+                    }
+                }
+
+                _logger.LogInformation(
+                    $"Agent {agentName} total aggregated value: {totalAgentValue} USD");
+                
+                // TODO: Save aggregated agent balance to database
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, $"Error processing agent {agentEntry.Key}");
+            }
+        }
+    }
+}
\ No newline at end of file