package store

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// sortedFile is a fixed-width, sorted flat file with binary search.
//
// Records reach it through three sources:
//
//	memory  - appended by Put since the last flush;
//	runs    - immutable, individually sorted files (path.buf.NNNNNN);
//	main    - the .dat file, produced by merging all runs into it.
//
// Each flush creates a new run instead of appending to one shared sidecar.
// That distinction matters: two independently sorted batches appended to one
// file are not themselves sorted, so binary search over that file silently
// misses records. Separate run files keep every source sorted, and a merge
// turns all of them into the main file. This is the LSM shape pogreb-style
// engines use, and the layout morly's index files follow.
type sortedFile struct {
	path   string
	recLen int
	cmpLen int

	f     *os.File
	count int64 // records in the main file

	buf    []byte
	bufN   int
	sorted bool

	runs    []runRef
	runData [][]byte
	runSeq  int
	del     map[string]struct{}

	mergeThresh int64
	bulk        bool
	serOff      int
	removed     func(uint64) bool
}

type runRef struct {
	path  string
	count int64
}

// flushAt bounds the memory buffer between explicit flushes.
const flushAt = 8 << 20

func openSorted(path string, recLen, cmpLen int) (*sortedFile, error) {
	os.Remove(path + ".tmp")
	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o644)
	if err != nil {
		return nil, err
	}
	info, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, err
	}
	size := info.Size()
	count := size / int64(recLen)
	if size%int64(recLen) != 0 {
		// A torn tail cannot be part of a complete record; the index can
		// always be rebuilt from the WAL.
		if err := f.Truncate(count * int64(recLen)); err != nil {
			f.Close()
			return nil, err
		}
	}

	s := &sortedFile{
		path:   path,
		recLen: recLen,
		cmpLen: cmpLen,
		f:      f,
		count:  count,
		sorted: true,
		del:    make(map[string]struct{}),
	}
	if err := s.discoverRuns(); err != nil {
		f.Close()
		return nil, err
	}

	thresh := size / 10
	if thresh < 1<<20 {
		thresh = 1 << 20
	}
	s.mergeThresh = thresh
	return s, nil
}

// discoverRuns finds the immutable run files, oldest first.
func (s *sortedFile) discoverRuns() error {
	names, err := filepath.Glob(s.path + ".buf*")
	if err != nil {
		return err
	}
	sort.Strings(names)
	for _, name := range names {
		info, err := os.Stat(name)
		if err != nil {
			continue
		}
		size := info.Size()
		count := size / int64(s.recLen)
		if size == 0 {
			os.Remove(name)
			continue
		}
		if size%int64(s.recLen) != 0 {
			// A partially written run is the tail of a crash; drop the
			// incomplete record so the run stays searchable.
			if err := os.Truncate(name, count*int64(s.recLen)); err != nil {
				return err
			}
		}
		s.runs = append(s.runs, runRef{path: name, count: count})
		if suffix := strings.TrimPrefix(name, s.path+".buf."); suffix != name {
			var n int
			if _, err := fmt.Sscanf(suffix, "%d", &n); err == nil && n >= s.runSeq {
				s.runSeq = n + 1
			}
		} else if s.runSeq == 0 {
			s.runSeq = 1
		}
	}
	s.runData = make([][]byte, len(s.runs))
	return nil
}

// Put appends a record. The record is copied.
func (s *sortedFile) Put(rec []byte) error {
	if len(rec) != s.recLen {
		panic("store: bad record length")
	}
	// A later Put cancels an earlier Delete for the same key.
	if len(s.del) > 0 {
		delete(s.del, string(rec[:s.cmpLen]))
	}
	s.buf = append(s.buf, rec...)
	s.bufN++
	s.sorted = false
	if len(s.buf) >= flushAt {
		return s.QuickFlush(false)
	}
	return nil
}

// Delete marks a key (its first cmpLen bytes) as removed.
func (s *sortedFile) Delete(key []byte) {
	k := make([]byte, s.cmpLen)
	copy(k, key)
	s.del[string(k)] = struct{}{}
}

func (s *sortedFile) isDeleted(rec []byte) bool {
	if len(s.del) > 0 {
		if _, ok := s.del[string(rec[:s.cmpLen])]; ok {
			return true
		}
	}
	if s.removed != nil && s.serOff >= 0 && len(rec) >= s.serOff+SerialLen {
		return s.removed(GetSerial(rec[s.serOff : s.serOff+SerialLen]))
	}
	return false
}

func (s *sortedFile) Dirty() bool { return s.bufN > 0 || len(s.del) > 0 || len(s.runs) > 0 }

func (s *sortedFile) Count() int64 { return s.count + int64(s.bufN) + s.runCount() }

func (s *sortedFile) runCount() int64 {
	var n int64
	for _, r := range s.runs {
		n += r.count
	}
	return n
}

func (s *sortedFile) runBytes() int64 {
	var n int64
	for _, r := range s.runs {
		n += r.count * int64(s.recLen)
	}
	return n
}

type recSorter struct {
	data   []byte
	recLen int
	cmpLen int
	tmp    []byte
}

func (r *recSorter) Len() int { return len(r.data) / r.recLen }
func (r *recSorter) Less(i, j int) bool {
	a := r.data[i*r.recLen : i*r.recLen+r.cmpLen]
	b := r.data[j*r.recLen : j*r.recLen+r.cmpLen]
	return bytes.Compare(a, b) < 0
}
func (r *recSorter) Swap(i, j int) {
	a := r.data[i*r.recLen : (i+1)*r.recLen]
	b := r.data[j*r.recLen : (j+1)*r.recLen]
	copy(r.tmp, a)
	copy(a, b)
	copy(b, r.tmp)
}

func (s *sortedFile) ensureSorted() {
	if s.sorted || s.bufN <= 1 {
		s.sorted = true
		return
	}
	sort.Sort(&recSorter{s.buf, s.recLen, s.cmpLen, make([]byte, s.recLen)})
	s.sorted = true
}

// QuickFlush writes the memory buffer as a new immutable run. sync controls
// fsync, which is what makes the checkpoint safe to advance.
func (s *sortedFile) QuickFlush(sync bool) error {
	if s.bufN == 0 {
		return nil
	}
	s.ensureSorted()
	name := fmt.Sprintf("%s.buf.%06d", s.path, s.runSeq)
	s.runSeq++
	tmp := name + ".tmp"
	f, err := os.OpenFile(tmp, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o644)
	if err != nil {
		return err
	}
	if _, err := f.Write(s.buf[:s.bufN*s.recLen]); err != nil {
		f.Close()
		return err
	}
	if sync {
		if err := f.Sync(); err != nil {
			f.Close()
			return err
		}
	}
	if err := f.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmp, name); err != nil {
		return err
	}
	count := int64(s.bufN)
	s.runs = append(s.runs, runRef{path: name, count: count})
	if s.bulk {
		// Bulk imports create many runs and defer every merge to Flush.
		// Keeping their bytes resident would make memory grow with the whole
		// import; loadRuns reads them back when a scan actually needs them.
		s.runData = append(s.runData, nil)
	} else {
		data := make([]byte, len(s.buf[:s.bufN*s.recLen]))
		copy(data, s.buf[:s.bufN*s.recLen])
		s.runData = append(s.runData, data)
	}
	s.buf = s.buf[:0]
	s.bufN = 0
	s.sorted = true

	if s.bulk {
		// Bulk imports write one sorted run per buffer and defer the single
		// merge to Flush. Merging every run as the main file grows turns a
		// linear import into repeated rewrites of the whole index.
		return nil
	}
	if s.runBytes() >= s.mergeThresh {
		return s.Flush()
	}
	return nil
}

// loadRuns reads every run into memory once. Runs are bounded by the merge
// threshold, so this is bounded memory, and the cache avoids re-reading the
// same run for every point lookup in a query.
func (s *sortedFile) loadRuns() error {
	if len(s.runData) == len(s.runs) {
		all := true
		for _, d := range s.runData {
			if d == nil {
				all = false
				break
			}
		}
		if all {
			return nil
		}
	}
	if len(s.runData) != len(s.runs) {
		s.runData = make([][]byte, len(s.runs))
	}
	for i, r := range s.runs {
		if s.runData[i] != nil {
			continue
		}
		data := make([]byte, r.count*int64(s.recLen))
		f, err := os.Open(r.path)
		if err != nil {
			return err
		}
		_, err = f.ReadAt(data, 0)
		f.Close()
		if err != nil && err != io.EOF {
			return err
		}
		s.runData[i] = data
	}
	return nil
}

func (s *sortedFile) cmpBound(rec, bound []byte) int {
	n := s.cmpLen
	if len(bound) < n {
		n = len(bound)
	}
	return bytes.Compare(rec[:n], bound[:n])
}

func (s *sortedFile) readAt(dst []byte, idx int64) error {
	_, err := s.f.ReadAt(dst, idx*int64(s.recLen))
	return err
}

// lowerBound returns the first main-file record >= key.
func (s *sortedFile) lowerBound(key []byte) int64 {
	lo, hi := int64(0), s.count
	rec := make([]byte, s.recLen)
	for lo < hi {
		mid := lo + (hi-lo)/2
		if err := s.readAt(rec, mid); err != nil {
			return s.count
		}
		if s.cmpBound(rec, key) < 0 {
			lo = mid + 1
		} else {
			hi = mid
		}
	}
	return lo
}

// Get returns the first record whose cmpLen prefix equals key.
func (s *sortedFile) Get(key []byte) ([]byte, bool, error) {
	// Memory is newest, then runs newest-to-oldest, then the main file.
	if s.bufN > 0 {
		s.ensureSorted()
		lo, hi := 0, s.bufN
		for lo < hi {
			mid := lo + (hi-lo)/2
			rec := s.buf[mid*s.recLen:]
			if s.cmpBound(rec, key) < 0 {
				lo = mid + 1
			} else {
				hi = mid
			}
		}
		if lo < s.bufN {
			rec := s.buf[lo*s.recLen : (lo+1)*s.recLen]
			if s.cmpBound(rec, key) == 0 && !s.isDeleted(rec) {
				out := make([]byte, s.recLen)
				copy(out, rec)
				return out, true, nil
			}
		}
	}
	if err := s.loadRuns(); err != nil {
		return nil, false, err
	}
	for i := len(s.runs) - 1; i >= 0; i-- {
		data := s.runData[i]
		n := s.runs[i].count
		lo, hi := int64(0), n
		for lo < hi {
			mid := lo + (hi-lo)/2
			rec := data[mid*int64(s.recLen):]
			if s.cmpBound(rec, key) < 0 {
				lo = mid + 1
			} else {
				hi = mid
			}
		}
		if lo < n {
			rec := data[lo*int64(s.recLen):]
			if s.cmpBound(rec, key) == 0 && !s.isDeleted(rec) {
				out := make([]byte, s.recLen)
				copy(out, rec[:s.recLen])
				return out, true, nil
			}
		}
	}
	idx := s.lowerBound(key)
	if idx < s.count {
		rec := make([]byte, s.recLen)
		if err := s.readAt(rec, idx); err != nil {
			return nil, false, err
		}
		if s.cmpBound(rec, key) == 0 {
			if s.isDeleted(rec) {
				return nil, false, nil
			}
			return rec, true, nil
		}
	}
	return nil, false, nil
}

// Scan emits every record in [start, end) in key order. A nil end means the
// end of the index. fn returning false stops the scan.
func (s *sortedFile) Scan(start, end []byte, fn func(rec []byte) bool) error {
	if s.bufN > 0 {
		s.ensureSorted()
	}
	if err := s.loadRuns(); err != nil {
		return err
	}

	var recs [][]byte
	inRange := func(rec []byte) bool {
		if start != nil && s.cmpBound(rec, start) < 0 {
			return false
		}
		if end != nil && s.cmpBound(rec, end) >= 0 {
			return false
		}
		return true
	}
	add := func(rec []byte) {
		if !inRange(rec) || s.isDeleted(rec) {
			return
		}
		out := make([]byte, s.recLen)
		copy(out, rec[:s.recLen])
		recs = append(recs, out)
	}

	from := int64(0)
	if start != nil {
		from = s.lowerBound(start)
	}
	chunkRecs := 256
	chunk := make([]byte, s.recLen*chunkRecs)
	pastEnd := func(rec []byte) bool { return end != nil && s.cmpBound(rec, end) >= 0 }

	// The main file is already sorted. With no runs or buffered records there
	// is nothing to merge, so stream records straight to fn. The general path
	// below collects and sorts the entire range; for a broad scan that can be
	// millions of records even when the caller stops after a few results.
	if len(s.runs) == 0 && s.bufN == 0 {
		for i := from; i < s.count; {
			n := int64(chunkRecs)
			if s.count-i < n {
				n = s.count - i
			}
			if _, err := s.f.ReadAt(chunk[:int(n)*s.recLen], i*int64(s.recLen)); err != nil && err != io.EOF {
				return err
			}
			for j := int64(0); j < n; j++ {
				rec := chunk[j*int64(s.recLen) : (j+1)*int64(s.recLen)]
				if pastEnd(rec) {
					return nil
				}
				if s.isDeleted(rec) {
					continue
				}
				if !fn(rec) {
					return nil
				}
			}
			i += n
		}
		return nil
	}

	for i := from; i < s.count; {
		n := int64(chunkRecs)
		if s.count-i < n {
			n = s.count - i
		}
		if _, err := s.f.ReadAt(chunk[:int(n)*s.recLen], i*int64(s.recLen)); err != nil && err != io.EOF {
			return err
		}
		stop := false
		for j := int64(0); j < n; j++ {
			rec := chunk[j*int64(s.recLen) : (j+1)*int64(s.recLen)]
			if pastEnd(rec) {
				stop = true
				break
			}
			add(rec)
		}
		if stop {
			break
		}
		i += n
	}
	for i := range s.runs {
		data := s.runData[i]
		for j := int64(0); j < s.runs[i].count; j++ {
			rec := data[j*int64(s.recLen):]
			if pastEnd(rec) {
				break
			}
			add(rec)
		}
	}
	for i := 0; i < s.bufN; i++ {
		rec := s.buf[i*s.recLen : (i+1)*s.recLen]
		if pastEnd(rec) {
			break
		}
		add(rec)
	}

	if len(recs) == 0 {
		return nil
	}
	sort.Slice(recs, func(i, j int) bool {
		c := bytes.Compare(recs[i][:s.cmpLen], recs[j][:s.cmpLen])
		if c != 0 {
			return c < 0
		}
		return bytes.Compare(recs[i], recs[j]) < 0
	})
	var prev []byte
	for _, rec := range recs {
		if prev != nil && bytes.Equal(prev, rec) {
			continue
		}
		prev = rec
		if !fn(rec) {
			return nil
		}
	}
	return nil
}

// Flush merges the main file and all runs into a new main file, dropping
// deleted and duplicate records.
func (s *sortedFile) Flush() error {
	if s.bufN > 0 {
		if err := s.QuickFlush(false); err != nil {
			return err
		}
	}
	if len(s.runs) == 0 && len(s.del) == 0 {
		return nil
	}
	if err := s.loadRuns(); err != nil {
		return err
	}

	// Combine the runs into one sorted buffer. Their total size is bounded by
	// mergeThresh, and the main file is streamed rather than loaded.
	total := 0
	for _, d := range s.runData {
		total += len(d)
	}
	merged := make([]byte, 0, total)
	for _, d := range s.runData {
		merged = append(merged, d...)
	}
	if len(merged) > 0 {
		sort.Sort(&recSorter{merged, s.recLen, s.cmpLen, make([]byte, s.recLen)})
	}

	tmpPath := s.path + ".tmp"
	tmp, err := os.OpenFile(tmpPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o644)
	if err != nil {
		return err
	}
	w := bufio.NewWriterSize(tmp, 1<<16)

	const chunkRecs = 1024
	mainBuf := make([]byte, chunkRecs*s.recLen)
	mainPos := int64(0)
	mainRec := func() ([]byte, error) {
		if mainPos >= s.count {
			return nil, nil
		}
		idx := mainPos % chunkRecs
		if idx == 0 {
			n := int64(chunkRecs)
			if s.count-mainPos < n {
				n = s.count - mainPos
			}
			if _, err := s.f.ReadAt(mainBuf[:int(n)*s.recLen], mainPos*int64(s.recLen)); err != nil && err != io.EOF {
				return nil, err
			}
		}
		rec := mainBuf[idx*int64(s.recLen) : (idx+1)*int64(s.recLen)]
		mainPos++
		return rec, nil
	}
	mergePos := int64(0)
	mergeRec := func() []byte {
		if mergePos >= int64(len(merged)) {
			return nil
		}
		rec := merged[mergePos : mergePos+int64(s.recLen)]
		mergePos += int64(s.recLen)
		return rec
	}

	prev := make([]byte, 0, s.recLen)
	emit := func(rec []byte) error {
		if s.isDeleted(rec) {
			return nil
		}
		if len(prev) == s.recLen && bytes.Equal(prev, rec) {
			return nil
		}
		prev = append(prev[:0], rec...)
		_, err := w.Write(rec)
		return err
	}

	mr, err := mainRec()
	if err != nil {
		tmp.Close()
		return err
	}
	ur := mergeRec()
	for mr != nil || ur != nil {
		var rec []byte
		switch {
		case mr == nil:
			rec = ur
			ur = mergeRec()
		case ur == nil:
			rec = mr
			mr, err = mainRec()
			if err != nil {
				tmp.Close()
				return err
			}
		default:
			c := bytes.Compare(mr[:s.cmpLen], ur[:s.cmpLen])
			switch {
			case c < 0:
				rec = mr
				mr, err = mainRec()
				if err != nil {
					tmp.Close()
					return err
				}
			case c > 0:
				rec = ur
				ur = mergeRec()
			default:
				if bytes.Compare(mr, ur) <= 0 {
					rec = mr
					mr, err = mainRec()
					if err != nil {
						tmp.Close()
						return err
					}
				} else {
					rec = ur
					ur = mergeRec()
				}
			}
		}
		if err := emit(rec); err != nil {
			tmp.Close()
			return err
		}
	}

	if err := w.Flush(); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := s.f.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmpPath, s.path); err != nil {
		return err
	}
	f, err := os.OpenFile(s.path, os.O_RDWR, 0o644)
	if err != nil {
		return err
	}
	s.f = f
	info, err := f.Stat()
	if err != nil {
		return err
	}
	s.count = info.Size() / int64(s.recLen)

	for _, r := range s.runs {
		os.Remove(r.path)
	}
	s.runs = nil
	s.runData = nil
	s.del = make(map[string]struct{})
	return nil
}

func (s *sortedFile) Close() error {
	if err := s.Flush(); err != nil {
		return err
	}
	if err := s.f.Sync(); err != nil {
		return err
	}
	if err := s.f.Close(); err != nil {
		return err
	}
	for _, r := range s.runs {
		os.Remove(r.path)
	}
	return nil
}

// Clear removes every record from the main file and every run.
func (s *sortedFile) Clear() error {
	s.buf = s.buf[:0]
	s.bufN = 0
	s.sorted = true
	s.del = make(map[string]struct{})
	for _, r := range s.runs {
		os.Remove(r.path)
	}
	s.runs = nil
	s.runData = nil
	if err := s.f.Truncate(0); err != nil {
		return err
	}
	s.count = 0
	return nil
}