package arena_test

import (
	"errors"
	"runtime"
	"testing"
	"unsafe"

	"git.brut.systems/judah/xx/arena"
	"git.brut.systems/judah/xx/mem"
)

func BenchmarkAlloc_New_Small(b *testing.B) {
	var last *int
	for i := range b.N {
		v := new(int)
		*v = i
		last = v

		if i%1000 == 0 {
			last = nil
		}
	}

	runtime.KeepAlive(last)
}

func BenchmarkAlloc_Closure_Small(b *testing.B) {
	alloc := arena.Ring[int](16)

	var last *int
	for i := range b.N {
		v := arena.New[int](alloc)
		*v = i
		last = v

		if i%1000 == 0 {
			arena.Reset(alloc)
		}
	}

	runtime.KeepAlive(last)
}

func BenchmarkAlloc_Interface_Small(b *testing.B) {
	alloc := NewLinear(16 * mem.Kilobyte)

	var last *int
	for i := range b.N {
		v := New[int](alloc)
		*v = i
		last = v

		if i%1000 == 0 {
			Reset(alloc)
		}
	}

	runtime.KeepAlive(last)
}

type large struct{ a, b, c, d, e, f, g, h, i int }

func BenchmarkAlloc_New_Large(b *testing.B) {
	var last *large
	for i := range b.N {
		v := new(large)
		v.e = i
		last = v

		if i%1000 == 0 {
			last = nil
		}
	}

	runtime.KeepAlive(last)
}

func BenchmarkAlloc_Closure_Large(b *testing.B) {
	alloc := arena.Linear(128 * mem.Kilobyte)

	var last *large
	for i := range b.N {
		v := arena.New[large](alloc)
		v.e = i
		last = v

		if i%1000 == 0 {
			arena.Reset(alloc)
		}
	}

	runtime.KeepAlive(last)
}

func BenchmarkAlloc_Interface_Large(b *testing.B) {
	alloc := NewLinear(128 * mem.Kilobyte)

	var last *large
	for i := range b.N {
		v := New[large](alloc)
		v.e = i
		last = v

		if i%1000 == 0 {
			Reset(alloc)
		}
	}

	runtime.KeepAlive(last)
}

func BenchmarkAlloc_Closure_HotPath(b *testing.B) {
	alloc := arena.Linear(256)

	var (
		lastlarge *large
		lastsmall *int
	)
	for i := range b.N {
		if i%2 == 0 {
			lastsmall = arena.New[int](alloc)
		} else {
			lastlarge = arena.New[large](alloc)
		}

		arena.Reset(alloc)
	}

	runtime.KeepAlive(lastlarge)
	runtime.KeepAlive(lastsmall)
}

func BenchmarkAlloc_Interface_HotPath(b *testing.B) {
	alloc := NewLinear(256)

	var (
		lastlarge *large
		lastsmall *int
	)
	for i := range b.N {
		if i%2 == 0 {
			lastsmall = New[int](alloc)
		} else {
			lastlarge = New[large](alloc)
		}

		Reset(alloc)
	}

	runtime.KeepAlive(lastlarge)
	runtime.KeepAlive(lastsmall)
}

func BenchmarkAlloc_Closure_Wrapped(b *testing.B) {
	alloc := arena.Pinned(arena.Pinned(arena.Pinned(arena.Linear(256))))

	var (
		lastlarge *large
		lastsmall *int
	)
	for i := range b.N {
		if i%2 == 0 {
			lastsmall = arena.New[int](alloc)
		} else {
			lastlarge = arena.New[large](alloc)
		}

		arena.Reset(alloc)
	}

	runtime.KeepAlive(lastlarge)
	runtime.KeepAlive(lastsmall)
}

func BenchmarkAlloc_Interface_Wrapped(b *testing.B) {
	alloc := NewPinned(NewPinned(NewPinned(NewLinear(256))))

	var (
		lastlarge *large
		lastsmall *int
	)
	for i := range b.N {
		if i%2 == 0 {
			lastsmall = New[int](alloc)
		} else {
			lastlarge = New[large](alloc)
		}

		Reset(alloc)
	}

	runtime.KeepAlive(lastlarge)
	runtime.KeepAlive(lastsmall)
}

type Allocator interface {
	Proc(a arena.Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error)
}

func New[T any](a Allocator) *T {
	ptr, err := a.Proc(arena.ACTION_ALLOC, mem.Sizeof[T](), mem.Alignof[T](), nil)
	if err != nil {
		panic(err)
	}

	return (*T)(ptr)
}

func Reset(a Allocator) {
	if _, err := a.Proc(arena.ACTION_RESET, 0, 0, nil); err != nil {
		panic(err)
	}
}

type Linear struct {
	data    []byte
	maxsize uintptr
	offset  uintptr
}

func NewLinear(maxsize uintptr) *Linear {
	return &Linear{
		data:    make([]byte, maxsize),
		maxsize: maxsize,
	}
}

func (l *Linear) Proc(a arena.Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
	switch a {
	case arena.ACTION_ALLOC:
		aligned := mem.AlignForward(size, align)
		if l.offset+aligned > l.maxsize {
			return nil, errors.New("linear: out of memory")
		}

		ptr := &l.data[l.offset]
		l.offset += aligned
		return unsafe.Pointer(ptr), nil

	case arena.ACTION_RESET:
		clear(l.data)
		l.offset = 0

	case arena.ACTION_SAVE:
		*watermark = l.offset
	case arena.ACTION_RESTORE:
		l.offset = *watermark

	default:
		panic("unimplemented action: " + a.String())
	}

	return nil, nil
}

type Pinned struct {
	arena  Allocator
	pinner runtime.Pinner
}

func NewPinned(arena Allocator) *Pinned {
	return &Pinned{arena: arena}
}

func (p *Pinned) Proc(a arena.Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
	ptr, err := p.arena.Proc(a, size, align, watermark)
	if err != nil {
		return ptr, err
	}

	if a == arena.ACTION_RESET {
		p.pinner.Unpin()
	} else {
		p.pinner.Pin(ptr)
	}

	return ptr, err
}