package stack

import (
	"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

	PushCall(newPc, returnPc int, env mem.Ptr) error
	PopCall() (int, 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)

	calls = append(calls, callFrame{
		pc:       0,
		returnPc: 0,
		base:     0,
	})

	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) PushCall(newPc, returnPc int, env mem.Ptr) error {
	if stack.CallDepth() == 1000 {
		return ErrReachedMaxCallDepth
	}

	stack.calls = append(stack.calls, callFrame{
		pc:       newPc,
		returnPc: returnPc,
		base:     stack.Len(),
		env:      env,
	})

	return nil
}

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

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

	return call.returnPc, 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 {
	pc       int     // Beginning of the called function.
	returnPc int     // 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.
}