package pcap

import (
	"cmp"
	"errors"
	"io"
	"slices"
	"sync"
	"sync/atomic"
)

// A reorderEntry is 16 bytes, which bounds the memory needed per packet.
type reorderEntry struct {
	ts  uint64 // seconds << 32 | fraction (µs or ns, consistent within a file)
	pos uint64 // record offset << lenBits | record length (including the 16 byte record header)
}

const (
	lenBits   = 28 // fits MaxRecordLen + 16
	maxOffset = 1<<(64-lenBits) - 1
)

func (e reorderEntry) off() int64 { return int64(e.pos >> lenBits) }
func (e reorderEntry) len() int   { return int(e.pos & (1<<lenBits - 1)) }

// Tuning for WriteTo: packets are copied in batches of up to batchBytes,
// using parallel reads since random reads from SSDs are much faster with
// several requests in flight. The next batch is read while the previous one
// is written.
const (
	batchPackets = 16384
	readers      = 16
)

var batchBytes = 8 << 20 // variable for tests

// Reorder sorts the packets of a pcap file by timestamp. NewReorder reads
// and indexes the file, so all parse errors are reported before WriteTo
// produces any output.
type Reorder struct {
	r          io.ReaderAt
	header     []byte
	entries    []reorderEntry
	end        int64 // offset just after the last complete record
	outOfOrder int
	truncated  bool
}

// NewReorder indexes the pcap file in r (of the given size) and sorts the
// index by timestamp. The sort is stable: packets with equal timestamps keep
// their original order. It needs 16 bytes of memory per packet.
func NewReorder(r io.ReaderAt, size int64) (*Reorder, error) {
	pr, err := NewReader(io.NewSectionReader(r, 0, size))
	if err != nil {
		return nil, err
	}
	ro := &Reorder{r: r, header: slices.Clone(pr.Header())}
	off := int64(len(ro.header))
	var latest uint64
	for {
		rec, err := pr.Next()
		if err == io.EOF {
			break
		}
		if errors.Is(err, ErrTruncated) {
			ro.truncated = true
			break
		}
		if err != nil {
			return nil, err
		}
		if off > maxOffset {
			return nil, errors.New("file too large to reorder")
		}
		ts := uint64(pr.order.Uint32(rec[0:4]))<<32 | uint64(pr.order.Uint32(rec[4:8]))
		if ts < latest {
			ro.outOfOrder++
		}
		latest = max(latest, ts)
		ro.entries = append(ro.entries, reorderEntry{ts: ts, pos: uint64(off)<<lenBits | uint64(len(rec))})
		off += int64(len(rec))

		if len(ro.entries) == 1024 {
			// Size the index from the average packet size so far, to avoid
			// repeatedly growing (and copying) a large slice.
			avg := (off - int64(len(ro.header))) / 1024
			est := size / avg
			ro.entries = slices.Grow(ro.entries, int(est+est/16))
		}
	}
	if len(ro.entries) == 0 {
		return nil, errors.New("the pcap file contains no packets")
	}
	ro.end = off
	ro.entries = slices.Clip(ro.entries)
	if ro.outOfOrder > 0 {
		slices.SortStableFunc(ro.entries, func(a, b reorderEntry) int { return cmp.Compare(a.ts, b.ts) })
	}
	return ro, nil
}

// Packets returns the number of packets.
func (ro *Reorder) Packets() int { return len(ro.entries) }

// OutOfOrder returns the number of packets with a timestamp earlier than a
// packet before them in the input.
func (ro *Reorder) OutOfOrder() int { return ro.outOfOrder }

// Truncated reports whether the input ended with an incomplete packet
// record, which is left out of the output.
func (ro *Reorder) Truncated() bool { return ro.truncated }

type batch struct {
	buf []byte
	err error
}

// WriteTo writes the file header followed by the packets in timestamp order.
func (ro *Reorder) WriteTo(w io.Writer) (int64, error) {
	if ro.outOfOrder == 0 {
		return io.Copy(w, io.NewSectionReader(ro.r, 0, ro.end))
	}

	n, err := w.Write(ro.header)
	total := int64(n)
	if err != nil {
		return total, err
	}

	// A producer goroutine reads batches into two alternating buffers while
	// this goroutine writes them out.
	full := make(chan batch)
	free := make(chan []byte, 2)
	free <- nil
	free <- nil
	done := make(chan struct{})
	var wg sync.WaitGroup
	defer func() {
		close(done)
		wg.Wait()
	}()
	wg.Go(func() {
		defer close(full)
		for i := 0; i < len(ro.entries); {
			var buf []byte
			select {
			case buf = <-free:
			case <-done:
				return
			}
			var err error
			i, buf, err = ro.readBatch(i, buf)
			select {
			case full <- batch{buf, err}:
			case <-done:
				return
			}
			if err != nil {
				return
			}
		}
	})

	for b := range full {
		if b.err != nil {
			return total, b.err
		}
		n, err := w.Write(b.buf)
		total += int64(n)
		if err != nil {
			return total, err
		}
		free <- b.buf
	}
	return total, nil
}

// readBatch reads the records of entries[i:j] into buf, choosing j so that
// the batch stays within batchBytes and batchPackets (but holds at least one
// record). It returns j and the filled buffer.
func (ro *Reorder) readBatch(i int, buf []byte) (int, []byte, error) {
	j, size := i, 0
	for j < len(ro.entries) && j-i < batchPackets && (j == i || size+ro.entries[j].len() <= batchBytes) {
		size += ro.entries[j].len()
		j++
	}
	entries := ro.entries[i:j]
	dst := make([]int, len(entries)+1)
	for k, e := range entries {
		dst[k+1] = dst[k] + e.len()
	}
	buf = slices.Grow(buf[:0], size)[:size]

	var (
		next atomic.Int64
		wg   sync.WaitGroup
		errs [readers]error
	)
	for r := range min(readers, len(entries)) {
		wg.Go(func() {
			for {
				k := int(next.Add(1) - 1)
				if k >= len(entries) {
					return
				}
				e := entries[k]
				if _, err := ro.r.ReadAt(buf[dst[k]:dst[k+1]], e.off()); err != nil {
					errs[r] = err
					return
				}
			}
		})
	}
	wg.Wait()
	return j, buf, errors.Join(errs[:]...)
}