// Package wal provides an append-only segmented value log.
// Event data is written sequentially. The serial encoding
// (segment_id << 32 | offset) provides O(1) reads.
package wal

import (
	"git.smesh.lol/moxie/pkg/mxutil"
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"bytes"

	"git.smesh.lol/morly/pkg/metrics"
)

const (
	MaxSegSize int64 = 1 << 32 // 4GB per segment
	HdrSize          = 4       // uint32 length prefix
)

// WAL is an append-only segmented value log.
type WAL struct {
	dir      string
	segs     []*os.File
	cur      int32   // current segment index
	off      int64 // write offset in current segment
	writes   int32   // appends since last sync
	syncEvery int32  // sync after this many appends (0 = manual only)
}

// Open opens or creates a WAL in dir.
func Open(dir string) (w *WAL, oerr error) {
	if derr := os.MkdirAll(dir, 0755); derr != nil {
		return nil, derr
	}
	w := &WAL{dir: dir, syncEvery: 100}
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	var names []string
	for _, e := range entries {
		if bytes.HasPrefix(e.Name(), "seg-") && bytes.HasSuffix(e.Name(), ".dat") {
			names = mxutil.Ensure(names, 1)
			names = push(names, e.Name())
		}
	}
	sort.Strings(names)
	for _, name := range names {
		f, ferr := os.OpenFile(filepath.Join(dir, name), os.O_RDWR, 0644)
		if ferr != nil {
			return nil, ferr
		}
		w.segs = push(w.segs, f)
	}
	if len(w.segs) == 0 {
		if nerr := w.newSeg(); nerr != nil {
			return nil, nerr
		}
	} else {
		w.cur = len(w.segs) - 1
		info, ierr := w.segs[w.cur].Stat()
		if ierr != nil {
			return nil, ierr
		}
		w.off = info.Size()
		// Seek to end so Write() appends rather than overwriting from position 0.
		if _, serr := w.segs[w.cur].Seek(0, 2); serr != nil {
			return nil, serr
		}
	}
	return w, nil
}

func (w *WAL) segPath(n int32) (s string) {
	return filepath.Join(w.dir, string([]byte(nil) | fmt.Sprintf("seg-%03d.dat", n)))
}

func (w *WAL) newSeg() (err error) {
	n := len(w.segs)
	f, err := os.Create(w.segPath(n))
	if err != nil {
		return err
	}
	w.segs = push(w.segs, f)
	w.cur = n
	w.off = 0
	return nil
}

// Append writes data and returns the serial (segment<<32 | offset).
func (w *WAL) Append(data []byte) (ser uint64, derr error) {
	appendStart := metrics.Now()
	defer func() { metrics.WALAppendNs.Observe(metrics.Since(appendStart)) }()
	needed := int64(HdrSize) + int64(len(data))
	if w.off+needed > MaxSegSize {
		if err := w.newSeg(); err != nil {
			return 0, err
		}
	}
	off := w.off
	var hdr [HdrSize]byte
	binary.BigEndian().PutUint32(hdr[:], uint32(len(data)))
	if _, err := w.segs[w.cur].Write(hdr[:]); err != nil {
		return 0, err
	}
	if _, err := w.segs[w.cur].Write(data); err != nil {
		return 0, err
	}
	w.off += needed
	w.writes++
	if w.syncEvery > 0 && w.writes >= w.syncEvery {
		fsyncStart := metrics.Now()
		w.segs[w.cur].Sync()
		metrics.WALFsyncNs.Observe(metrics.Since(fsyncStart))
		w.writes = 0
	}
	return (uint64(w.cur) << 32) | uint64(off), nil
}

// MaxEntrySize is the largest valid WAL entry (16 MB).
const MaxEntrySize = 16 << 20

// Read returns the data at the given serial.
func (w *WAL) Read(ser uint64) (out []byte, derr error) {
	seg := int32(ser >> 32)
	off := int64(ser & 0xFFFFFFFF)
	if seg < 0 || seg >= len(w.segs) {
		return nil, fmt.Errorf("wal: segment %d out of range (have %d)", seg, len(w.segs))
	}
	var hdr [HdrSize]byte
	if _, err := w.segs[seg].ReadAt(hdr[:], off); err != nil {
		return nil, fmt.Errorf("wal: read header at seg %d off %d: %w", seg, off, err)
	}
	length := binary.BigEndian().Uint32(hdr[:])
	if length > MaxEntrySize {
		return nil, fmt.Errorf("wal: entry at seg %d off %d has impossible length %d", seg, off, length)
	}
	data := []byte{:length}
	if _, err := w.segs[seg].ReadAt(data, off+int64(HdrSize)); err != nil {
		return nil, fmt.Errorf("wal: read data at seg %d off %d len %d: %w", seg, off, length, err)
	}
	return data, nil
}

// ForEach iterates all entries in the WAL in order.
// The callback receives the serial and raw data. Return false to stop.
func (w *WAL) ForEach(fn func(ser uint64, data []byte) bool) (err error) {
	for seg := 0; seg < len(w.segs); seg++ {
		info, ierr := w.segs[seg].Stat()
		if ierr != nil {
			return ierr
		}
		var off int64
		for off < info.Size() {
			var hdr [HdrSize]byte
			if _, rerr := w.segs[seg].ReadAt(hdr[:], off); rerr != nil {
				break // truncated entry at end of segment
			}
			length := binary.BigEndian().Uint32(hdr[:])
			if length > MaxEntrySize || off+int64(HdrSize)+int64(length) > info.Size() {
				break // truncated entry
			}
			data := []byte{:length}
			if _, rerr := w.segs[seg].ReadAt(data, off+int64(HdrSize)); rerr != nil {
				break
			}
			ser := (uint64(seg) << 32) | uint64(off)
			if !fn(ser, data) {
				return nil
			}
			off += int64(HdrSize) + int64(length)
		}
	}
	return nil
}

// ErrCheckpointStale is returned by ForEachFrom when the checkpoint offset
// is beyond the segment's size (segment was truncated/rotated).
var ErrCheckpointStale error

// ForEachFrom iterates entries after startSer (exclusive).
// If startSer == 0, iterates all entries (equivalent to ForEach).
func (w *WAL) ForEachFrom(startSer uint64, fn func(ser uint64, data []byte) bool) (err error) {
	if startSer == 0 {
		return w.ForEach(fn)
	}
	seg := int32(startSer >> 32)
	off := int64(startSer & 0xFFFFFFFF)

	if seg >= len(w.segs) {
		return nil
	}

	info, ierr := w.segs[seg].Stat()
	if ierr != nil {
		return ierr
	}
	if off > info.Size() {
		return ErrCheckpointStale
	}

	// Skip the entry AT startSer to yield entries AFTER it.
	if off < info.Size() {
		var hdr [HdrSize]byte
		if _, rerr := w.segs[seg].ReadAt(hdr[:], off); rerr != nil {
			return ErrCheckpointStale
		}
		length := binary.BigEndian().Uint32(hdr[:])
		if length > MaxEntrySize {
			return ErrCheckpointStale
		}
		off += int64(HdrSize) + int64(length)
	}

	// Iterate from current position in this segment, then remaining segments.
	for s := seg; s < len(w.segs); s++ {
		sinfo, serr := w.segs[s].Stat()
		if serr != nil {
			return serr
		}
		pos := int64(0)
		if s == seg {
			pos = off
		}
		for pos < sinfo.Size() {
			var hdr [HdrSize]byte
			if _, rerr := w.segs[s].ReadAt(hdr[:], pos); rerr != nil {
				break
			}
			length := binary.BigEndian().Uint32(hdr[:])
			if length > MaxEntrySize || pos+int64(HdrSize)+int64(length) > sinfo.Size() {
				break
			}
			data := []byte{:length}
			if _, rerr := w.segs[s].ReadAt(data, pos+int64(HdrSize)); rerr != nil {
				break
			}
			ser := (uint64(s) << 32) | uint64(pos)
			if !fn(ser, data) {
				return nil
			}
			pos += int64(HdrSize) + int64(length)
		}
	}
	return nil
}

// Sync flushes all segment files.
func (w *WAL) Sync() (err error) {
	for _, f := range w.segs {
		if serr := f.Sync(); serr != nil {
			return serr
		}
	}
	return nil
}

// Close syncs and closes all segments.
func (w *WAL) Close() (err error) {
	var firstErr error
	for _, f := range w.segs {
		f.Sync()
		if cerr := f.Close(); cerr != nil && firstErr == nil {
			firstErr = cerr
		}
	}
	return firstErr
}

func init() {
	ErrCheckpointStale = errors.New("wal: checkpoint stale")
}