sendico/api/edge/bff/internal/server/walletapiimp/balance.go

package walletapiimp

import (
	"context"
	"crypto/tls"
	"net/http"
	"strings"
	"sync"

	"github.com/tech/sendico/pkg/api/http/response"
	"github.com/tech/sendico/pkg/discovery"
	"github.com/tech/sendico/pkg/merrors"
	"github.com/tech/sendico/pkg/model"
	"github.com/tech/sendico/pkg/mservice"
	connectorv1 "github.com/tech/sendico/pkg/proto/connector/v1"
	"github.com/tech/sendico/server/interface/api/sresponse"
	mutil "github.com/tech/sendico/server/internal/mutil/param"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
)

func (a *WalletAPI) getWalletBalance(r *http.Request, account *model.Account, token *sresponse.TokenData) http.HandlerFunc {
	orgRef, err := a.oph.GetRef(r)
	if err != nil {
		a.logger.Warn("Failed to parse organization reference for wallet balance", zap.Error(err), zap.String(a.oph.Name(), a.oph.GetID(r)))
		return response.BadReference(a.logger, a.Name(), a.oph.Name(), a.oph.GetID(r), err)
	}
	walletRef := strings.TrimSpace(a.wph.GetID(r))
	if walletRef == "" {
		return response.BadReference(a.logger, a.Name(), a.wph.Name(), a.wph.GetID(r), merrors.InvalidArgument("wallet reference is required"))
	}

	ctx := r.Context()
	res, err := a.enf.Enforce(ctx, a.balancesPermissionRef, account.ID, orgRef, bson.NilObjectID, model.ActionRead)
	if err != nil {
		a.logger.Warn("Failed to check wallet balance permissions", zap.Error(err), mutil.PLog(a.oph, r), zap.String("wallet_ref", walletRef))
		return response.Auto(a.logger, a.Name(), err)
	}
	if !res {
		a.logger.Debug("Access denied when reading wallet balance", mutil.PLog(a.oph, r), zap.String("wallet_ref", walletRef))
		return response.AccessDenied(a.logger, a.Name(), "wallet balance read permission denied")
	}

	if a.discovery == nil {
		return response.Internal(a.logger, mservice.ChainGateway, merrors.Internal("discovery client is not configured"))
	}

	// Discover CRYPTO rail gateways
	lookupCtx, cancel := context.WithTimeout(ctx, discoveryLookupTimeout)
	defer cancel()

	lookupResp, err := a.discovery.Lookup(lookupCtx)
	if err != nil {
		a.logger.Warn("Failed to lookup discovery registry", zap.Error(err))
		return response.Auto(a.logger, a.Name(), err)
	}

	// Filter gateways by CRYPTO rail
	cryptoGateways := filterCryptoGateways(lookupResp.Gateways)
	if len(cryptoGateways) == 0 {
		a.logger.Debug("No CRYPTO rail gateways found in discovery")
		return response.Auto(a.logger, a.Name(), merrors.NoData("no crypto gateways available"))
	}
	a.logger.Debug("Resolved CRYPTO gateways for wallet balance lookup",
		zap.String("organization_ref", orgRef.Hex()),
		zap.String("wallet_ref", walletRef),
		zap.Int("gateway_count", len(cryptoGateways)))

	route, routeErr := a.walletRoute(ctx, orgRef.Hex(), walletRef)
	if routeErr != nil {
		a.logger.Warn("Failed to resolve wallet route", zap.Error(routeErr), zap.String("wallet_ref", walletRef), zap.String("organization_ref", orgRef.Hex()))
	}
	if route != nil {
		a.logger.Debug("Resolved stored wallet route",
			zap.String("organization_ref", orgRef.Hex()),
			zap.String("wallet_ref", walletRef),
			zap.String("route_network", route.Network),
			zap.String("route_gateway_id", route.GatewayID))
		preferred := findGatewayForRoute(cryptoGateways, route)
		if preferred != nil {
			a.logger.Debug("Using preferred gateway from stored wallet route",
				zap.String("organization_ref", orgRef.Hex()),
				zap.String("wallet_ref", walletRef),
				zap.String("gateway_id", preferred.ID),
				zap.String("network", preferred.Network),
				zap.String("invoke_uri", preferred.InvokeURI))
			bal, preferredErr := a.queryGatewayBalance(ctx, *preferred, walletRef)
			if preferredErr == nil && bal != nil {
				a.rememberWalletRoute(ctx, orgRef.Hex(), walletRef, preferred.Network, preferred.ID)
				a.logger.Debug("Wallet balance resolved via preferred gateway",
					zap.String("organization_ref", orgRef.Hex()),
					zap.String("wallet_ref", walletRef),
					zap.String("gateway_id", preferred.ID),
					zap.String("network", preferred.Network))
				return sresponse.WalletBalanceFromConnector(a.logger, bal, token)
			}

			if preferredErr != nil {
				a.logger.Debug("Preferred gateway balance lookup failed, falling back to fan-out",
					zap.String("wallet_ref", walletRef),
					zap.String("network", route.Network),
					zap.String("gateway_id", preferred.ID),
					zap.String("invoke_uri", preferred.InvokeURI),
					zap.Error(preferredErr))
			} else {
				a.logger.Debug("Preferred gateway returned empty balance, falling back to fan-out",
					zap.String("wallet_ref", walletRef),
					zap.String("network", route.Network),
					zap.String("gateway_id", preferred.ID),
					zap.String("invoke_uri", preferred.InvokeURI))
			}

			cryptoGateways = dropGatewayByInvokeURI(cryptoGateways, preferred.InvokeURI)
			if len(cryptoGateways) == 0 {
				if preferredErr != nil {
					a.logger.Warn("Failed to fetch wallet balance from preferred gateway", zap.Error(preferredErr), zap.String("wallet_ref", walletRef))
					return response.Auto(a.logger, a.Name(), preferredErr)
				}
				a.logger.Warn("Wallet balance not found on preferred gateway", zap.String("wallet_ref", walletRef))
				return response.Auto(a.logger, a.Name(), merrors.NoData("wallet not found"))
			}
		} else {
			a.logger.Warn("Stored wallet route did not match any healthy discovery gateway",
				zap.String("organization_ref", orgRef.Hex()),
				zap.String("wallet_ref", walletRef),
				zap.String("route_network", route.Network),
				zap.String("route_gateway_id", route.GatewayID))
		}
	} else {
		a.logger.Debug("Stored wallet route not found; using gateway fallback",
			zap.String("organization_ref", orgRef.Hex()),
			zap.String("wallet_ref", walletRef))
	}

	// Fall back to querying remaining gateways in parallel.
	a.logger.Debug("Starting fallback wallet balance fan-out",
		zap.String("organization_ref", orgRef.Hex()),
		zap.String("wallet_ref", walletRef),
		zap.Int("gateway_count", len(cryptoGateways)))
	bal, err := a.queryBalanceFromGateways(ctx, cryptoGateways, orgRef.Hex(), walletRef)
	if err != nil {
		a.logger.Warn("Failed to fetch wallet balance from gateways", zap.Error(err), zap.String("wallet_ref", walletRef))
		return response.Auto(a.logger, a.Name(), err)
	}

	if bal == nil {
		a.logger.Warn("Wallet balance not found on any gateway", zap.String("wallet_ref", walletRef))
		return response.Auto(a.logger, a.Name(), merrors.NoData("wallet not found"))
	}

	return sresponse.WalletBalanceFromConnector(a.logger, bal, token)
}

func (a *WalletAPI) queryBalanceFromGateways(ctx context.Context, gateways []discovery.GatewaySummary, organizationRef string, walletRef string) (*connectorv1.Balance, error) {
	var mu sync.Mutex
	var wg sync.WaitGroup
	var result *connectorv1.Balance
	var lastErr error
	selectedGatewayID := ""
	selectedNetwork := ""

	a.logger.Debug("Querying wallet balance across gateways",
		zap.String("organization_ref", organizationRef),
		zap.String("wallet_ref", walletRef),
		zap.Int("gateway_count", len(gateways)))

	for _, gw := range gateways {
		wg.Add(1)
		go func(gateway discovery.GatewaySummary) {
			defer wg.Done()

			bal, err := a.queryGatewayBalance(ctx, gateway, walletRef)
			if err != nil {
				a.logger.Debug("Failed to query gateway for balance",
					zap.String("gateway_id", gateway.ID),
					zap.String("invoke_uri", gateway.InvokeURI),
					zap.String("wallet_ref", walletRef),
					zap.Error(err))
				mu.Lock()
				lastErr = err
				mu.Unlock()
				return
			}

			if bal != nil {
				mu.Lock()
				if result == nil {
					result = bal
					a.rememberWalletRoute(ctx, organizationRef, walletRef, gateway.Network, gateway.ID)
					selectedGatewayID = gateway.ID
					selectedNetwork = gateway.Network
					a.logger.Debug("Found wallet balance on gateway",
						zap.String("gateway_id", gateway.ID),
						zap.String("network", gateway.Network),
						zap.String("wallet_ref", walletRef))
				}
				mu.Unlock()
			}
		}(gw)
	}

	wg.Wait()

	if result != nil {
		a.logger.Debug("Wallet balance fan-out completed with result",
			zap.String("organization_ref", organizationRef),
			zap.String("wallet_ref", walletRef),
			zap.String("gateway_id", selectedGatewayID),
			zap.String("network", selectedNetwork))
		return result, nil
	}
	if lastErr != nil {
		a.logger.Debug("Wallet balance fan-out completed with errors",
			zap.String("organization_ref", organizationRef),
			zap.String("wallet_ref", walletRef),
			zap.Error(lastErr))
		return nil, lastErr
	}
	a.logger.Debug("Wallet balance fan-out completed without result",
		zap.String("organization_ref", organizationRef),
		zap.String("wallet_ref", walletRef))
	return nil, nil
}

func (a *WalletAPI) queryGatewayBalance(ctx context.Context, gateway discovery.GatewaySummary, walletRef string) (*connectorv1.Balance, error) {
	// Create connection with timeout
	dialCtx, cancel := context.WithTimeout(ctx, a.dialTimeout)
	defer cancel()

	var dialOpts []grpc.DialOption
	if a.insecure {
		dialOpts = append(dialOpts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	} else {
		dialOpts = append(dialOpts, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{})))
	}

	conn, err := grpc.DialContext(dialCtx, gateway.InvokeURI, dialOpts...)
	if err != nil {
		return nil, merrors.InternalWrap(err, "dial gateway")
	}
	defer conn.Close()

	client := connectorv1.NewConnectorServiceClient(conn)

	// Call with timeout
	callCtx, callCancel := context.WithTimeout(ctx, a.callTimeout)
	defer callCancel()

	req := &connectorv1.GetBalanceRequest{
		AccountRef: &connectorv1.AccountRef{
			AccountId: walletRef,
		},
	}

	resp, err := client.GetBalance(callCtx, req)
	if err != nil {
		return nil, err
	}

	return resp.GetBalance(), nil
}