package xdns

import (
	"sync"
	"time"
)

const (
	fragTTL   = 4 * time.Second
	fragSize  = 4096
	fragCount = 4096
)

type FragKey struct {
	clientID ClientID
	fragID   byte
}

type FragEntry struct {
	data     [][]byte
	size     int
	len      int
	total    byte
	deadline time.Time
}

type FragManager struct {
	m  map[FragKey]*FragEntry
	ch chan struct{}
	mu sync.Mutex
}

func NewFragManager() *FragManager {
	m := &FragManager{
		m:  make(map[FragKey]*FragEntry),
		ch: make(chan struct{}),
	}
	go m.gc()
	return m
}

func (m *FragManager) closed() bool {
	select {
	case <-m.ch:
		return true
	default:
		return false
	}
}

func (m *FragManager) remove(key FragKey) {
	delete(m.m, key)
}

func (m *FragManager) tryRemove() {
	if len(m.m) < fragCount {
		return
	}
	var key FragKey
	var entry *FragEntry
	first := true
	for k, e := range m.m {
		if first || e.deadline.Before(entry.deadline) {
			key = k
			entry = e
			first = false
		}
	}
	m.remove(key)
}

func (m *FragManager) gc() {
	ticker := time.NewTicker(fragTTL / 2)
	defer ticker.Stop()
	for {
		select {
		case <-m.ch:
			return
		case now := <-ticker.C:
			m.mu.Lock()
			for k, e := range m.m {
				if now.After(e.deadline) {
					m.remove(k)
				}
			}
			m.mu.Unlock()
		}
	}
}

func (m *FragManager) Feed(out []byte, key FragKey, fragIdx, fragN byte, data []byte) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.closed() {
		return 0
	}

	if fragN < 2 {
		return 0
	}

	now := time.Now()
	entry := m.m[key]
	if entry == nil || now.After(entry.deadline) {
		if entry == nil {
			m.tryRemove()
		} else {
			m.remove(key)
		}
		entry = &FragEntry{
			data:     make([][]byte, fragN),
			total:    fragN,
			deadline: now.Add(fragTTL),
		}
		m.m[key] = entry
	}

	if fragN != entry.total {
		return 0
	}
	if fragIdx >= entry.total {
		return 0
	}
	if entry.data[fragIdx] != nil {
		return 0
	}
	if entry.size+len(data) > fragSize {
		return 0
	}

	cp := make([]byte, len(data))
	copy(cp, data)

	entry.data[fragIdx] = cp
	entry.size += len(data)
	entry.len++
	entry.deadline = now.Add(fragTTL)

	if entry.len < int(entry.total) {
		return 0
	}

	out = out[:0]
	for i := range entry.data {
		out = append(out, entry.data[i]...)
	}
	m.remove(key)
	return len(out)
}

func (m *FragManager) Close() {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.closed() {
		return
	}
	close(m.ch)
	for k := range m.m {
		delete(m.m, k)
	}
}