package mem_test

import (
	"jinx/pkg/lang/vm/mem"
	"jinx/pkg/lang/vm/value"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestAllocation(t *testing.T) {
	m := mem.New()

	one, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)

	two, err := m.Allocate(mem.CellKindArray)
	assert.NoError(t, err)

	assert.Equal(t, one, mem.Ptr(1))
	assert.Equal(t, two, mem.Ptr(2))
}

func TestSpecificAllocation(t *testing.T) {
	m := mem.New()

	assert.NoError(t, m.AllocateAt(mem.Ptr(123), mem.CellKindString))
	assert.Error(t, m.AllocateAt(mem.NullPtr, mem.CellKindString))
	assert.Error(t, m.AllocateAt(mem.Ptr(124), mem.CellKindEmpty))
	assert.Error(t, m.AllocateAt(mem.Ptr(123), mem.CellKindArray))
}

func TestGetSet(t *testing.T) {
	m := mem.New()

	one, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(one, value.StringCell("one")))

	two, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(two, value.StringCell("two")))

	oneCell, err := m.Get(one)
	assert.NoError(t, err)
	assert.Equal(t, oneCell.(value.StringCell).Get(), "one")

	twoCell, err := m.Get(two)
	assert.NoError(t, err)
	assert.Equal(t, twoCell.(value.StringCell).Get(), "two")

	assert.NoError(t, m.Set(one, value.StringCell("one, but changed")))

	oneCell, err = m.Get(one)
	assert.NoError(t, err)
	assert.Equal(t, oneCell.(value.StringCell).Get(), "one, but changed")
}

func TestNullIsForbidden(t *testing.T) {
	m := mem.New()

	_, err := m.Get(mem.NullPtr)
	assert.Error(t, err)

	assert.Error(t, m.Set(mem.NullPtr, value.StringCell("one")))

	assert.Error(t, m.Retain(mem.NullPtr))
	assert.Error(t, m.Release(mem.NullPtr))
	assert.Equal(t, 0, m.RefCount(mem.NullPtr))
}

func TestCantGetNorSetUnallocatedCell(t *testing.T) {
	m := mem.New()

	_, err := m.Get(mem.Ptr(123))
	assert.Error(t, err)

	err = m.Set(mem.Ptr(123), value.StringCell("one"))
	assert.Error(t, err)
}

func TestCantAllocateMarkerCells(t *testing.T) {
	m := mem.New()

	_, err := m.Allocate(mem.CellKindEmpty)
	assert.Error(t, err)

	_, err = m.Allocate(mem.CellKindForbidden)
	assert.Error(t, err)
}

func TestCantAssignWrongCellKind(t *testing.T) {
	m := mem.New()

	one, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.Error(t, m.Set(one, value.EnvCell(value.NewEnv())))
}

func TestIsAndKind(t *testing.T) {
	m := mem.New()

	one, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)

	assert.True(t, m.Is(one, mem.CellKindString))
	assert.False(t, m.Is(one, mem.CellKindArray))

	assert.Equal(t, mem.CellKindString, m.Kind(one))
}

func TestIsAndKindAlwaysAgree(t *testing.T) {
	m := mem.New()

	_, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	_, err = m.Allocate(mem.CellKindArray)
	assert.NoError(t, err)
	_, err = m.Allocate(mem.CellKindOutlet)
	assert.NoError(t, err)
	_, err = m.Allocate(mem.CellKindEnv)
	assert.NoError(t, err)

	for i := 0; i < 100; i++ {
		ptr := mem.Ptr(i)
		assert.True(t, m.Is(ptr, m.Kind(ptr)))
	}
}

func TestRetainRelease(t *testing.T) {
	m := mem.New()

	ptr, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(ptr, value.StringCell("an allocated string")))

	assert.True(t, m.Is(ptr, mem.CellKindString))
	assert.Equal(t, 1, m.RefCount(ptr))

	assert.NoError(t, m.Retain(ptr))
	assert.NoError(t, m.Retain(ptr))
	assert.NoError(t, m.Retain(ptr))

	assert.True(t, m.Is(ptr, mem.CellKindString))
	assert.Equal(t, 4, m.RefCount(ptr))

	assert.NoError(t, m.Release(ptr))
	assert.NoError(t, m.Release(ptr))
	assert.NoError(t, m.Release(ptr))

	assert.True(t, m.Is(ptr, mem.CellKindString))
	cell, err := m.Get(ptr)
	assert.NoError(t, err)
	assert.Equal(t, "an allocated string", cell.(value.StringCell).Get())
	assert.Equal(t, 1, m.RefCount(ptr))

	assert.NoError(t, m.Release(ptr))

	assert.Equal(t, 0, m.RefCount(ptr))
	assert.True(t, m.Is(ptr, mem.CellKindEmpty))
	_, err = m.Get(ptr)
	assert.Error(t, err)
}

func TestCantReleaseTooMuch(t *testing.T) {
	m := mem.New()

	ptr, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(ptr, value.StringCell("an allocated string")))

	assert.NoError(t, m.Retain(ptr))
	assert.Equal(t, 2, m.RefCount(ptr))

	assert.NoError(t, m.Release(ptr))
	assert.NoError(t, m.Release(ptr))
	assert.Error(t, m.Release(ptr))
}

func TestCantRetainNorReleaseUnallocated(t *testing.T) {
	m := mem.New()

	assert.Error(t, m.Retain(mem.Ptr(123)))
	assert.Error(t, m.Release(mem.Ptr(123)))
}

func TestChainDrop(t *testing.T) {
	m := mem.New()

	arr := []value.Value{
		unwrap(t, func() (value.Value, error) { return value.NewString(m, "one") }),
		unwrap(t, func() (value.Value, error) { return value.NewString(m, "two") }),

		unwrap(t, func() (value.Value, error) {
			return value.NewArray(m, []value.Value{
				unwrap(t, func() (value.Value, error) { return value.NewString(m, "three") }),
			})
		}),
	}

	val, err := value.NewArray(m, arr)
	assert.NoError(t, err)

	assert.Equal(t, mem.CellKindString, m.Kind(mem.Ptr(1)))
	assert.Equal(t, mem.CellKindString, m.Kind(mem.Ptr(2)))
	assert.Equal(t, mem.CellKindString, m.Kind(mem.Ptr(3)))
	assert.Equal(t, mem.CellKindArray, m.Kind(mem.Ptr(4)))
	assert.Equal(t, mem.CellKindArray, m.Kind(mem.Ptr(5)))
	assert.Equal(t, mem.CellKindEmpty, m.Kind(mem.Ptr(6)))

	assert.NoError(t, val.Drop(m))

	for _, ptr := range []mem.Ptr{1, 2, 3, 4, 5, 6} {
		assert.Equal(t, mem.CellKindEmpty, m.Kind(ptr))
	}
}

func TestFreeCellsGetReused(t *testing.T) {
	m := mem.New()

	one, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(one, value.StringCell("old one")))

	two, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(two, value.StringCell("old two")))

	assert.NoError(t, m.Release(one))

	newOne, err := m.Allocate(mem.CellKindString)
	assert.NoError(t, err)
	assert.NoError(t, m.Set(newOne, value.StringCell("new one")))

	assert.Equal(t, mem.CellKindString, m.Kind(newOne))
	newCell, err := m.Get(newOne)
	assert.NoError(t, err)
	assert.Equal(t, "new one", newCell.(value.StringCell).Get())
	assert.Equal(t, 1, m.RefCount(newOne))
	assert.Equal(t, one, newOne)
}

func unwrap[X any](t *testing.T, f func() (X, error)) X {
	x, err := f()
	assert.NoError(t, err)
	return x
}