package vm

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

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

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

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

	return Stack{
		data:  data,
		calls: calls,
	}
}

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

func (stack *Stack) 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 *Stack) 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 *Stack) 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 *Stack) Local(offset int) (value.Value, error) {
	if stack.ReachedBaseOfCall() {
		return value.Value{}, ErrStackUnderflow
	}

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

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

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

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

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

func (stack *Stack) 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 *Stack) 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 *Stack) CurrentCallEnv() mem.Ptr {
	if stack.CallDepth() == 0 || stack.TopCall().env.IsNull() {
		return mem.NullPtr
	}

	return stack.TopCall().env
}

func (stack *Stack) ShiftTopCallBase(by int) error {
	call := stack.TopCall()
	newBase := call.base - by

	if newBase < 0 {
		return ErrCallBaseCantBeNegative
	}

	call.base = newBase
	return nil
}

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

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

func (stack *Stack) ReachedBaseOfCall() bool {
	return stack.TopCall().base == stack.Len()
}

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

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.
}