package dispatcher

import (
	"bytes"
	"crypto/sha256"
	"errors"
	"sync"
	"time"
)

const MaxPendingMessages = 64

type pendingMessage struct {
	digest  [32]byte
	total   uint16
	parts   map[uint16][]byte
	created time.Time
}

type Reassembler struct {
	mu      sync.Mutex
	pending map[[16]byte]*pendingMessage
	ttl     time.Duration
	now     func() time.Time
}

func NewReassembler(ttl time.Duration) *Reassembler {
	return &Reassembler{
		pending: make(map[[16]byte]*pendingMessage),
		ttl:     ttl,
		now:     time.Now,
	}
}

func (r *Reassembler) Add(frame Frame) ([]byte, bool, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.expireLocked()
	pending, ok := r.pending[frame.ID]
	if !ok {
		if len(r.pending) >= MaxPendingMessages {
			return nil, false, errors.New("too many pending messages")
		}
		pending = &pendingMessage{
			digest:  frame.Digest,
			total:   frame.Total,
			parts:   make(map[uint16][]byte, frame.Total),
			created: r.now(),
		}
		r.pending[frame.ID] = pending
	}
	if pending.total != frame.Total || pending.digest != frame.Digest {
		delete(r.pending, frame.ID)
		return nil, false, errors.New("inconsistent frame metadata")
	}
	if _, duplicate := pending.parts[frame.Index]; !duplicate {
		pending.parts[frame.Index] = bytes.Clone(frame.Data)
	}
	if len(pending.parts) != int(pending.total) {
		return nil, false, nil
	}

	var assembled bytes.Buffer
	for index := uint16(0); index < pending.total; index++ {
		part, exists := pending.parts[index]
		if !exists {
			return nil, false, nil
		}
		if assembled.Len()+len(part) > MaxEnvelopeBytes {
			delete(r.pending, frame.ID)
			return nil, false, errors.New("reassembled envelope is too large")
		}
		assembled.Write(part)
	}
	delete(r.pending, frame.ID)
	result := assembled.Bytes()
	if sha256.Sum256(result) != pending.digest {
		return nil, false, errors.New("reassembled envelope digest mismatch")
	}
	return result, true, nil
}

func (r *Reassembler) expireLocked() {
	cutoff := r.now().Add(-r.ttl)
	for id, pending := range r.pending {
		if pending.created.Before(cutoff) {
			delete(r.pending, id)
		}
	}
}