package bitcoinElectrum

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sync/atomic"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/btcutil/hdkeychain"
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/txscript"
)

// Derives a P2WPKH address at a specific index from a branch key.
func deriveAddress(branchKey *hdkeychain.ExtendedKey, index uint32) (btcutil.Address, error) {
	childKey, err := branchKey.Derive(index)
	if err != nil {
		return nil, err
	}
	pubKey, err := childKey.ECPubKey()
	if err != nil {
		return nil, err
	}
	// Note: For zpub/vpub, use NewAddressWitnessPubKeyHash.
	// For ypub/upub, use NewAddressScriptHash with a P2WPKH script.
	// For xpub/tpub, use NewAddressPubKeyHash.
	return btcutil.NewAddressWitnessPubKeyHash(btcutil.Hash160(pubKey.SerializeCompressed()), &chaincfg.MainNetParams)
}

// Converts a btcutil.Address to an Electrum scripthash.
func addressToScriptHash(address btcutil.Address) (string, error) {
	script, err := txscript.PayToAddrScript(address)
	if err != nil {
		return "", err
	}
	hash := sha256.Sum256(script)
	// Reverse byte order for Electrum protocol
	for i, j := 0, len(hash)-1; i < j; i, j = i+1, j-1 {
		hash[i], hash[j] = hash[j], hash[i]
	}
	return hex.EncodeToString(hash[:]), nil
}

// Scans a derivation branch (m/0 or m/1) for balances until the gap limit is reached.
func scanBranch(masterKey *hdkeychain.ExtendedKey, branch uint32, gapLimit int, spvServer string) (int64, error) {
	branchKey, err := masterKey.Derive(branch)
	if err != nil {
		return 0, fmt.Errorf("error deriving branch %d: %v", branch, err)
	}

	var branchTotalBalance int64
	unusedAddressCount := 0

	for i := uint32(0); ; i++ {
		if unusedAddressCount >= gapLimit {
			break
		}

		// 3. Derive child address
		address, err := deriveAddress(branchKey, i)
		if err != nil {
			return 0, fmt.Errorf("could not derive address at index %d on branch %d: %v", i, branch, err)
		}

		// 4. Get balance from Electrum server
		balance, err := getAddressBalance(spvServer, address)
		if err != nil {
			return 0, fmt.Errorf("error getting balance for %s: %v", address.EncodeAddress(), err)
		}

		total := balance.Result.Confirmed + balance.Result.Unconfirmed
		if total > 0 {
			atomic.AddInt64(&branchTotalBalance, total)
			unusedAddressCount = 0 // Reset gap counter on finding a used address
		} else {
			unusedAddressCount++
		}
	}
	return branchTotalBalance, nil
}