package bwtest

import (
	"encoding/binary"
	"errors"
	"hash/fnv"
	"math"
	"time"
)

const (
	MagicBytes uint32 = 0x42575453 // "BWTS"
	ProtoVer   uint8  = 1

	PacketTypeData   uint8 = 1
	PacketTypeReport uint8 = 2
	PacketTypeStop   uint8 = 3

	HeaderSize = 36
)

var (
	ErrInvalidMagic = errors.New("invalid bandwidth test packet magic")
	ErrShortPacket  = errors.New("packet shorter than minimum header size")
)

// PacketHeader represents the 36-byte binary header on every UDP test datagram
type PacketHeader struct {
	Magic         uint32
	Version       uint8
	PacketType    uint8
	StreamID      uint8
	Flags         uint8
	SessionHash   uint64
	Sequence      uint64
	TimestampNano int64
	PayloadLength uint32
}

// HashSessionID generates a fast 64-bit FNV-1a hash of the session ID string
func HashSessionID(sessionID string) uint64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(sessionID))
	return h.Sum64()
}

// EncodeHeader serializes the packet header into the provided byte slice (must be >= 36 bytes)
func EncodeHeader(b []byte, hdr PacketHeader) {
	binary.BigEndian.PutUint32(b[0:4], MagicBytes)
	b[4] = ProtoVer
	b[5] = hdr.PacketType
	b[6] = hdr.StreamID
	b[7] = hdr.Flags
	binary.BigEndian.PutUint64(b[8:16], hdr.SessionHash)
	binary.BigEndian.PutUint64(b[16:24], hdr.Sequence)
	binary.BigEndian.PutUint64(b[24:32], uint64(hdr.TimestampNano))
	binary.BigEndian.PutUint32(b[32:36], hdr.PayloadLength)
}

// DecodeHeader deserializes the binary packet header from byte slice
func DecodeHeader(b []byte) (PacketHeader, error) {
	if len(b) < HeaderSize {
		return PacketHeader{}, ErrShortPacket
	}
	magic := binary.BigEndian.Uint32(b[0:4])
	if magic != MagicBytes {
		return PacketHeader{}, ErrInvalidMagic
	}
	return PacketHeader{
		Magic:         magic,
		Version:       b[4],
		PacketType:    b[5],
		StreamID:      b[6],
		Flags:         b[7],
		SessionHash:   binary.BigEndian.Uint64(b[8:16]),
		Sequence:      binary.BigEndian.Uint64(b[16:24]),
		TimestampNano: int64(binary.BigEndian.Uint64(b[24:32])),
		PayloadLength: binary.BigEndian.Uint32(b[32:36]),
	}, nil
}

// JitterCalculator computes RFC 3550 statistical interarrival jitter
type JitterCalculator struct {
	hasPrev     bool
	prevTransit int64
	jitterFloat float64 // in nanoseconds
}

// Update incorporates a newly received packet's sender timestamp and arrival time
func (jc *JitterCalculator) Update(senderTsNano, arrivalTsNano int64) float64 {
	transit := arrivalTsNano - senderTsNano
	if !jc.hasPrev {
		jc.prevTransit = transit
		jc.hasPrev = true
		return 0
	}
	diff := math.Abs(float64(transit - jc.prevTransit))
	jc.prevTransit = transit
	// RFC 3550: J(i) = J(i-1) + (|D(i-1,i)| - J(i-1))/16
	jc.jitterFloat += (diff - jc.jitterFloat) / 16.0
	return jc.jitterFloat / 1e6 // Return in milliseconds
}

// RateLimiter implements a high-precision token bucket for pacing outbound test traffic
type RateLimiter struct {
	bps           int64
	tokens        float64
	maxTokens     float64
	lastCheckTime time.Time
}

// NewRateLimiter initializes a rate limiter for target bits per second (0 = no limit)
func NewRateLimiter(targetBps int64) *RateLimiter {
	if targetBps <= 0 {
		return nil
	}
	bytesPerSec := float64(targetBps) / 8.0
	return &RateLimiter{
		bps:           targetBps,
		tokens:        bytesPerSec * 0.05, // 50ms burst allowance
		maxTokens:     bytesPerSec * 0.10, // 100ms max burst
		lastCheckTime: time.Now(),
	}
}

// Wait paces the sender goroutine to ensure bandwidth limit is respected
func (rl *RateLimiter) Wait(bytes int) {
	if rl == nil || rl.bps <= 0 {
		return
	}
	neededTokens := float64(bytes)
	bytesPerSec := float64(rl.bps) / 8.0

	for {
		now := time.Now()
		elapsed := now.Sub(rl.lastCheckTime).Seconds()
		rl.lastCheckTime = now

		rl.tokens += elapsed * bytesPerSec
		if rl.tokens > rl.maxTokens {
			rl.tokens = rl.maxTokens
		}

		if rl.tokens >= neededTokens {
			rl.tokens -= neededTokens
			return
		}

		deficit := neededTokens - rl.tokens
		sleepSec := deficit / bytesPerSec
		if sleepSec > 0.05 {
			sleepSec = 0.05
		}
		time.Sleep(time.Duration(sleepSec * float64(time.Second)))
	}
}