a301_server/internal/chain/service.go

package chain

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

	"a301_server/pkg/apperror"

	"github.com/tolelom/tolchain/core"
	tocrypto "github.com/tolelom/tolchain/crypto"
	"github.com/tolelom/tolchain/wallet"
	"golang.org/x/crypto/hkdf"
)

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)
	passwordVerifier func(userID uint, password string) 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
}

func (s *Service) SetPasswordVerifier(fn func(userID uint, password string) error) {
	s.passwordVerifier = fn
}

func (s *Service) ExportPrivKey(userID uint, password string) (string, error) {
	if s.passwordVerifier == nil {
		return "", fmt.Errorf("password verifier not configured")
	}
	if err := s.passwordVerifier(userID, password); err != nil {
		return "", err
	}
	w, _, err := s.loadUserWallet(userID)
	if err != nil {
		return "", err
	}
	return w.PrivKey().Hex(), nil
}

// resolveUsername converts a username to the user's on-chain pubKeyHex.
// If the user exists but has no wallet (e.g. legacy user or failed creation),
// a wallet is auto-created on the fly.
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 {
		// 지갑이 없으면 자동 생성 시도
		var createErr error
		uw, createErr = s.CreateWallet(userID)
		if createErr != nil {
			if apperror.IsDuplicateEntry(createErr) {
				// unique constraint 위반 — 다른 고루틴이 먼저 생성 완료
				uw, err = s.repo.FindByUserID(userID)
				if err != nil {
					return "", fmt.Errorf("wallet auto-creation failed: %w", err)
				}
			} else {
				return "", fmt.Errorf("wallet auto-creation failed: %w", createErr)
			}
		} else {
			log.Printf("INFO: auto-created wallet for userID=%d (username=%s)", userID, username)
		}
	}
	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) derivePerWalletKey(salt []byte, userID uint) ([]byte, error) {
	info := []byte("wallet:" + strconv.FormatUint(uint64(userID), 10))
	r := hkdf.New(sha256.New, s.encKeyBytes, salt, info)
	key := make([]byte, 32)
	if _, err := io.ReadFull(r, key); err != nil {
		return nil, fmt.Errorf("HKDF key derivation failed: %w", err)
	}
	return key, nil
}

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
}

func (s *Service) encryptPrivKeyV2(privKey tocrypto.PrivateKey, userID uint) (cipherHex, nonceHex, saltHex string, err error) {
	salt := make([]byte, 16)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return "", "", "", err
	}
	key, err := s.derivePerWalletKey(salt, userID)
	if err != nil {
		return "", "", "", err
	}
	block, err := aes.NewCipher(key)
	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), hex.EncodeToString(salt), nil
}

func (s *Service) decryptPrivKeyV2(cipherHex, nonceHex, saltHex string, userID uint) (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
	}
	salt, err := hex.DecodeString(saltHex)
	if err != nil {
		return nil, err
	}
	key, err := s.derivePerWalletKey(salt, userID)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(key)
	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 Migration ----

// MigrateWalletKeys re-encrypts all v1 wallets using HKDF per-wallet keys.
// Each wallet is migrated individually; failures are logged and skipped.
func (s *Service) MigrateWalletKeys() error {
	wallets, err := s.repo.FindAllByKeyVersion(1)
	if err != nil {
		return fmt.Errorf("query v1 wallets: %w", err)
	}
	if len(wallets) == 0 {
		return nil
	}
	log.Printf("INFO: migrating %d v1 wallets to v2 (HKDF)", len(wallets))
	var migrated, failed int
	for _, uw := range wallets {
		privKey, err := s.decryptPrivKey(uw.EncryptedPrivKey, uw.EncNonce)
		if err != nil {
			log.Printf("ERROR: v1 decrypt failed for walletID=%d userID=%d: %v", uw.ID, uw.UserID, err)
			failed++
			continue
		}
		cipherHex, nonceHex, saltHex, err := s.encryptPrivKeyV2(privKey, uw.UserID)
		if err != nil {
			log.Printf("ERROR: v2 encrypt failed for walletID=%d userID=%d: %v", uw.ID, uw.UserID, err)
			failed++
			continue
		}
		if err := s.repo.UpdateEncryption(uw.ID, cipherHex, nonceHex, saltHex, 2); err != nil {
			log.Printf("ERROR: DB update failed for walletID=%d userID=%d: %v", uw.ID, uw.UserID, err)
			failed++
			continue
		}
		migrated++
	}
	log.Printf("INFO: wallet migration complete: %d migrated, %d failed", migrated, failed)
	return 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, saltHex, err := s.encryptPrivKeyV2(w.PrivKey(), userID)
	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,
		KeyVersion:       2,
		HKDFSalt:         saltHex,
	}
	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)
	}
	var privKey tocrypto.PrivateKey
	if uw.KeyVersion >= 2 {
		privKey, err = s.decryptPrivKeyV2(uw.EncryptedPrivKey, uw.EncNonce, uw.HKDFSalt, uw.UserID)
	} else {
		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)
}