// Command sweep is the offline-signed USDT 归集 tool (Phase D). It splits the
// work across the air gap so the mnemonic never touches an online machine:
//
//	sweep plan      (online)  list derived receiving addresses that hold USDT
//	sweep build     (online)  build unsigned transfers -> unsigned.json (no key)
//	sweep sign      (OFFLINE) verify + sign with the mnemonic -> signed.json
//	sweep broadcast (online)  submit signed.json to the chain
//
// The mnemonic is read from env PAY_SWEEP_MNEMONIC (set on the air-gapped box
// from Bitwarden), never a CLI arg and never on the hot service. `sign`
// independently re-verifies every transaction (recomputes the txid from
// raw_data, checks the recipient+amount+contract are embedded, and checks the
// derived key matches the owner) so a compromised online builder cannot trick it
// into signing a payment to someone else.
package main

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/wangjia/pangolin/pay/internal/tron"
	"github.com/wangjia/pangolin/pay/internal/wallet"
)

type item struct {
	Index    uint32          `json:"index"`
	Owner    string          `json:"owner"`
	Cold     string          `json:"cold"`
	Amount   int64           `json:"amount"` // micro-USDT
	Unsigned tron.UnsignedTx `json:"unsigned"`
	Signed   *tron.SignedTx  `json:"signed,omitempty"`
}

func main() {
	if len(os.Args) < 2 {
		fail("usage: sweep <plan|build|sign|broadcast> [flags]")
	}
	switch os.Args[1] {
	case "plan":
		cmdPlan(os.Args[2:])
	case "build":
		cmdBuild(os.Args[2:])
	case "sign":
		cmdSign(os.Args[2:])
	case "broadcast":
		cmdBroadcast(os.Args[2:])
	default:
		fail("unknown subcommand %q (plan|build|sign|broadcast)", os.Args[1])
	}
}

func env(k, def string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return def
}

func usdtContract() string { return env("USDT_CONTRACT", tron.USDTContractMainnet) }

func newClient() *tron.Client {
	return tron.NewClient(env("TRONGRID_BASE", ""), usdtContract(), os.Getenv("TRONGRID_API_KEY"))
}

// cmdPlan (online): scan indices 0..max, print those with a USDT balance.
func cmdPlan(args []string) {
	fs := flag.NewFlagSet("plan", flag.ExitOnError)
	max := fs.Uint("max", 50, "highest address index to scan")
	_ = fs.Parse(args)
	xpub := mustEnv("PAY_ACCOUNT_XPUB")
	c := newClient()
	ctx := context.Background()
	var total int64
	for i := uint32(0); i <= uint32(*max); i++ {
		addr, err := wallet.AddressFromAccountXpub(xpub, 0, i)
		if err != nil {
			fail("derive %d: %v", i, err)
		}
		bal, err := c.TRC20Balance(ctx, addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "warn: balance %d %s: %v\n", i, addr, err)
			continue
		}
		if bal > 0 {
			total += bal
			fmt.Printf("index=%d\taddr=%s\tusdt=%s\n", i, addr, usdt(bal))
		}
	}
	fmt.Printf("TOTAL: %s USDT\n", usdt(total))
}

// cmdBuild (online): build unsigned transfers of each address's full balance to
// the cold wallet. No private key used.
func cmdBuild(args []string) {
	fs := flag.NewFlagSet("build", flag.ExitOnError)
	cold := fs.String("cold", "", "cold wallet destination address (T...)")
	max := fs.Uint("max", 50, "highest address index to scan")
	feeLimit := fs.Int64("fee-limit", 30_000_000, "fee limit in sun (1e-6 TRX)")
	_ = fs.Parse(args)
	if *cold == "" {
		fail("--cold is required")
	}
	if _, err := wallet.DecodeTronAddress(*cold); err != nil {
		fail("bad --cold: %v", err)
	}
	xpub := mustEnv("PAY_ACCOUNT_XPUB")
	c := newClient()
	ctx := context.Background()
	var out []item
	for i := uint32(0); i <= uint32(*max); i++ {
		owner, err := wallet.AddressFromAccountXpub(xpub, 0, i)
		if err != nil {
			fail("derive %d: %v", i, err)
		}
		bal, err := c.TRC20Balance(ctx, owner)
		if err != nil || bal <= 0 {
			continue
		}
		ut, err := c.BuildTransfer(ctx, owner, *cold, bal, *feeLimit)
		if err != nil {
			fail("build %d %s: %v", i, owner, err)
		}
		out = append(out, item{Index: i, Owner: owner, Cold: *cold, Amount: bal, Unsigned: *ut})
		fmt.Fprintf(os.Stderr, "built index=%d %s -> %s %s USDT\n", i, owner, *cold, usdt(bal))
	}
	emit(out)
}

// cmdSign (OFFLINE): verify each tx independently, then sign with the mnemonic.
func cmdSign(args []string) {
	fs := flag.NewFlagSet("sign", flag.ExitOnError)
	cold := fs.String("cold", "", "expected cold destination (guards against tampering)")
	_ = fs.Parse(args)
	mnemonic := mustEnv("PAY_SWEEP_MNEMONIC") // set on the air-gapped box from Bitwarden
	items := read()
	contractBodyHex, err := wallet.TronAddressBodyHex(usdtContract())
	if err != nil {
		fail("usdt contract: %v", err)
	}
	for i := range items {
		it := &items[i]

		if *cold != "" && it.Cold != *cold {
			fail("index %d: cold %s != expected %s", it.Index, it.Cold, *cold)
		}
		// 1) raw_data integrity: recompute txid, must equal the claimed one.
		txid, err := tron.TxID(it.Unsigned.RawDataHex)
		if err != nil {
			fail("index %d: txid: %v", it.Index, err)
		}
		if !strings.EqualFold(hex.EncodeToString(txid), it.Unsigned.TxID) {
			fail("index %d: txid mismatch — raw_data tampered", it.Index)
		}
		// 2) recipient + amount: the exact ABI param must be embedded in raw_data.
		wantParam, err := tron.ABIEncodeTransferParams(it.Cold, it.Amount)
		if err != nil {
			fail("index %d: abi: %v", it.Index, err)
		}
		if !strings.Contains(strings.ToLower(it.Unsigned.RawDataHex), strings.ToLower(wantParam)) {
			fail("index %d: recipient/amount not found in raw_data — refusing to sign", it.Index)
		}
		// 3) contract: the USDT contract body must be in raw_data (right token).
		if !strings.Contains(strings.ToLower(it.Unsigned.RawDataHex), strings.ToLower(contractBodyHex)) {
			fail("index %d: USDT contract not found in raw_data — refusing to sign", it.Index)
		}
		// 4) key: the derived address for this index must equal the owner.
		addr, err := wallet.AddressFromMnemonic(mnemonic, "", 0, 0, it.Index)
		if err != nil {
			fail("index %d: derive addr: %v", it.Index, err)
		}
		if addr != it.Owner {
			fail("index %d: derived %s != owner %s — wrong mnemonic/index", it.Index, addr, it.Owner)
		}

		priv, err := wallet.PrivKeyHexFromMnemonic(mnemonic, "", 0, 0, it.Index)
		if err != nil {
			fail("index %d: privkey: %v", it.Index, err)
		}
		sig, err := tron.SignRawData(it.Unsigned.RawDataHex, priv)
		if err != nil {
			fail("index %d: sign: %v", it.Index, err)
		}
		it.Signed = &tron.SignedTx{
			TxID:       it.Unsigned.TxID,
			RawData:    it.Unsigned.RawData,
			RawDataHex: it.Unsigned.RawDataHex,
			Visible:    it.Unsigned.Visible,
			Signature:  []string{sig},
		}
		fmt.Fprintf(os.Stderr, "signed index=%d %s -> %s %s USDT\n", it.Index, it.Owner, it.Cold, usdt(it.Amount))
	}
	emit(items)
}

// cmdBroadcast (online): submit each signed tx.
func cmdBroadcast(args []string) {
	_ = flag.NewFlagSet("broadcast", flag.ExitOnError).Parse(args)
	items := read()
	c := newClient()
	ctx := context.Background()
	for i := range items {
		it := &items[i]
		if it.Signed == nil {
			fail("index %d: not signed", it.Index)
		}
		txid, err := c.Broadcast(ctx, it.Signed)
		if err != nil {
			fail("index %d: broadcast: %v", it.Index, err)
		}
		fmt.Printf("broadcast index=%d %s USDT -> tx %s\n", it.Index, usdt(it.Amount), txid)
		time.Sleep(200 * time.Millisecond) // be gentle with the endpoint
	}
}

// --- helpers ---

func usdt(micro int64) string {
	return strconv.FormatFloat(float64(micro)/1e6, 'f', 6, 64)
}

func mustEnv(k string) string {
	v := os.Getenv(k)
	if v == "" {
		fail("env %s is required", k)
	}
	return v
}

func read() []item {
	var items []item
	if err := json.NewDecoder(os.Stdin).Decode(&items); err != nil {
		fail("read json from stdin: %v", err)
	}
	return items
}

func emit(items []item) {
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	if err := enc.Encode(items); err != nil {
		fail("write json: %v", err)
	}
}

func fail(format string, a ...any) {
	fmt.Fprintf(os.Stderr, "sweep: "+format+"\n", a...)
	os.Exit(1)
}