managing-apps/src/Managing.Application/Workers/BalanceTrackingWorker.cs

using Managing.Application.Abstractions.Repositories;
using Managing.Application.Abstractions.Services;
using Managing.Application.ManageBot.Commands;
using Managing.Domain.Bots;
using Managing.Domain.Statistics;
using MediatR;
using Microsoft.Extensions.Logging;
using static Managing.Common.Enums;

namespace Managing.Application.Workers;

public class BalanceTrackingWorker : BaseWorker<BalanceTrackingWorker>
{
    private readonly IMediator _mediator;
    private readonly IAccountService _accountService;
    private readonly IAgentBalanceRepository _agentBalanceRepository;
    private bool _isInitialized;

    public BalanceTrackingWorker(
        ILogger<BalanceTrackingWorker> logger,
        IServiceProvider serviceProvider,
        IMediator mediator,
        IAccountService accountService,
        IAgentBalanceRepository agentBalanceRepository)
        : base(
            WorkerType.BalanceTracking,
            logger,
            TimeSpan.FromHours(1),
            serviceProvider)
    {
        _mediator = mediator;
        _accountService = accountService;
        _agentBalanceRepository = agentBalanceRepository;
        _isInitialized = false;
    }

    protected override async Task Run(CancellationToken cancellationToken)
    {
        if (!_isInitialized)
        {
            _logger.LogInformation("Waiting 5 minutes for bots to initialize before starting balance tracking...");
            await Task.Delay(TimeSpan.FromMinutes(3), cancellationToken);
            _isInitialized = true;
        }

        _logger.LogInformation("Starting balance tracking...");

        // Get all active bots
        var bots = await _mediator.Send(new GetBotsByStatusCommand(BotStatus.Running));

        var botCount = bots.Count();
        if (botCount == 0)
        {
            _logger.LogWarning("No active bots found. Skipping balance tracking.");
            return;
        }

        _logger.LogInformation($"Found {botCount} active bots. Proceeding with balance tracking.");
        await TrackBalances(bots);
        _logger.LogInformation("Completed balance tracking");
    }

    private async Task TrackBalances(IEnumerable<Bot> bots)
    {
        // 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;

                // Check if we need to update this agent's balance
                var lastBalance = (await _agentBalanceRepository.GetAgentBalances(
                    agentName,
                    DateTime.UtcNow.AddDays(-1),
                    DateTime.UtcNow)).OrderByDescending(b => b.Time).FirstOrDefault();

                if (lastBalance != null && DateTime.UtcNow.Subtract(lastBalance.Time).TotalHours < 24)
                {
                    _logger.LogInformation(
                        $"Skipping agent {agentName} - Last balance update was {lastBalance.Time:g} UTC");
                    continue;
                }

                decimal totalAgentValue = 0;
                decimal totalBotAllocatedBalance = 0;
                decimal totalAccountUsdValue = 0;
                decimal botsAllocationUsdValue = 0;
                decimal totalPnL = 0;

                _logger.LogInformation($"Processing agent: {agentName} with {agentBots.Count} bots");

                // Calculate total allocated balance for all bots
                foreach (var bot in agentBots)
                {
                    totalBotAllocatedBalance += bot.Volume;
                    _logger.LogInformation(
                        $"Bot {bot.Name} allocated balance: {bot.Volume} 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");
                        }

                        totalAccountUsdValue += 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)");
                        }
                    }
                }

                // Process all bots in a single iteration
                foreach (var bot in agentBots)
                {
                    totalPnL += bot.Pnl;
                }

                totalAgentValue = totalAccountUsdValue + botsAllocationUsdValue;

                _logger.LogInformation(
                    $"Agent {agentName} total aggregated value: {totalAgentValue} USD (Account: {totalAccountUsdValue} USD, Bot Wallet: {botsAllocationUsdValue} USD)");

                // Create and save the agent balance
                var agentBalance = new AgentBalance
                {
                    AgentName = agentName,
                    TotalValue = totalAgentValue,
                    TotalAccountUsdValue = totalAccountUsdValue,
                    BotsAllocationUsdValue = botsAllocationUsdValue,
                    PnL = totalPnL,
                    Time = DateTime.UtcNow
                };

                _agentBalanceRepository.InsertAgentBalance(agentBalance);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, $"Error processing agent {agentEntry.Key}");
            }
        }
    }
}