package memtable

import (
	"bytes"
	"math/rand"
	"sync"
	"sync/atomic"
)

const (
	maxLevel    = 20
	probability = 0.25
)

// SkipList is an ordered in-memory key-value index.
//
// Writers are serialized by mu. Readers never take mu: every next pointer is
// published with atomic.Pointer.Store and read with atomic.Pointer.Load, giving
// release/acquire ordering for fully initialized nodes.
type SkipList struct {
	arena *Arena
	mu    sync.Mutex
	head  *skipNode

	height atomic.Int32
	length atomic.Uint64
}

type skipNode struct {
	next [maxLevel]atomic.Pointer[skipNode]

	key      []byte
	value    []byte
	sequence uint64
	height   int
	pending  bool
}

// NewSkipList creates an empty skip list. The arena is retained for the
// memtable API and future arena-backed nodes; Phase 1 stores nodes on the Go
// heap so atomic.Pointer can provide correct lock-free reads.
func NewSkipList(arena *Arena) *SkipList {
	head := &skipNode{height: maxLevel}
	sl := &SkipList{
		arena: arena,
		head:  head,
	}
	sl.height.Store(1)
	return sl
}

// Put inserts or replaces key with value and sequence.
func (sl *SkipList) Put(key []byte, value []byte, sequence uint64, pending bool) error {
	sl.mu.Lock()
	defer sl.mu.Unlock()

	var prev [maxLevel]*skipNode
	found := sl.findGreaterOrEqual(key, &prev)

	if found != nil && bytes.Equal(found.key, key) {
		sl.replaceLocked(&prev, found, key, value, sequence, pending)
		return nil
	}

	height := randomHeight()
	sl.ensureHeightLocked(height, &prev)
	node := newSkipNode(key, value, sequence, pending, height)

	for level := range height {
		node.next[level].Store(prev[level].next[level].Load())
	}
	for level := range height {
		prev[level].next[level].Store(node)
	}
	sl.length.Add(1)
	return nil
}

// Get returns the latest published value for key. It performs no locking.
func (sl *SkipList) Get(key []byte) (found bool, value []byte, sequence uint64) {
	node := sl.findGreaterOrEqual(key, nil)
	if node == nil || !bytes.Equal(node.key, key) || node.pending {
		return false, nil, 0
	}
	return true, cloneBytes(node.value), node.sequence
}

// Len returns the number of distinct keys in the skip list.
func (sl *SkipList) Len() uint64 {
	return sl.length.Load()
}

// NewIterator creates a forward iterator over published entries.
func (sl *SkipList) NewIterator() *Iterator {
	it := &Iterator{node: sl.head.next[0].Load()}
	it.skipPending()
	return it
}

// Iterator is a lock-free forward iterator over SkipList entries.
type Iterator struct {
	node *skipNode
}

// Valid reports whether the iterator points at an entry.
func (it *Iterator) Valid() bool {
	return it.node != nil
}

// Next advances the iterator to the next published entry.
func (it *Iterator) Next() {
	if it.node == nil {
		return
	}
	it.node = it.node.next[0].Load()
	it.skipPending()
}

// Key returns a copy of the current key.
func (it *Iterator) Key() []byte {
	if it.node == nil {
		return nil
	}
	return cloneBytes(it.node.key)
}

// Value returns a copy of the current value.
func (it *Iterator) Value() []byte {
	if it.node == nil {
		return nil
	}
	return cloneBytes(it.node.value)
}

// Sequence returns the current entry's sequence number.
func (it *Iterator) Sequence() uint64 {
	if it.node == nil {
		return 0
	}
	return it.node.sequence
}

func (it *Iterator) skipPending() {
	for it.node != nil && it.node.pending {
		it.node = it.node.next[0].Load()
	}
}

func (sl *SkipList) replaceLocked(prev *[maxLevel]*skipNode, old *skipNode, key []byte, value []byte, sequence uint64, pending bool) {
	height := randomHeight()
	sl.ensureHeightLocked(height, prev)
	node := newSkipNode(key, value, sequence, pending, height)

	for level := range maxLevel {
		oldNext := old.next[level].Load()
		if level < height {
			node.next[level].Store(oldNext)
			prev[level].next[level].Store(node)
			continue
		}

		if level < old.height {
			prev[level].next[level].Store(oldNext)
		}
	}
}

func (sl *SkipList) ensureHeightLocked(height int, prev *[maxLevel]*skipNode) {
	currentHeight := int(sl.height.Load())
	if height <= currentHeight {
		return
	}

	for level := currentHeight; level < height; level++ {
		prev[level] = sl.head
	}
	sl.height.Store(int32(height))
}

func (sl *SkipList) findGreaterOrEqual(key []byte, prev *[maxLevel]*skipNode) *skipNode {
	x := sl.head
	level := int(sl.height.Load()) - 1
	for level >= 0 {
		next := x.next[level].Load()
		if next != nil && bytes.Compare(next.key, key) < 0 {
			x = next
			continue
		}
		if prev != nil {
			prev[level] = x
		}
		level--
	}
	return x.next[0].Load()
}

func newSkipNode(key []byte, value []byte, sequence uint64, pending bool, height int) *skipNode {
	return &skipNode{
		key:      cloneBytes(key),
		value:    cloneBytes(value),
		sequence: sequence,
		height:   height,
		pending:  pending,
	}
}

func randomHeight() int {
	height := 1
	for height < maxLevel && rand.Float64() < probability {
		height++
	}
	return height
}

func cloneBytes(src []byte) []byte {
	if src == nil {
		return nil
	}
	dst := make([]byte, len(src))
	copy(dst, src)
	return dst
}