package main

import (
	"sync"
	"time"
)

// fragmentSet holds the received fragments for one (counter, type) pair.
type fragmentSet struct {
	total     int
	received  int
	fragments [][]byte // indexed by fragmentIdx
	lastSeen  time.Time
}

// Reassembler reassembles fragmented UDP payloads.
// Key: uint64(counter)<<8 | uint64(pktType)
type Reassembler struct {
	mu      sync.Mutex
	sets    map[uint64]*fragmentSet
	expiry  time.Duration // how long to keep incomplete sets
}

// NewReassembler creates a Reassembler that expires stale fragment sets after ttl.
func NewReassembler(ttl time.Duration) *Reassembler {
	r := &Reassembler{
		sets:   make(map[uint64]*fragmentSet),
		expiry: ttl,
	}
	go r.gcLoop()
	return r
}

func reassemblyKey(counter uint32, pktType uint8) uint64 {
	return uint64(counter)<<8 | uint64(pktType)
}

// AddFragment registers one fragment. Returns the reassembled payload and true
// when all fragments for this (counter, type) have been received, otherwise
// returns nil, false.
func (r *Reassembler) AddFragment(h PacketHeader, payload []byte) ([]byte, bool) {
	key := reassemblyKey(h.Counter, h.Type)
	total := int(h.TotalFragments)
	idx := int(h.FragmentIdx)

	// Fast path: single-fragment packet.
	if total == 1 && idx == 0 {
		return payload, true
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	fs, ok := r.sets[key]
	if !ok {
		fs = &fragmentSet{
			total:     total,
			fragments: make([][]byte, total),
		}
		r.sets[key] = fs
	}

	if idx >= len(fs.fragments) {
		// Stale or corrupt – discard.
		return nil, false
	}
	if fs.fragments[idx] == nil {
		fs.fragments[idx] = payload
		fs.received++
	}
	fs.lastSeen = time.Now()

	if fs.received < fs.total {
		return nil, false
	}

	// All fragments received – concatenate in order.
	delete(r.sets, key)

	total_len := 0
	for _, f := range fs.fragments {
		total_len += len(f)
	}
	out := make([]byte, 0, total_len)
	for _, f := range fs.fragments {
		out = append(out, f...)
	}
	return out, true
}

// gcLoop periodically removes fragment sets that have been incomplete for too long.
func (r *Reassembler) gcLoop() {
	ticker := time.NewTicker(r.expiry / 2)
	defer ticker.Stop()
	for range ticker.C {
		r.mu.Lock()
		now := time.Now()
		for k, fs := range r.sets {
			if now.Sub(fs.lastSeen) > r.expiry {
				delete(r.sets, k)
			}
		}
		r.mu.Unlock()
	}
}