managing-apps/src/Managing.Application/Bots/Grains/AgentGrain.cs

using Managing.Application.Abstractions;
using Managing.Application.Abstractions.Grains;
using Managing.Application.Abstractions.Models;
using Managing.Application.Abstractions.Services;
using Managing.Application.Bots.Models;
using Managing.Application.Orleans;
using Managing.Domain.Statistics;
using Microsoft.Extensions.Logging;
using static Managing.Common.Enums;

namespace Managing.Application.Bots.Grains;

/// <summary>
/// Orleans grain for Agent operations.
/// Uses custom trading placement with load balancing and built-in fallback.
/// </summary>
[TradingPlacement] // Use custom trading placement with load balancing
public class AgentGrain : Grain, IAgentGrain
{
    private readonly IPersistentState<AgentGrainState> _state;
    private readonly ILogger<AgentGrain> _logger;
    private readonly IBotService _botService;
    private readonly IAgentService _agentService;
    private readonly IExchangeService _exchangeService;
    private readonly IUserService _userService;
    private readonly IAccountService _accountService;
    private readonly ITradingService _tradingService;

    public AgentGrain(
        [PersistentState("agent-state", "agent-store")]
        IPersistentState<AgentGrainState> state,
        ILogger<AgentGrain> logger,
        IBotService botService,
        IAgentService agentService,
        IExchangeService exchangeService,
        IUserService userService,
        IAccountService accountService,
        ITradingService tradingService)
    {
        _state = state;
        _logger = logger;
        _botService = botService;
        _agentService = agentService;
        _exchangeService = exchangeService;
        _userService = userService;
        _accountService = accountService;
        _tradingService = tradingService;
    }

    public override Task OnActivateAsync(CancellationToken cancellationToken)
    {
        _logger.LogInformation("AgentGrain activated for user {UserId}", this.GetPrimaryKeyLong());
        return base.OnActivateAsync(cancellationToken);
    }

    public async Task InitializeAsync(int userId, string agentName)
    {
        _state.State.AgentName = agentName;
        await _state.WriteStateAsync();
        
        // Create an empty AgentSummary for the new agent
        var emptySummary = new AgentSummary
        {
            UserId = userId,
            AgentName = agentName,
            TotalPnL = 0,
            TotalROI = 0,
            Wins = 0,
            Losses = 0,
            Runtime = null,
            CreatedAt = DateTime.UtcNow,
            UpdatedAt = DateTime.UtcNow,
            ActiveStrategiesCount = 0,
            TotalVolume = 0,
            TotalBalance = 0
        };
        
        await _agentService.SaveOrUpdateAgentSummary(emptySummary);
        _logger.LogInformation("Agent {UserId} initialized with name {AgentName} and empty summary", userId, agentName);
    }

    public async Task UpdateAgentNameAsync(string agentName)
    {
        _state.State.AgentName = agentName;
        await _state.WriteStateAsync();
        _logger.LogInformation("Agent {UserId} updated with name {AgentName}", this.GetPrimaryKeyLong(), agentName);
    }

    public async Task OnAgentSummaryUpdateAsync(AgentSummaryUpdateEvent updateEvent)
    {
        try
        {
            _logger.LogInformation("Received agent summary update event for user {UserId}, event type: {EventType}",
                this.GetPrimaryKeyLong(), updateEvent.EventType);
            
            // Only update summary if the event is for this agent's bots
            if (_state.State.BotIds.Contains(updateEvent.BotId))
            {
                await UpdateSummary();
            }
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error processing agent summary update event for user {UserId}",
                this.GetPrimaryKeyLong());
        }
    }

    public async Task UpdateSummary()
    {
        try
        {
            // Get all bots for this agent
            var bots = await _botService.GetBotsByIdsAsync(_state.State.BotIds);

            // Calculate aggregated statistics from bot data
            var totalPnL = bots.Sum(b => b.Pnl);
            var totalWins = bots.Sum(b => b.TradeWins);
            var totalLosses = bots.Sum(b => b.TradeLosses);

            // Calculate ROI based on total volume traded with proper division by zero handling
            var totalVolume = bots.Sum(b => b.Volume);
            decimal totalROI;

            if (totalVolume > 0)
            {
                totalROI = (totalPnL / totalVolume) * 100;
            }
            else if (totalVolume == 0 && totalPnL == 0)
            {
                // No trading activity yet
                totalROI = 0;
            }
            else if (totalVolume == 0 && totalPnL != 0)
            {
                // Edge case: PnL exists but no volume (shouldn't happen in normal cases)
                _logger.LogWarning("Agent {UserId} has PnL {PnL} but zero volume", this.GetPrimaryKeyLong(), totalPnL);
                totalROI = 0;
            }
            else
            {
                // Fallback for any other edge cases
                totalROI = 0;
            }

            // Calculate Runtime based on the farthest date from bot startup times
            DateTime? runtime = null;
            if (bots.Any())
            {
                runtime = bots.Max(b => b.StartupTime);
            }

            // Calculate total balance (USDC + open positions value)
            decimal totalBalance = 0;
            try
            {
                var userId = (int)this.GetPrimaryKeyLong();
                var user = await _userService.GetUserByIdAsync(userId);

                if (user != null)
                {
                    var userAccounts = await _accountService.GetAccountsByUserAsync(user, hideSecrets: true, true);

                    foreach (var account in userAccounts)
                    {
                        // Get USDC balance
                        var usdcBalances = await _exchangeService.GetBalances(account);
                        var usdcBalance = usdcBalances.FirstOrDefault(b => b.TokenName?.ToUpper() == "USDC")?.Amount ??
                                          0;
                        totalBalance += usdcBalance;
                    }

                    // Get positions for all bots using their GUIDs as InitiatorIdentifier
                    var botPositions =
                        await _tradingService.GetPositionsByInitiatorIdentifiersAsync(_state.State.BotIds);

                    foreach (var position in botPositions.Where(p => !p.IsFinished()))
                    {
                        totalBalance += position.Open.Price * position.Open.Quantity;
                        totalBalance += position.ProfitAndLoss?.Realized ?? 0;
                    }
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error calculating total balance for agent {UserId}", this.GetPrimaryKeyLong());
                totalBalance = 0; // Set to 0 if calculation fails
            }

            var summary = new AgentSummary
            {
                UserId = (int)this.GetPrimaryKeyLong(),
                AgentName = _state.State.AgentName,
                TotalPnL = totalPnL,
                Wins = totalWins,
                Losses = totalLosses,
                TotalROI = totalROI,
                Runtime = runtime,
                ActiveStrategiesCount = bots.Count(b => b.Status == BotStatus.Running),
                TotalVolume = totalVolume,
                TotalBalance = totalBalance,
            };

            // Save summary to database
            await _agentService.SaveOrUpdateAgentSummary(summary);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error calculating agent summary for user {UserId}", this.GetPrimaryKeyLong());
        }
    }

    public async Task RegisterBotAsync(Guid botId)
    {
        if (_state.State.BotIds.Add(botId))
        {
            await _state.WriteStateAsync();
            _logger.LogInformation("Bot {BotId} registered to Agent {UserId}", botId, this.GetPrimaryKeyLong());
            
            // Update summary after registering bot
            await UpdateSummary();
        }
    }

    public async Task UnregisterBotAsync(Guid botId)
    {
        if (_state.State.BotIds.Remove(botId))
        {
            await _state.WriteStateAsync();
            _logger.LogInformation("Bot {BotId} unregistered from Agent {UserId}", botId, this.GetPrimaryKeyLong());
            
            // Update summary after unregistering bot
            await UpdateSummary();
        }
    }
}