DebTik/internal/bwtest/client.go

package bwtest

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

var (
	ErrTestAlreadyRunning = errors.New("a bandwidth test is already in progress")
	ErrServerRejected     = errors.New("remote server rejected bandwidth test request")
)

// BwTestClient manages executing bandwidth tests against a remote BwTest server
type BwTestClient struct {
	mu            sync.RWMutex
	running       bool
	currentCancel context.CancelFunc

	status  TestStatus
	samples []*LiveSample

	sampleCh chan LiveSample
}

// NewClient creates a new BwTestClient instance
func NewClient() *BwTestClient {
	return &BwTestClient{
		sampleCh: make(chan LiveSample, 100),
	}
}

// GetStatus returns the current or most recent test state
func (c *BwTestClient) GetStatus() TestStatus {
	c.mu.RLock()
	defer c.mu.RUnlock()

	st := c.status
	if len(c.samples) > 0 {
		st.CurrentSample = c.samples[len(c.samples)-1]
		st.RecentSamples = c.samples
	}
	return st
}

// Subscribe returns a channel of live telemetry samples
func (c *BwTestClient) Subscribe() <-chan LiveSample {
	return c.sampleCh
}

// Stop terminates any currently active bandwidth test
func (c *BwTestClient) Stop() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.running && c.currentCancel != nil {
		c.currentCancel()
	}
}

// Run executes a bandwidth test asynchronously
func (c *BwTestClient) Run(ctx context.Context, req SessionRequest) error {
	c.mu.Lock()
	if c.running {
		c.mu.Unlock()
		return ErrTestAlreadyRunning
	}

	if req.Target == "" {
		req.Target = "127.0.0.1"
	}
	if req.Port <= 0 {
		req.Port = 2000
	}
	if req.Protocol != "tcp" {
		req.Protocol = "udp"
	}
	if req.Direction != "tx" && req.Direction != "rx" && req.Direction != "both" {
		req.Direction = "both"
	}
	if req.PacketSize <= 0 {
		req.PacketSize = 1472
	}
	if req.Streams <= 0 {
		req.Streams = 1
	}
	if req.Streams > 8 {
		req.Streams = 8
	}
	if req.Duration <= 0 {
		req.Duration = 10
	}

	testCtx, cancel := context.WithCancel(context.Background())
	c.currentCancel = cancel
	c.running = true
	c.samples = make([]*LiveSample, 0, 120)

	sessionID := fmt.Sprintf("bw-%d-%d", time.Now().Unix(), rand.Intn(10000))
	c.status = TestStatus{
		Running:        true,
		SessionID:      sessionID,
		Target:         req.Target,
		Port:           req.Port,
		Protocol:       req.Protocol,
		Direction:      req.Direction,
		BandwidthLimit: req.BandwidthLimit,
		PacketSize:     req.PacketSize,
		Streams:        req.Streams,
		Duration:       req.Duration,
		StartTime:      time.Now(),
	}
	c.mu.Unlock()

	go c.executeTest(testCtx, cancel, req, sessionID)
	return nil
}

func (c *BwTestClient) executeTest(ctx context.Context, cancel context.CancelFunc, req SessionRequest, sessionID string) {
	defer func() {
		cancel()
		c.mu.Lock()
		c.running = false
		c.status.Running = false
		c.finalizeSummary()
		c.mu.Unlock()
	}()

	targetAddr := fmt.Sprintf("%s:%d", req.Target, req.Port)
	dialer := net.Dialer{Timeout: 5 * time.Second}
	controlConn, err := dialer.DialContext(ctx, "tcp", targetAddr)
	if err != nil {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Failed to connect to control port %s: %v", targetAddr, err)
		c.mu.Unlock()
		return
	}
	defer controlConn.Close()

	// If UDP, bind a local UDP socket
	var udpConn *net.UDPConn
	var clientUDPPort int
	if req.Protocol == "udp" {
		uaddr, err := net.ResolveUDPAddr("udp", ":0")
		if err == nil {
			udpConn, err = net.ListenUDP("udp", uaddr)
		}
		if err != nil {
			c.mu.Lock()
			c.status.Error = fmt.Sprintf("Failed to create local UDP socket: %v", err)
			c.mu.Unlock()
			return
		}
		defer udpConn.Close()
		_ = udpConn.SetReadBuffer(8 * 1024 * 1024)
		if la, ok := udpConn.LocalAddr().(*net.UDPAddr); ok {
			clientUDPPort = la.Port
		}
	}

	// Send handshake request
	hs := SessionHandshake{
		Type:           "request",
		SessionID:      sessionID,
		Protocol:       req.Protocol,
		Direction:      req.Direction,
		BandwidthLimit: req.BandwidthLimit,
		PacketSize:     req.PacketSize,
		Streams:        req.Streams,
		Duration:       req.Duration,
		UDPPort:        clientUDPPort,
	}
	b, _ := json.Marshal(hs)
	if _, err := controlConn.Write(append(b, '\n')); err != nil {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Handshake send failed: %v", err)
		c.mu.Unlock()
		return
	}

	// Read accept response
	reader := bufio.NewReader(controlConn)
	line, err := reader.ReadBytes('\n')
	if err != nil {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Server response error: %v", err)
		c.mu.Unlock()
		return
	}

	var acceptResp SessionHandshake
	if err := json.Unmarshal(line, &acceptResp); err != nil || acceptResp.Type != "accept" {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Server rejected test: %s", acceptResp.Message)
		c.mu.Unlock()
		return
	}

	serverUDPPort := acceptResp.UDPPort
	if serverUDPPort <= 0 {
		serverUDPPort = req.Port
	}

	// Statistics counters
	var (
		localTxBytes     atomic.Int64
		localTxPackets   atomic.Int64
		localRxBytes     atomic.Int64
		localRxPackets   atomic.Int64
		localLostPackets atomic.Int64
		serverRxBytes    atomic.Int64
		serverRxPackets  atomic.Int64
		serverLostPackets atomic.Int64

		curJitterMs atomic.Uint64 // float64 in bits
		jitterMu    sync.Mutex
		jitterCalc  JitterCalculator
		maxSeqRecv  atomic.Uint64
	)

	sessionHash := HashSessionID(sessionID)

	// Background worker for server feedback reports over TCP control socket
	go func() {
		for {
			line, err := reader.ReadBytes('\n')
			if err != nil {
				return
			}
			var rep SessionHandshake
			if err := json.Unmarshal(line, &rep); err == nil && rep.Type == "report" {
				serverRxBytes.Store(rep.RxBytes)
				serverRxPackets.Store(rep.RxPackets)
				serverLostPackets.Store(rep.LostPackets)
			}
		}
	}()

	var wg sync.WaitGroup

	// UDP Data path
	if req.Protocol == "udp" {
		serverUDPAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", req.Target, serverUDPPort))
		if err != nil {
			c.mu.Lock()
			c.status.Error = fmt.Sprintf("Failed to resolve server UDP address: %v", err)
			c.mu.Unlock()
			return
		}

		// Transmitter (TX or Both)
		if req.Direction == "tx" || req.Direction == "both" {
			for s := 0; s < req.Streams; s++ {
				wg.Add(1)
				streamID := uint8(s)
				go func() {
					defer wg.Done()
					limiter := NewRateLimiter(req.BandwidthLimit / int64(req.Streams))
					buf := make([]byte, req.PacketSize)
					var seq uint64

					hdr := PacketHeader{
						Magic:         MagicBytes,
						Version:       ProtoVer,
						PacketType:    PacketTypeData,
						StreamID:      streamID,
						SessionHash:   sessionHash,
						PayloadLength: uint32(req.PacketSize - HeaderSize),
					}

					for {
						select {
						case <-ctx.Done():
							return
						default:
						}

						if limiter != nil {
							limiter.Wait(req.PacketSize)
						}

						seq++
						hdr.Sequence = seq
						hdr.TimestampNano = time.Now().UnixNano()
						EncodeHeader(buf, hdr)

						n, err := udpConn.WriteToUDP(buf, serverUDPAddr)
						if err == nil {
							localTxBytes.Add(int64(n))
							localTxPackets.Add(1)
						}
					}
				}()
			}
		}

		// Receiver (RX or Both)
		if req.Direction == "rx" || req.Direction == "both" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				rxBuf := make([]byte, 65535)

				for {
					select {
					case <-ctx.Done():
						return
					default:
					}

					n, _, err := udpConn.ReadFromUDP(rxBuf)
					if err != nil {
						select {
						case <-ctx.Done():
							return
						default:
							time.Sleep(10 * time.Millisecond)
							continue
						}
					}

					if n < HeaderSize {
						continue
					}

					hdr, err := DecodeHeader(rxBuf[:n])
					if err != nil || hdr.SessionHash != sessionHash {
						continue
					}

					nowNano := time.Now().UnixNano()
					localRxBytes.Add(int64(n))
					localRxPackets.Add(1)

					// Jitter calculation
					jitterMu.Lock()
					j := jitterCalc.Update(hdr.TimestampNano, nowNano)
					jitterMu.Unlock()
					curJitterMs.Store(uint64(j * 1000)) // ms * 1000

					// Packet loss tracking
					lastMax := maxSeqRecv.Load()
					if hdr.Sequence > lastMax {
						if lastMax > 0 && hdr.Sequence > lastMax+1 {
							gap := int64(hdr.Sequence - lastMax - 1)
							localLostPackets.Add(gap)
						}
						maxSeqRecv.Store(hdr.Sequence)
					}
				}
			}()
		}
	} else if req.Protocol == "tcp" {
		// TCP Data path
		if req.Direction == "tx" || req.Direction == "both" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				chunk := make([]byte, 64*1024)
				limiter := NewRateLimiter(req.BandwidthLimit)

				for {
					select {
					case <-ctx.Done():
						return
					default:
					}

					if limiter != nil {
						limiter.Wait(len(chunk))
					}

					n, err := controlConn.Write(chunk)
					if err != nil {
						return
					}
					localTxBytes.Add(int64(n))
				}
			}()
		}

		if req.Direction == "rx" || req.Direction == "both" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				buf := make([]byte, 64*1024)
				for {
					select {
					case <-ctx.Done():
						return
					default:
					}

					n, err := controlConn.Read(buf)
					if err != nil {
						return
					}
					localRxBytes.Add(int64(n))
				}
			}()
		}
	}

	// Telemetry sampling loop (every 500ms)
	sampleTicker := time.NewTicker(500 * time.Millisecond)
	defer sampleTicker.Stop()

	var lastTxBytes, lastRxBytes int64
	lastTime := time.Now()
	startTime := lastTime

	timeoutChan := make(<-chan time.Time)
	if req.Duration > 0 {
		timer := time.NewTimer(time.Duration(req.Duration) * time.Second)
		defer timer.Stop()
		timeoutChan = timer.C
	}

	for {
		select {
		case <-ctx.Done():
			// Notify remote server to stop
			stopMsg := SessionHandshake{Type: "stop", SessionID: sessionID}
			sb, _ := json.Marshal(stopMsg)
			_, _ = controlConn.Write(append(sb, '\n'))
			return

		case <-timeoutChan:
			stopMsg := SessionHandshake{Type: "stop", SessionID: sessionID}
			sb, _ := json.Marshal(stopMsg)
			_, _ = controlConn.Write(append(sb, '\n'))
			return

		case now := <-sampleTicker.C:
			intervalSec := now.Sub(lastTime).Seconds()
			if intervalSec <= 0 {
				intervalSec = 0.5
			}
			lastTime = now
			elapsedSec := now.Sub(startTime).Seconds()

			var curTxBytes, curRxBytes int64
			var curLostPackets int64

			if req.Direction == "tx" {
				curTxBytes = localTxBytes.Load()
				curRxBytes = serverRxBytes.Load()
				curLostPackets = serverLostPackets.Load()
			} else if req.Direction == "rx" {
				curRxBytes = localRxBytes.Load()
				curLostPackets = localLostPackets.Load()
			} else { // both
				curTxBytes = localTxBytes.Load()
				curRxBytes = localRxBytes.Load()
				curLostPackets = localLostPackets.Load() + serverLostPackets.Load()
			}

			deltaTx := curTxBytes - lastTxBytes
			deltaRx := curRxBytes - lastRxBytes
			lastTxBytes = curTxBytes
			lastRxBytes = curRxBytes

			txBps := int64(float64(deltaTx*8) / intervalSec)
			rxBps := int64(float64(deltaRx*8) / intervalSec)

			txMbps := float64(txBps) / 1e6
			rxMbps := float64(rxBps) / 1e6

			// Total packets & loss %
			totalExpectedPkts := localRxPackets.Load() + localLostPackets.Load()
			var lossPct float64
			if totalExpectedPkts > 0 {
				lossPct = (float64(curLostPackets) / float64(totalExpectedPkts)) * 100.0
				if lossPct > 100.0 {
					lossPct = 100.0
				}
			}

			jitter := float64(curJitterMs.Load()) / 1000.0

			sample := &LiveSample{
				Timestamp:   now,
				ElapsedSec:  elapsedSec,
				TxBps:       txBps,
				RxBps:       rxBps,
				TxMbps:      txMbps,
				RxMbps:      rxMbps,
				TxPackets:   localTxPackets.Load(),
				RxPackets:   localRxPackets.Load(),
				TxBytes:     curTxBytes,
				RxBytes:     curRxBytes,
				LostPackets: curLostPackets,
				LossPct:     lossPct,
				JitterMs:    jitter,
			}

			c.mu.Lock()
			c.samples = append(c.samples, sample)
			if len(c.samples) > 200 {
				c.samples = c.samples[len(c.samples)-200:]
			}
			c.status.ElapsedSec = elapsedSec
			c.status.CurrentSample = sample
			c.mu.Unlock()

			// Non-blocking publish to subscriber channel
			select {
			case c.sampleCh <- *sample:
			default:
			}
		}
	}
}

func (c *BwTestClient) finalizeSummary() {
	if len(c.samples) == 0 {
		return
	}

	first := c.samples[0]
	last := c.samples[len(c.samples)-1]
	duration := last.ElapsedSec

	var sumTxMbps, sumRxMbps float64
	var peakTxMbps, peakRxMbps float64
	var sumLossPct, sumJitter float64

	for _, s := range c.samples {
		sumTxMbps += s.TxMbps
		sumRxMbps += s.RxMbps
		sumLossPct += s.LossPct
		sumJitter += s.JitterMs
		if s.TxMbps > peakTxMbps {
			peakTxMbps = s.TxMbps
		}
		if s.RxMbps > peakRxMbps {
			peakRxMbps = s.RxMbps
		}
	}

	count := float64(len(c.samples))
	avgTx := sumTxMbps / count
	avgRx := sumRxMbps / count
	avgLoss := sumLossPct / count
	avgJitter := sumJitter / count

	c.status.Summary = &TestSummary{
		DurationSec:      duration,
		StartTime:        c.status.StartTime,
		EndTime:          time.Now(),
		TotalTxBytes:     last.TxBytes,
		TotalRxBytes:     last.RxBytes,
		AvgTxMbps:        avgTx,
		AvgRxMbps:        avgRx,
		PeakTxMbps:       peakTxMbps,
		PeakRxMbps:       peakRxMbps,
		TotalTxPackets:   last.TxPackets,
		TotalRxPackets:   last.RxPackets,
		TotalLostPackets: last.LostPackets,
		AvgLossPct:       avgLoss,
		AvgJitterMs:      avgJitter,
	}
	_ = first
}