a301_server/internal/chain/service.go

package chain

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"sync"
	"time"

	"github.com/tolelom/tolchain/core"
	tocrypto "github.com/tolelom/tolchain/crypto"
	"github.com/tolelom/tolchain/wallet"
)

type Service struct {
	repo           *Repository
	client         *Client
	chainID        string
	operatorWallet *wallet.Wallet
	encKeyBytes    []byte // 32-byte AES-256 key
	userResolver   func(username string) (uint, error)
	operatorMu     sync.Mutex // serialises operator-nonce transactions
	userMu         sync.Map   // per-user mutex (keyed by userID uint)
}

// SetUserResolver sets the callback that resolves username → userID.
func (s *Service) SetUserResolver(fn func(username string) (uint, error)) {
	s.userResolver = fn
}

// resolveUsername converts a username to the user's on-chain pubKeyHex.
func (s *Service) resolveUsername(username string) (string, error) {
	if s.userResolver == nil {
		return "", fmt.Errorf("user resolver not configured")
	}
	userID, err := s.userResolver(username)
	if err != nil {
		return "", fmt.Errorf("user not found")
	}
	uw, err := s.repo.FindByUserID(userID)
	if err != nil {
		return "", fmt.Errorf("wallet not found")
	}
	return uw.PubKeyHex, nil
}

func NewService(
	repo *Repository,
	client *Client,
	chainID string,
	operatorKeyHex string,
	walletEncKeyHex string,
) (*Service, error) {
	encKey, err := hex.DecodeString(walletEncKeyHex)
	if err != nil || len(encKey) != 32 {
		return nil, fmt.Errorf("WALLET_ENCRYPTION_KEY must be 64 hex chars (32 bytes)")
	}

	var opWallet *wallet.Wallet
	if operatorKeyHex != "" {
		privKey, err := tocrypto.PrivKeyFromHex(operatorKeyHex)
		if err != nil {
			return nil, fmt.Errorf("invalid OPERATOR_KEY_HEX: %w", err)
		}
		opWallet = wallet.New(privKey)
	}

	return &Service{
		repo:           repo,
		client:         client,
		chainID:        chainID,
		operatorWallet: opWallet,
		encKeyBytes:    encKey,
	}, nil
}

// ---- Wallet Encryption (AES-256-GCM) ----

func (s *Service) encryptPrivKey(privKey tocrypto.PrivateKey) (cipherHex, nonceHex string, err error) {
	block, err := aes.NewCipher(s.encKeyBytes)
	if err != nil {
		return "", "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", "", err
	}
	cipherText := gcm.Seal(nil, nonce, []byte(privKey), nil)
	return hex.EncodeToString(cipherText), hex.EncodeToString(nonce), nil
}

func (s *Service) decryptPrivKey(cipherHex, nonceHex string) (tocrypto.PrivateKey, error) {
	cipherText, err := hex.DecodeString(cipherHex)
	if err != nil {
		return nil, err
	}
	nonce, err := hex.DecodeString(nonceHex)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(s.encKeyBytes)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	plaintext, err := gcm.Open(nil, nonce, cipherText, nil)
	if err != nil {
		return nil, fmt.Errorf("wallet decryption failed: %w", err)
	}
	return tocrypto.PrivateKey(plaintext), nil
}

// ---- Wallet Management ----

// CreateWallet generates a new keypair, encrypts it, and stores in DB.
func (s *Service) CreateWallet(userID uint) (*UserWallet, error) {
	w, err := wallet.Generate()
	if err != nil {
		return nil, fmt.Errorf("key generation failed: %w", err)
	}

	cipherHex, nonceHex, err := s.encryptPrivKey(w.PrivKey())
	if err != nil {
		return nil, fmt.Errorf("key encryption failed: %w", err)
	}

	uw := &UserWallet{
		UserID:           userID,
		PubKeyHex:        w.PubKey(),
		Address:          w.Address(),
		EncryptedPrivKey: cipherHex,
		EncNonce:         nonceHex,
	}
	if err := s.repo.Create(uw); err != nil {
		return nil, fmt.Errorf("wallet save failed: %w", err)
	}
	return uw, nil
}

func (s *Service) GetWallet(userID uint) (*UserWallet, error) {
	return s.repo.FindByUserID(userID)
}

// loadUserWallet decrypts a user's private key and returns a wallet.Wallet.
func (s *Service) loadUserWallet(userID uint) (*wallet.Wallet, string, error) {
	uw, err := s.repo.FindByUserID(userID)
	if err != nil {
		return nil, "", fmt.Errorf("wallet not found: %w", err)
	}
	privKey, err := s.decryptPrivKey(uw.EncryptedPrivKey, uw.EncNonce)
	if err != nil {
		log.Printf("WARNING: wallet decryption failed for userID=%d: %v", userID, err)
		return nil, "", fmt.Errorf("wallet decryption failed")
	}
	return wallet.New(privKey), uw.PubKeyHex, nil
}

func (s *Service) getNonce(address string) (uint64, error) {
	bal, err := s.client.GetBalance(address)
	if err != nil {
		return 0, fmt.Errorf("get nonce failed: %w", err)
	}
	return bal.Nonce, nil
}

// txConfirmTimeout is the maximum time to wait for a transaction to be
// included in a block. PoA block intervals are typically a few seconds,
// so 15s provides ample margin.
const txConfirmTimeout = 15 * time.Second

// submitTx sends a signed transaction and waits for block confirmation.
// Returns the confirmed status or an error (including TxError for on-chain failures).
func (s *Service) submitTx(tx any) (*TxStatusResult, error) {
	return s.client.SendTxAndWait(tx, txConfirmTimeout)
}

// ---- Query Methods ----

func (s *Service) GetBalance(userID uint) (*BalanceResult, error) {
	uw, err := s.repo.FindByUserID(userID)
	if err != nil {
		return nil, fmt.Errorf("wallet not found: %w", err)
	}
	return s.client.GetBalance(uw.PubKeyHex)
}

func (s *Service) GetAssets(userID uint, offset, limit int) (json.RawMessage, error) {
	uw, err := s.repo.FindByUserID(userID)
	if err != nil {
		return nil, fmt.Errorf("wallet not found: %w", err)
	}
	return s.client.GetAssetsByOwner(uw.PubKeyHex, offset, limit)
}

func (s *Service) GetAsset(assetID string) (json.RawMessage, error) {
	return s.client.GetAsset(assetID)
}

func (s *Service) GetInventory(userID uint) (json.RawMessage, error) {
	uw, err := s.repo.FindByUserID(userID)
	if err != nil {
		return nil, fmt.Errorf("wallet not found: %w", err)
	}
	return s.client.GetInventory(uw.PubKeyHex)
}

func (s *Service) GetMarketListings(offset, limit int) (json.RawMessage, error) {
	return s.client.GetActiveListings(offset, limit)
}

func (s *Service) GetListing(listingID string) (json.RawMessage, error) {
	return s.client.GetListing(listingID)
}

// getUserMu returns a per-user mutex, creating one if it doesn't exist.
func (s *Service) getUserMu(userID uint) *sync.Mutex {
	v, _ := s.userMu.LoadOrStore(userID, &sync.Mutex{})
	return v.(*sync.Mutex)
}

// ---- User Transaction Methods ----

// userTx handles the common boilerplate for user transactions:
// acquire per-user mutex → load wallet → get nonce → build tx → submit.
func (s *Service) userTx(userID uint, buildFn func(w *wallet.Wallet, nonce uint64) (any, error)) (*TxStatusResult, error) {
	mu := s.getUserMu(userID)
	mu.Lock()
	defer mu.Unlock()
	w, pubKey, err := s.loadUserWallet(userID)
	if err != nil {
		return nil, err
	}
	nonce, err := s.getNonce(pubKey)
	if err != nil {
		return nil, err
	}
	tx, err := buildFn(w, nonce)
	if err != nil {
		return nil, fmt.Errorf("build tx failed: %w", err)
	}
	return s.submitTx(tx)
}

func (s *Service) Transfer(userID uint, to string, amount uint64) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.Transfer(s.chainID, to, amount, nonce, 0)
	})
}

func (s *Service) TransferAsset(userID uint, assetID, to string) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.TransferAsset(s.chainID, assetID, to, nonce, 0)
	})
}

func (s *Service) ListOnMarket(userID uint, assetID string, price uint64) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.ListMarket(s.chainID, assetID, price, nonce, 0)
	})
}

func (s *Service) BuyFromMarket(userID uint, listingID string) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.BuyMarket(s.chainID, listingID, nonce, 0)
	})
}

func (s *Service) CancelListing(userID uint, listingID string) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.CancelListing(s.chainID, listingID, nonce, 0)
	})
}

func (s *Service) EquipItem(userID uint, assetID, slot string) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.EquipItem(s.chainID, assetID, slot, nonce, 0)
	})
}

func (s *Service) UnequipItem(userID uint, assetID string) (*TxStatusResult, error) {
	return s.userTx(userID, func(w *wallet.Wallet, nonce uint64) (any, error) {
		return w.UnequipItem(s.chainID, assetID, nonce, 0)
	})
}

// ---- Operator Transaction Methods ----

func (s *Service) ensureOperator() error {
	if s.operatorWallet == nil {
		return fmt.Errorf("operator wallet not configured")
	}
	return nil
}

func (s *Service) getOperatorNonce() (uint64, error) {
	if err := s.ensureOperator(); err != nil {
		return 0, err
	}
	return s.getNonce(s.operatorWallet.PubKey())
}

// operatorTx handles the common boilerplate for operator transactions:
// acquire operator mutex → ensure operator → get nonce → build tx → submit.
func (s *Service) operatorTx(buildFn func(nonce uint64) (any, error)) (*TxStatusResult, error) {
	s.operatorMu.Lock()
	defer s.operatorMu.Unlock()
	if err := s.ensureOperator(); err != nil {
		return nil, err
	}
	nonce, err := s.getOperatorNonce()
	if err != nil {
		return nil, err
	}
	tx, err := buildFn(nonce)
	if err != nil {
		return nil, fmt.Errorf("build tx failed: %w", err)
	}
	return s.submitTx(tx)
}

func (s *Service) MintAsset(templateID, ownerPubKey string, properties map[string]any) (*TxStatusResult, error) {
	return s.operatorTx(func(nonce uint64) (any, error) {
		return s.operatorWallet.MintAsset(s.chainID, templateID, ownerPubKey, properties, nonce, 0)
	})
}

func (s *Service) GrantReward(recipientPubKey string, tokenAmount uint64, assets []core.MintAssetPayload) (*TxStatusResult, error) {
	return s.operatorTx(func(nonce uint64) (any, error) {
		return s.operatorWallet.GrantReward(s.chainID, recipientPubKey, tokenAmount, assets, nonce, 0)
	})
}

func (s *Service) RegisterTemplate(id, name string, schema map[string]any, tradeable bool) (*TxStatusResult, error) {
	return s.operatorTx(func(nonce uint64) (any, error) {
		return s.operatorWallet.RegisterTemplate(s.chainID, id, name, schema, tradeable, nonce, 0)
	})
}

// ---- Username-based Methods (for game server) ----

func (s *Service) GrantRewardByUsername(username string, tokenAmount uint64, assets []core.MintAssetPayload) (*TxStatusResult, error) {
	pubKey, err := s.resolveUsername(username)
	if err != nil {
		return nil, err
	}
	return s.GrantReward(pubKey, tokenAmount, assets)
}

func (s *Service) MintAssetByUsername(templateID, username string, properties map[string]any) (*TxStatusResult, error) {
	pubKey, err := s.resolveUsername(username)
	if err != nil {
		return nil, err
	}
	return s.MintAsset(templateID, pubKey, properties)
}

func (s *Service) GetBalanceByUsername(username string) (*BalanceResult, error) {
	pubKey, err := s.resolveUsername(username)
	if err != nil {
		return nil, err
	}
	return s.client.GetBalance(pubKey)
}

func (s *Service) GetAssetsByUsername(username string, offset, limit int) (json.RawMessage, error) {
	pubKey, err := s.resolveUsername(username)
	if err != nil {
		return nil, err
	}
	return s.client.GetAssetsByOwner(pubKey, offset, limit)
}

func (s *Service) GetInventoryByUsername(username string) (json.RawMessage, error) {
	pubKey, err := s.resolveUsername(username)
	if err != nil {
		return nil, err
	}
	return s.client.GetInventory(pubKey)
}