bc-netts-energy/internal/service/energy_service.go

package service

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"time"

	appconfig "github.com/D36u99er/bc-netts-energy/internal/config"
	"github.com/D36u99er/bc-netts-energy/internal/netts"

	"log/slog"
)

var tronAddressPattern = regexp.MustCompile(`^T[1-9A-HJ-NP-Za-km-z]{33}$`)

// EnsureEnergyRequest describes a request to guarantee sufficient energy.
type EnsureEnergyRequest struct {
	FromAddress string
	ToAddress   string
	Amount      string
	CallbackURL string
}

// EnsureEnergyResponse summarises the orchestration result.
type EnsureEnergyResponse struct {
	FromAddress       string                `json:"from_address"`
	ToAddress         string                `json:"to_address"`
	Amount            string                `json:"amount"`
	AddressAdded      bool                  `json:"address_added"`
	RecommendedEnergy int                   `json:"recommended_energy"`
	EnergyNeeded      int                   `json:"energy_needed"`
	CyclesBefore      int                   `json:"cycles_before"`
	CyclesAfter       int                   `json:"cycles_after"`
	CyclesPurchased   int                   `json:"cycles_purchased,omitempty"`
	TotalCycles       int                   `json:"total_cycles,omitempty"`
	OrderID           string                `json:"order_id,omitempty"`
	OrderStatus       string                `json:"order_status,omitempty"`
	TotalCost         float64               `json:"total_cost,omitempty"`
	NextDelegation    int64                 `json:"next_delegation_time,omitempty"`
	Analysis          netts.TransferDetails `json:"analysis"`
	OrderRaw          *netts.OrderResult    `json:"-"`
	Status            *netts.AddressStatus  `json:"-"`
}

// EnergyService coordinates Netts energy rentals.
type NettsClient interface {
	AnalyzeUSDT(ctx context.Context, req netts.AnalyzeUSDTRequest) (*netts.AnalyzeUSDTData, error)
	GetAddressStatus(ctx context.Context, address string) (*netts.AddressStatus, error)
	AddHostAddress(ctx context.Context, address, callback string) (*netts.AddAddressResult, error)
	OrderCycles(ctx context.Context, address string, cycles int) (*netts.OrderResult, error)
}

type EnergyService struct {
	cfg         appconfig.EnergyConfig
	nettsClient NettsClient
	logger      *slog.Logger
}

// NewEnergyService instantiates an EnergyService.
func NewEnergyService(cfg appconfig.EnergyConfig, client NettsClient, logger *slog.Logger) *EnergyService {
	return &EnergyService{
		cfg:         cfg,
		nettsClient: client,
		logger:      logger,
	}
}

// EnsureEnergy ensures an address has sufficient energy and cycles.
func (s *EnergyService) EnsureEnergy(ctx context.Context, req EnsureEnergyRequest) (*EnsureEnergyResponse, error) {
	if err := s.validateRequest(req); err != nil {
		return nil, err
	}

	amount := req.Amount
	if amount == "" {
		amount = s.cfg.DefaultAnalyzeValue
	}

	analysis, err := s.nettsClient.AnalyzeUSDT(ctx, netts.AnalyzeUSDTRequest{
		Sender:   req.FromAddress,
		Receiver: req.ToAddress,
		Amount:   amount,
	})
	if err != nil {
		return nil, fmt.Errorf("analyze transfer: %w", err)
	}

	transfer := analysis.TransferDetails
	energyNeeded := transfer.EnergyNeeded
	if energyNeeded == 0 {
		energyNeeded = transfer.RecommendedEnergy
	}

	s.logger.InfoContext(ctx, "analysis completed",
		"from", req.FromAddress,
		"to", req.ToAddress,
		"recommended_energy", transfer.RecommendedEnergy,
		"energy_needed", energyNeeded,
	)

	status, addressAdded, err := s.ensureHostMode(ctx, req.FromAddress, req.CallbackURL)
	if err != nil {
		return nil, err
	}

	var (
		cyclesBefore = 0
		cyclesAfter  = 0
		order        *netts.OrderResult
	)

	if status != nil {
		cyclesBefore = status.CyclesRemaining
		if shouldOrderCycles(status, s.cfg.MinCycles) {
			targetCycles := s.cfg.TargetCycles
			if targetCycles < s.cfg.MinCycles {
				targetCycles = s.cfg.MinCycles
			}
			cyclesToBuy := targetCycles - safeCycleValue(status.CyclesRemaining)
			if cyclesToBuy > 0 {
				order, err = s.orderCycles(ctx, req.FromAddress, cyclesToBuy)
				if err != nil {
					return nil, err
				}
				if s.cfg.PostOrderWait.Duration() > 0 {
					select {
					case <-time.After(s.cfg.PostOrderWait.Duration()):
					case <-ctx.Done():
						return nil, ctx.Err()
					}
				}
				status, _, err = s.ensureHostMode(ctx, req.FromAddress, req.CallbackURL)
				if err != nil {
					return nil, err
				}
			}
		}
	}

	if status != nil {
		cyclesAfter = status.CyclesRemaining
	}

	resp := &EnsureEnergyResponse{
		FromAddress:       req.FromAddress,
		ToAddress:         req.ToAddress,
		Amount:            amount,
		AddressAdded:      addressAdded,
		RecommendedEnergy: transfer.RecommendedEnergy,
		EnergyNeeded:      energyNeeded,
		CyclesBefore:      cyclesBefore,
		CyclesAfter:       cyclesAfter,
		Analysis:          transfer,
		OrderRaw:          order,
		Status:            status,
	}

	if order != nil {
		resp.CyclesPurchased = order.CyclesPurchased
		resp.TotalCycles = order.TotalCycles
		resp.OrderID = order.OrderID
		resp.TotalCost = order.TotalCost
		resp.OrderStatus = order.Status
		resp.NextDelegation = order.NextDelegation
	}

	return resp, nil
}

func (s *EnergyService) ensureHostMode(ctx context.Context, address, callbackURL string) (*netts.AddressStatus, bool, error) {
	status, err := s.nettsClient.GetAddressStatus(ctx, address)
	if err == nil {
		return status, false, nil
	}
	if !errors.Is(err, netts.ErrAddressNotFound) {
		return nil, false, fmt.Errorf("get address status: %w", err)
	}
	if !s.cfg.AutoAddHost {
		return nil, false, fmt.Errorf("address %s not in host mode and autoAddHost disabled", address)
	}
	if _, err := s.nettsClient.AddHostAddress(ctx, address, callbackURL); err != nil {
		return nil, false, fmt.Errorf("add host address: %w", err)
	}
	s.logger.InfoContext(ctx, "address added to Host Mode", "address", address)

	status, err = s.nettsClient.GetAddressStatus(ctx, address)
	if err != nil {
		return nil, true, fmt.Errorf("get address status after add: %w", err)
	}
	return status, true, nil
}

func (s *EnergyService) orderCycles(ctx context.Context, address string, cycles int) (*netts.OrderResult, error) {
	if cycles <= 0 {
		return nil, nil
	}
	if cycles > 1000 {
		cycles = 1000
	}
	order, err := s.nettsClient.OrderCycles(ctx, address, cycles)
	if err != nil {
		return nil, fmt.Errorf("order cycles: %w", err)
	}
	s.logger.InfoContext(ctx, "cycles ordered",
		"address", address,
		"cycles_purchased", order.CyclesPurchased,
		"total_cycles", order.TotalCycles,
		"order_id", order.OrderID,
	)
	return order, nil
}

func (s *EnergyService) validateRequest(req EnsureEnergyRequest) error {
	if !tronAddressPattern.MatchString(req.FromAddress) {
		return fmt.Errorf("invalid from_address: %s", req.FromAddress)
	}
	if !tronAddressPattern.MatchString(req.ToAddress) {
		return fmt.Errorf("invalid to_address: %s", req.ToAddress)
	}
	return nil
}

func shouldOrderCycles(status *netts.AddressStatus, minCycles int) bool {
	if status == nil {
		return true
	}
	if status.CyclesRemaining < 0 {
		// Infinity mode (-1) requires no manual orders.
		return false
	}
	return status.CyclesRemaining < minCycles
}

func safeCycleValue(value int) int {
	if value < 0 {
		return 0
	}
	return value
}