package stack

import (
	"jinx/pkg/lang/vm/code"
	"jinx/pkg/lang/vm/mem"
	"jinx/pkg/lang/vm/value"
)

type Stack interface {
	Push(value value.Value)
	Pop() (value.Value, error)

	Get(index int) (value.Value, error)
	Set(index int, value value.Value) error
	Local(offset int) (value.Value, error)
	Top() (value.Value, error)

	Len() int
	IsEmpty() bool

	PutRootCall(basePos code.Pos) error
	PushCall(newPos, returnPos code.Pos, env mem.Ptr) error
	PopCall() (code.Pos, error)

	CurrentCallEnv() mem.Ptr
	CallDepth() int
	ReachedBaseOfCall() bool
	LocalToStackIndex(local int) int
}

type stackImpl struct {
	data  []value.Value
	calls []callFrame
}

func New() Stack {
	data := make([]value.Value, 0, 64)
	calls := make([]callFrame, 0, 8)

	return &stackImpl{
		data:  data,
		calls: calls,
	}
}

func (stack *stackImpl) Push(value value.Value) {
	stack.data = append(stack.data, value)
}

func (stack *stackImpl) Pop() (value.Value, error) {
	if stack.IsEmpty() || stack.ReachedBaseOfCall() {
		return value.Value{}, ErrStackUnderflow
	}

	v, err := stack.Top()
	if err != nil {
		return value.Value{}, err
	}

	stack.data = stack.data[:stack.Len()-1]
	return v, nil
}

func (stack *stackImpl) Get(index int) (value.Value, error) {
	if index < 0 || index >= stack.Len() {
		return value.Value{}, ErrStackIndexOutOfBounds{Index: index, Len: stack.Len()}
	}

	return stack.data[index], nil
}

func (stack *stackImpl) Set(index int, value value.Value) error {
	if index < 0 || index >= stack.Len() {
		return ErrStackIndexOutOfBounds{Index: index, Len: stack.Len()}
	}

	stack.data[index] = value
	return nil
}

func (stack *stackImpl) Local(offset int) (value.Value, error) {
	if stack.ReachedBaseOfCall() {
		return value.Value{}, ErrStackUnderflow
	}

	return stack.Get(stack.LocalToStackIndex(offset))
}

func (stack *stackImpl) Top() (value.Value, error) {
	if stack.IsEmpty() || stack.ReachedBaseOfCall() {
		return value.Value{}, ErrStackUnderflow
	}

	return stack.data[stack.Len()-1], nil
}

func (stack *stackImpl) Len() int {
	return len(stack.data)
}

func (stack *stackImpl) IsEmpty() bool {
	return len(stack.data) == 0
}

func (stack *stackImpl) PutRootCall(basePos code.Pos) error {
	frame := callFrame{
		pos:       basePos,
		returnPos: basePos,
		base:      0,
		env:       mem.NullPtr,
	}

	if stack.CallDepth() == 0 {
		stack.calls = append(stack.calls, frame)

		return nil
	}

	if stack.CallDepth() != 1 {
		return ErrNotAtRootCallFrame
	}

	stack.calls[0] = frame

	return nil
}

func (stack *stackImpl) PushCall(newPos, returnPos code.Pos, env mem.Ptr) error {
	if stack.CallDepth() == 1000 {
		return ErrReachedMaxCallDepth
	}

	stack.calls = append(stack.calls, callFrame{
		pos:       newPos,
		returnPos: returnPos,
		base:      stack.Len(),
		env:       env,
	})

	return nil
}

func (stack *stackImpl) PopCall() (code.Pos, error) {
	if stack.CallDepth() == 0 {
		return code.Pos{}, ErrReachedRootCallFrame
	}

	call := stack.topCall()
	stack.calls = stack.calls[:stack.CallDepth()-1]

	return call.returnPos, nil
}

func (stack *stackImpl) CurrentCallEnv() mem.Ptr {
	if stack.CallDepth() == 0 || stack.topCall().env.IsNull() {
		return mem.NullPtr
	}

	return stack.topCall().env
}

func (stack *stackImpl) CallDepth() int {
	return len(stack.calls)
}

func (stack *stackImpl) ReachedBaseOfCall() bool {
	return stack.topCall().base == stack.Len()
}

func (stack *stackImpl) LocalToStackIndex(local int) int {
	if stack.CallDepth() == 0 {
		return local
	}
	return stack.topCall().base + local
}

func (stack *stackImpl) topCall() *callFrame {
	return &stack.calls[len(stack.calls)-1]
}

type callFrame struct {
	pos       code.Pos // Beginning of the called function.
	returnPos code.Pos // Where to return to after the called function returns.
	base      int      // Base of the local variables on the data stack.
	env       mem.Ptr  // Environment of the called function.
}