package checker

// voice.go — predictive voice diagnostics.
//
// runVoiceQualityBurst fires a burst of STUN binding requests through
// an open SOCKS5 UDP relay, then derives packet-loss / jitter /
// percentile-RTT from the replies. A single round-trip says the relay
// accepts UDP; a 30-packet burst tells you whether voice will actually
// hold together.

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"math"
	"net"
	"sort"
	"sync"
	"time"
)

// VoiceQualityResult is the outcome of a UDP burst through a SOCKS5
// relay. All fields are zero on a hard failure (no replies at all).
type VoiceQualityResult struct {
	Sent     int
	Received int
	LossPct  float64 // 0..100
	JitterMS float64 // mean abs of inter-arrival deltas in ms
	P50RTTMS float64 // median round-trip in ms
	P95RTTMS float64 // 95th percentile (informational, not gated)
}

// runVoiceQualityBurst sends `count` STUN binding requests through the
// already-open SOCKS5 UDP relay (relayAddr) to stunHost:stunPort,
// spaced `interval` apart. It listens on udpConn until
// `time.Now() + max(interval, 200ms)` after the last send, then returns
// the aggregate result.
//
// Each outbound datagram has the SOCKS5 UDP header
// (RSV 00 00, FRAG 00, ATYP 01, DST_IPv4(4), DST_PORT(2)) followed by
// a 20-byte STUN binding request. We track each request by its
// transaction ID. Replies are stripped of their 10-byte SOCKS5 UDP
// header before being handed to ParseBindingResponse.
//
// Returns an error only when ctx is cancelled or stunHost can't be
// resolved to IPv4. A 100% loss is NOT an error — the caller decides
// what status to assign; we just report Sent=count, Received=0.
func runVoiceQualityBurst(
	ctx context.Context,
	udpConn net.PacketConn,
	relayAddr *net.UDPAddr,
	stunHost string,
	stunPort uint16,
	count int,
	interval time.Duration,
) (VoiceQualityResult, error) {
	if count <= 0 {
		return VoiceQualityResult{}, errors.New("voice-quality: burst count must be > 0")
	}

	// Resolve stunHost to IPv4.
	ips, err := net.DefaultResolver.LookupIP(ctx, "ip4", stunHost)
	if err != nil {
		return VoiceQualityResult{}, fmt.Errorf("voice-quality: lookup %s: %w", stunHost, err)
	}
	var stunIP4 net.IP
	for _, ip := range ips {
		if v4 := ip.To4(); v4 != nil {
			stunIP4 = v4
			break
		}
	}
	if stunIP4 == nil {
		return VoiceQualityResult{}, fmt.Errorf("voice-quality: no IPv4 for %s", stunHost)
	}

	// Per-tx state: send-time + arrival-time.
	type entry struct {
		sentAt    time.Time
		arrivedAt time.Time
		received  bool
	}
	var (
		mu       sync.Mutex
		entries  = make(map[[12]byte]*entry, count)
		arrivals = make([]time.Time, 0, count) // for jitter (in arrival order)
		rtts     = make([]float64, 0, count)   // milliseconds
	)

	// Reader goroutine: loops on ReadFrom until deadline expires.
	doneRead := make(chan struct{})
	go func() {
		defer close(doneRead)
		buf := make([]byte, 1500)
		for {
			n, _, rerr := udpConn.ReadFrom(buf)
			if rerr != nil {
				// Deadline expired or conn closed — exit.
				return
			}
			if n < 10 {
				continue
			}
			// Validate SOCKS5 UDP wrapper, derive header length.
			if buf[0] != 0x00 || buf[1] != 0x00 || buf[2] != 0x00 {
				continue
			}
			var hdrLen int
			switch buf[3] {
			case 0x01:
				hdrLen = 10
			case 0x04:
				hdrLen = 22
			case 0x03:
				if n < 5 {
					continue
				}
				hdrLen = 4 + 1 + int(buf[4]) + 2
			default:
				continue
			}
			if n < hdrLen+20 {
				continue
			}
			stunReply := buf[hdrLen:n]
			// Pull the transaction ID out of the STUN header so we
			// can look up the matching send-time. ParseBindingResponse
			// rejects mismatched txIDs, so we feed it the *expected*
			// id from the entries map.
			var txID [12]byte
			copy(txID[:], stunReply[8:20])

			now := time.Now()
			mu.Lock()
			ent, ok := entries[txID]
			if !ok || ent.received {
				mu.Unlock()
				continue
			}
			if _, _, perr := ParseBindingResponse(stunReply, txID); perr != nil {
				mu.Unlock()
				continue
			}
			ent.arrivedAt = now
			ent.received = true
			arrivals = append(arrivals, now)
			rtts = append(rtts, float64(now.Sub(ent.sentAt).Microseconds())/1000.0)
			mu.Unlock()
		}
	}()

	// Build base SOCKS5 UDP header (RSV+FRAG+ATYP+IP+PORT). STUN body
	// is per-packet (fresh tx id each).
	hdr := make([]byte, 0, 10)
	hdr = append(hdr, 0x00, 0x00, 0x00, 0x01)
	hdr = append(hdr, stunIP4...)
	var portBuf [2]byte
	binary.BigEndian.PutUint16(portBuf[:], stunPort)
	hdr = append(hdr, portBuf[:]...)

	// Send burst.
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	sent := 0
sendLoop:
	for sent < count {
		// Make a fresh tx id and STUN request.
		txID, terr := NewTransactionID()
		if terr != nil {
			break
		}
		stunReq := EncodeBindingRequest(txID)
		dgram := make([]byte, 0, len(hdr)+len(stunReq))
		dgram = append(dgram, hdr...)
		dgram = append(dgram, stunReq...)

		// Record send-time *before* the write. Note: we register the
		// entry into the map BEFORE Write so the reader can never get a
		// reply for an unknown tx (would happen on a very fast localhost
		// echo).
		now := time.Now()
		mu.Lock()
		entries[txID] = &entry{sentAt: now}
		mu.Unlock()

		if _, werr := udpConn.WriteTo(dgram, relayAddr); werr != nil {
			// Write failure aborts the burst — but we still wait for
			// any in-flight replies. Treat as "sent so far".
			break
		}
		sent++

		if sent >= count {
			break sendLoop
		}

		// Wait for next tick OR ctx cancel.
		select {
		case <-ticker.C:
		case <-ctx.Done():
			break sendLoop
		}
	}

	// Wait window for stragglers — at least 200ms past last send.
	wait := interval
	if wait < 200*time.Millisecond {
		wait = 200 * time.Millisecond
	}
	deadline := time.Now().Add(wait)
	_ = udpConn.SetReadDeadline(deadline)

	// Wait for reader to exit. ctx cancel still races: bound by deadline.
	select {
	case <-doneRead:
	case <-ctx.Done():
		// Force the reader to exit ASAP by setting a past deadline.
		_ = udpConn.SetReadDeadline(time.Unix(0, 1))
		<-doneRead
	}
	// Reset deadline so subsequent users of the conn aren't surprised.
	_ = udpConn.SetReadDeadline(time.Time{})

	// Compute aggregates.
	mu.Lock()
	defer mu.Unlock()

	received := len(rtts)
	res := VoiceQualityResult{
		Sent:     sent,
		Received: received,
	}
	if sent > 0 {
		res.LossPct = float64(sent-received) / float64(sent) * 100.0
	}
	if received >= 2 {
		// Sort arrivals to compute inter-arrival jitter in chronological order.
		// arrivals is already chronological (appended as packets came in).
		var diffs []float64
		for i := 1; i < len(arrivals); i++ {
			d := float64(arrivals[i].Sub(arrivals[i-1]).Microseconds()) / 1000.0
			diffs = append(diffs, d)
		}
		// mean abs of consecutive deltas of inter-arrival diffs.
		if len(diffs) >= 2 {
			var sum float64
			for i := 1; i < len(diffs); i++ {
				sum += math.Abs(diffs[i] - diffs[i-1])
			}
			res.JitterMS = sum / float64(len(diffs)-1)
		} else if len(diffs) == 1 {
			// Only two arrivals — single delta, no second-order jitter.
			res.JitterMS = 0
		}
	}
	if received > 0 {
		// percentile.
		sorted := make([]float64, len(rtts))
		copy(sorted, rtts)
		sort.Float64s(sorted)
		p50idx := len(sorted) / 2
		if p50idx >= len(sorted) {
			p50idx = len(sorted) - 1
		}
		res.P50RTTMS = sorted[p50idx]
		p95idx := int(0.95 * float64(len(sorted)))
		if p95idx >= len(sorted) {
			p95idx = len(sorted) - 1
		}
		res.P95RTTMS = sorted[p95idx]
	}

	return res, nil
}