package udp

import (
	"sync"
	"time"
)

// RetransmitEntry tracks a packet awaiting acknowledgment
type RetransmitEntry struct {
	Packet    *ProtocolPacket // The packet to retransmit
	Sequence  uint16          // Packet sequence number
	Timestamp time.Time       // When packet was last sent
	Attempts  int             // Number of transmission attempts
}

// RetransmitQueue manages reliable packet delivery with exponential backoff
type RetransmitQueue struct {
	entries     map[uint16]*RetransmitEntry // Pending packets by sequence
	mutex       sync.RWMutex                // Thread-safe access
	baseTimeout time.Duration               // Base retransmission timeout
	maxAttempts int                         // Maximum retry attempts
	maxTimeout  time.Duration               // Maximum timeout cap
}

// NewRetransmitQueue creates a queue with default settings
func NewRetransmitQueue() *RetransmitQueue {
	return &RetransmitQueue{
		entries:     make(map[uint16]*RetransmitEntry),
		baseTimeout: 500 * time.Millisecond,
		maxAttempts: 5,
		maxTimeout:  5 * time.Second,
	}
}

// NewRetransmitQueueWithConfig creates a queue with custom settings
func NewRetransmitQueueWithConfig(baseTimeout, maxTimeout time.Duration, maxAttempts int) *RetransmitQueue {
	return &RetransmitQueue{
		entries:     make(map[uint16]*RetransmitEntry),
		baseTimeout: baseTimeout,
		maxAttempts: maxAttempts,
		maxTimeout:  maxTimeout,
	}
}

// Add queues a packet for potential retransmission
func (rq *RetransmitQueue) Add(packet *ProtocolPacket, sequence uint16) {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()

	rq.entries[sequence] = &RetransmitEntry{
		Packet:    packet,
		Sequence:  sequence,
		Timestamp: time.Now(),
		Attempts:  1,
	}
}

// Acknowledge removes a packet from the retransmit queue
func (rq *RetransmitQueue) Acknowledge(sequence uint16) bool {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()

	_, existed := rq.entries[sequence]
	delete(rq.entries, sequence)
	return existed
}

// GetExpired returns packets that need retransmission
func (rq *RetransmitQueue) GetExpired() []*RetransmitEntry {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()

	now := time.Now()
	var expired []*RetransmitEntry

	for seq, entry := range rq.entries {
		timeout := rq.calculateTimeout(entry.Attempts)

		if now.Sub(entry.Timestamp) > timeout {
			if entry.Attempts >= rq.maxAttempts {
				// Give up after max attempts
				delete(rq.entries, seq)
			} else {
				// Schedule for retransmission
				entry.Attempts++
				entry.Timestamp = now
				expired = append(expired, entry)
			}
		}
	}

	return expired
}

// calculateTimeout computes timeout with exponential backoff
func (rq *RetransmitQueue) calculateTimeout(attempts int) time.Duration {
	timeout := rq.baseTimeout * time.Duration(attempts*attempts) // Quadratic backoff
	if timeout > rq.maxTimeout {
		timeout = rq.maxTimeout
	}
	return timeout
}

// Clear removes all pending packets
func (rq *RetransmitQueue) Clear() {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()
	rq.entries = make(map[uint16]*RetransmitEntry)
}

// Size returns the number of pending packets
func (rq *RetransmitQueue) Size() int {
	rq.mutex.RLock()
	defer rq.mutex.RUnlock()
	return len(rq.entries)
}

// GetPendingSequences returns all sequence numbers awaiting acknowledgment
func (rq *RetransmitQueue) GetPendingSequences() []uint16 {
	rq.mutex.RLock()
	defer rq.mutex.RUnlock()

	sequences := make([]uint16, 0, len(rq.entries))
	for seq := range rq.entries {
		sequences = append(sequences, seq)
	}
	return sequences
}

// IsEmpty returns true if no packets are pending
func (rq *RetransmitQueue) IsEmpty() bool {
	rq.mutex.RLock()
	defer rq.mutex.RUnlock()
	return len(rq.entries) == 0
}

// SetBaseTimeout updates the base retransmission timeout
func (rq *RetransmitQueue) SetBaseTimeout(timeout time.Duration) {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()
	rq.baseTimeout = timeout
}

// SetMaxAttempts updates the maximum retry attempts
func (rq *RetransmitQueue) SetMaxAttempts(attempts int) {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()
	rq.maxAttempts = attempts
}

// SetMaxTimeout updates the maximum timeout cap
func (rq *RetransmitQueue) SetMaxTimeout(timeout time.Duration) {
	rq.mutex.Lock()
	defer rq.mutex.Unlock()
	rq.maxTimeout = timeout
}

// GetStats returns retransmission statistics
func (rq *RetransmitQueue) GetStats() RetransmitStats {
	rq.mutex.RLock()
	defer rq.mutex.RUnlock()

	stats := RetransmitStats{
		PendingCount: len(rq.entries),
		BaseTimeout:  rq.baseTimeout,
		MaxAttempts:  rq.maxAttempts,
		MaxTimeout:   rq.maxTimeout,
	}

	// Calculate attempt distribution
	for _, entry := range rq.entries {
		if entry.Attempts == 1 {
			stats.FirstAttempts++
		} else {
			stats.Retransmissions++
		}
	}

	return stats
}

// RetransmitStats contains retransmission queue statistics
type RetransmitStats struct {
	PendingCount    int           // Total pending packets
	FirstAttempts   int           // Packets on first attempt
	Retransmissions int           // Packets being retransmitted
	BaseTimeout     time.Duration // Base timeout setting
	MaxAttempts     int           // Maximum attempts setting
	MaxTimeout      time.Duration // Maximum timeout setting
}