┌─────────────────────────────────────────────────────────────────┐
│                    BATCH SESSION FLOW                            │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│  1. JoinBatchSession(intent) ──► Connects to ASP stream          │
│                          │                                       │
│                          ▼                                       │
│  2. BatchStarted ──► Receive batch parameters (round ID, etc.)  │
│                          │                                       │
│                          ▼                                       │
│  3. TreeSigningStarted ──► Receive unsigned VTXO tree           │
│                          │                                       │
│                          ▼                                       │
│  4. Create tree signer, generate nonces, send to server         │
│                          │                                       │
│                          ▼                                       │
│  5. TreeNoncesAggregated ──► Receive aggregated nonces          │
│                          │                                       │
│                          ▼                                       │
│  6. Sign tree with aggregated nonces, send partial sigs         │
│                          │                                       │
│                          ▼                                       │
│  7. BatchFinalization ──► Round is being finalized              │
│                          │                                       │
│                          ▼                                       │
│  8. BatchFinalized ──► Round complete, VTXOs created            │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘

type BatchEventsHandler interface {
    // Called when batch starts - receive round parameters
    HandleBatchStarted(
        roundID string,
        unsignedVtxoTree *tree.TxTree,
        batchOutSweepTapscript []byte,
    ) error

    // Called when tree is ready for signing
    HandleTreeSigningStarted(
        cosigners []*btcec.PublicKey,
        unsignedVtxoTree *tree.TxTree,
    ) (nonces tree.TreeNonces, err error)

    // Called when all nonces are aggregated
    HandleTreeNoncesAggregated(
        aggregatedNonces tree.TreeNonces,
    ) (sigs tree.TreePartialSigs, err error)

    // Called during finalization
    HandleBatchFinalization(
        signedVtxoTree *tree.TxTree,
        connectorTxs []*wire.MsgTx,
        roundTx string,
    ) error

    // Called when batch is fully finalized
    HandleBatchFinalized(
        roundTxid string,
    ) error
}

func JoinBatchSession(
    ctx context.Context,
    arkClient client.TransportClient,  // gRPC/REST client
    handler BatchEventsHandler,         // Event handler
    intent client.Intent,               // User's settlement intent
) error {
    // Create batch event stream from server
    eventCh, err := arkClient.JoinBatchSession(ctx, intent)
    if err != nil {
        return err
    }

    // Current state in the state machine
    currentState := start

    // Process events from stream
    for event := range eventCh {
        if event.Err != nil {
            return event.Err
        }

        switch e := event.Event.(type) {
        case client.BatchStarted:
            // Handle batch start...
        case client.TreeSigningStarted:
            // Handle tree signing...
        // ... other events
        }
    }

    return nil
}

const (
    start                 = "start"
    batchStarted          = "batchStarted"
    treeSigningStarted    = "treeSigningStarted"
    treeNoncesAggregated  = "treeNoncesAggregated"
    batchFinalization     = "batchFinalization"
)

// Valid state transitions
switch currentState {
case start:
    // Can transition to: batchStarted
    if event is BatchStarted {
        currentState = batchStarted
    }

case batchStarted:
    // Can transition to: treeSigningStarted
    if event is TreeSigningStarted {
        currentState = treeSigningStarted
    }

case treeSigningStarted:
    // Can transition to: treeNoncesAggregated
    if event is TreeNoncesAggregated {
        currentState = treeNoncesAggregated
    }

case treeNoncesAggregated:
    // Can transition to: batchFinalization
    if event is BatchFinalization {
        currentState = batchFinalization
    }

case batchFinalization:
    // Terminal state - BatchFinalized ends the loop
}

case client.BatchStarted:
    if currentState != start {
        return fmt.Errorf(
            "invalid state transition: %s -> batchStarted",
            currentState,
        )
    }

    // Parse the unsigned VTXO tree
    unsignedVtxoTree, err := tree.FromFlatTxTree(e.UnsignedVtxoTree)
    if err != nil {
        return err
    }

    // Notify handler
    if err := handler.HandleBatchStarted(
        e.RoundID,
        unsignedVtxoTree,
        e.BatchOutSweepTapscript,
    ); err != nil {
        return err
    }

    currentState = batchStarted

case client.TreeSigningStarted:
    if currentState != batchStarted {
        return fmt.Errorf(
            "invalid state transition: %s -> treeSigningStarted",
            currentState,
        )
    }

    // Parse cosigner public keys
    cosigners := make([]*btcec.PublicKey, 0, len(e.Cosigners))
    for _, c := range e.Cosigners {
        pubkeyBytes, _ := hex.DecodeString(c)
        pubkey, _ := schnorr.ParsePubKey(pubkeyBytes)
        cosigners = append(cosigners, pubkey)
    }

    // Parse unsigned tree
    unsignedVtxoTree, _ := tree.FromFlatTxTree(e.UnsignedVtxoTree)

    // Handler generates nonces
    nonces, err := handler.HandleTreeSigningStarted(cosigners, unsignedVtxoTree)
    if err != nil {
        return err
    }

    // Send nonces to server
    serializedNonces := serializeNonces(nonces)
    if err := arkClient.SendTreeNonces(ctx, e.RoundID, serializedNonces); err != nil {
        return err
    }

    currentState = treeSigningStarted

case client.TreeNoncesAggregated:
    if currentState != treeSigningStarted {
        return fmt.Errorf(
            "invalid state transition: %s -> treeNoncesAggregated",
            currentState,
        )
    }

    // Parse aggregated nonces
    aggregatedNonces := deserializeNonces(e.AggregatedNonces)

    // Handler signs the tree
    sigs, err := handler.HandleTreeNoncesAggregated(aggregatedNonces)
    if err != nil {
        return err
    }

    // Send partial signatures to server
    serializedSigs := serializeSigs(sigs)
    if err := arkClient.SendTreeSigs(ctx, e.RoundID, serializedSigs); err != nil {
        return err
    }

    currentState = treeNoncesAggregated

type defaultBatchEventsHandler struct {
    wallet              wallet.WalletService
    explorer            explorer.Explorer
    signerSession       tree.SignerSession
    batchOutSweepClosure []byte
    cosigners           []*btcec.PublicKey
    unsignedVtxoTree    *tree.TxTree
    signedVtxoTree      *tree.TxTree
    connectorTxs        []*wire.MsgTx
    roundTx             string
    roundTxid           string
}

func NewDefaultBatchEventsHandler(
    wallet wallet.WalletService,
    explorer explorer.Explorer,
) BatchEventsHandler {
    return &defaultBatchEventsHandler{
        wallet:   wallet,
        explorer: explorer,
    }
}

func (h *defaultBatchEventsHandler) HandleBatchStarted(
    roundID string,
    unsignedVtxoTree *tree.TxTree,
    batchOutSweepTapscript []byte,
) error {
    h.unsignedVtxoTree = unsignedVtxoTree
    h.batchOutSweepClosure = batchOutSweepTapscript
    return nil
}

func (h *defaultBatchEventsHandler) HandleTreeSigningStarted(
    cosigners []*btcec.PublicKey,
    unsignedVtxoTree *tree.TxTree,
) (tree.TreeNonces, error) {
    ctx := context.Background()
    h.cosigners = cosigners
    h.unsignedVtxoTree = unsignedVtxoTree

    // Create a new tree signer session using the wallet
    signerSession, err := h.wallet.NewVtxoTreeSigner(ctx, "")
    if err != nil {
        return nil, err
    }
    h.signerSession = signerSession

    // Calculate batch output amount from tree root
    batchOutAmount := h.unsignedVtxoTree.Root.Outputs[0].Value

    // Initialize the signer session
    if err := signerSession.Init(
        h.batchOutSweepClosure,
        batchOutAmount,
        h.unsignedVtxoTree,
    ); err != nil {
        return nil, err
    }

    // Generate and return nonces
    return signerSession.GetNonces()
}

func (h *defaultBatchEventsHandler) HandleTreeNoncesAggregated(
    aggregatedNonces tree.TreeNonces,
) (tree.TreePartialSigs, error) {
    // Set the aggregated nonces on the signer session
    h.signerSession.SetAggregatedNonces(aggregatedNonces)

    // Generate partial signatures for all tree transactions
    return h.signerSession.Sign()
}

func (h *defaultBatchEventsHandler) HandleBatchFinalization(
    signedVtxoTree *tree.TxTree,
    connectorTxs []*wire.MsgTx,
    roundTx string,
) error {
    h.signedVtxoTree = signedVtxoTree
    h.connectorTxs = connectorTxs
    h.roundTx = roundTx
    return nil
}

func (h *defaultBatchEventsHandler) HandleBatchFinalized(
    roundTxid string,
) error {
    h.roundTxid = roundTxid
    return nil
}

// Getter methods for results
func (h *defaultBatchEventsHandler) SignedVtxoTree() *tree.TxTree {
    return h.signedVtxoTree
}

func (h *defaultBatchEventsHandler) ConnectorTxs() []*wire.MsgTx {
    return h.connectorTxs
}

func (h *defaultBatchEventsHandler) RoundTx() string {
    return h.roundTx
}

func (h *defaultBatchEventsHandler) RoundTxid() string {
    return h.roundTxid
}

func serializeNonces(nonces tree.TreeNonces) map[string]string {
    result := make(map[string]string, len(nonces))
    for txid, nonce := range nonces {
        result[txid] = hex.EncodeToString(nonce.PubNonce[:])
    }
    return result
}

func deserializeNonces(serialized map[string]string) tree.TreeNonces {
    result := make(tree.TreeNonces, len(serialized))
    for txid, nonceHex := range serialized {
        nonceBytes, _ := hex.DecodeString(nonceHex)
        nonce := &tree.Musig2Nonce{}
        copy(nonce.PubNonce[:], nonceBytes)
        result[txid] = nonce
    }
    return result
}

func serializeSigs(sigs tree.TreePartialSigs) map[string]string {
    result := make(map[string]string, len(sigs))
    for txid, sig := range sigs {
        result[txid] = hex.EncodeToString(sig.S[:])
    }
    return result
}

// Inside arkClient.Settle():
func (a *arkClient) Settle(ctx context.Context, opts ...Option) (string, error) {
    // Create settlement intent
    intent := client.Intent{
        Inputs:  selectedCoins,
        Outputs: receivers,
    }

    // Create event handler
    handler := NewDefaultBatchEventsHandler(a.wallet, a.explorer)

    // Join the batch session and process all events
    if err := JoinBatchSession(ctx, a.client, handler, intent); err != nil {
        return "", err
    }

    // Get results from handler
    roundTxid := handler.RoundTxid()
    signedTree := handler.SignedVtxoTree()

    // Update local state with new VTXOs...

    return roundTxid, nil
}

// The Settle function internally uses JoinBatchSession
txid, err := client.Settle(ctx)
if err != nil {
    return err
}
fmt.Printf("Settled in round: %s\n", txid)

type myHandler struct {
    *defaultBatchEventsHandler
    onProgress func(state string)
}

func (h *myHandler) HandleBatchStarted(
    roundID string,
    unsignedVtxoTree *tree.TxTree,
    batchOutSweepTapscript []byte,
) error {
    h.onProgress("batch_started")
    return h.defaultBatchEventsHandler.HandleBatchStarted(
        roundID, unsignedVtxoTree, batchOutSweepTapscript,
    )
}

// Use custom handler
handler := &myHandler{
    defaultBatchEventsHandler: NewDefaultBatchEventsHandler(wallet, explorer),
    onProgress: func(state string) {
        log.Printf("Progress: %s", state)
    },
}

type progressHandler struct {
    BatchEventsHandler
    states []string
}

func (h *progressHandler) HandleBatchStarted(...) error {
    h.states = append(h.states, "started")
    return h.BatchEventsHandler.HandleBatchStarted(...)
}

func (h *progressHandler) HandleTreeSigningStarted(...) (tree.TreeNonces, error) {
    h.states = append(h.states, "signing")
    return h.BatchEventsHandler.HandleTreeSigningStarted(...)
}

// After JoinBatchSession completes:
fmt.Printf("States traversed: %v\n", handler.states)

err := JoinBatchSession(ctx, arkClient, handler, intent)
if err != nil {
    switch {
    case strings.Contains(err.Error(), "invalid state transition"):
        // Server sent events out of order - bug or network issue
        log.Error("Protocol violation - out of order events")
    case strings.Contains(err.Error(), "context canceled"):
        // Client canceled the operation
        log.Info("Batch session canceled by user")
    case strings.Contains(err.Error(), "round failed"):
        // Round was aborted by server
        log.Warn("Round failed - will retry in next round")
    default:
        log.Errorf("Batch session error: %v", err)
    }
}