package vm

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

func (vm *VM) execPushInt(x int64) {
	vm.stack.Push(value.NewInt(x))
}

func (vm *VM) execPushFloat(x float64) {
	vm.stack.Push(value.NewFloat(x))
}

func (vm *VM) execPushString(str string) error {
	val, err := value.NewString(vm.memory, str)
	if err != nil {
		return err
	}

	vm.stack.Push(val)
	return nil
}

func (vm *VM) execPushBool(b bool) {
	vm.stack.Push(value.NewBool(b))
}

func (vm *VM) execPushNull() {
	vm.stack.Push(value.NewNull())
}

func (vm *VM) execPushArray() error {
	val, err := value.NewArray(vm.memory, []value.Value{})
	if err != nil {
		return err
	}

	vm.stack.Push(val)
	return nil
}

func (vm *VM) execPushFunction(pc int) {
	// TODO: Make push ops into functions, where the argCount can be passed.
	vm.stack.Push(value.NewFunction(pc, 0))
}

func (vm *VM) execPushObject() error {
	obj, err := value.NewObject(vm.memory, mem.NullPtr)
	if err != nil {
		return err
	}

	vm.stack.Push(obj)
	return nil
}

func (vm *VM) execGetLocal(offset int) error {
	local, err := vm.stack.Local(int(offset))
	if err != nil {
		return err
	}

	v := local.Clone(vm.memory)

	vm.stack.Push(v)
	return nil
}

func (vm *VM) execSetLocal(offset int) error {
	new, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	local, err := vm.stack.Local(int(offset))
	if err != nil {
		return err
	}

	local.Drop(vm.memory)

	stackIndex := vm.stack.LocalToStackIndex(offset)
	return vm.stack.Set(stackIndex, new)
}

func (vm *VM) execPushType(name string) error {
	ref, err := value.NewType(vm.memory, name)
	if err != nil {
		return err
	}

	vm.stack.Push(ref)
	return nil
}

func (vm *VM) execGetMember(name string) error {
	parent, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	if parent.Type() == value.ObjectType {
		obj := parent.Data().(value.ObjectData)
		ptr := obj.Ptr()

		cell, err := vm.getMemCell(ptr, mem.CellKindObject, false)
		if err != nil {
			return err
		}

		objCell := cell.(value.ObjectCell)
		member, ok := objCell.Get()[name]
		if ok {
			member = member.Clone(vm.memory)

			vm.stack.Push(member)
			parent.Drop(vm.memory)
			return nil
		}
	}

	if parent.Outlet().IsNull() {
		outletPtr, err := vm.memory.Allocate(mem.CellKindOutlet)
		if err != nil {
			return err
		}
		parent = parent.WithOutlet(outletPtr)
	}

	if parent.Type() == value.TypeRefType {
		ref := parent.Data().(value.TypeRefData)
		ptr := ref.TypeRef()

		cell, err := vm.getMemCell(ptr, mem.CellKindType, false)
		if err != nil {
			return err
		}

		typ := cell.(value.TypeCell).Get()

		methodData, ok := typ.GetMethod(name)

		if !methodData.Env().IsNull() {
			panic("methods with environments not implemented yet")
		}

		if ok {
			newEnv := value.NewEnv()
			newEnv.Add(0, parent.Outlet())
			envPtr, err := newEnv.Allocate(vm.memory)
			if err != nil {
				return err
			}

			method := value.NewFunction(0, 0).WithData(methodData.WithEnv(envPtr))
			// method = method.Clone(vm.memory) will only be necessary when we support methods with environments.

			vm.stack.Push(method)
			parent.Drop(vm.memory)
			return nil
		}
	}

	typeCell, err := vm.getMemCell(parent.TypePtr(), mem.CellKindType, false)
	if err != nil {
		return err
	}

	t := typeCell.(value.TypeCell).Get()

	member, ok := t.GetMethod(name)
	if !ok {
		vm.stack.Push(value.NewNull())
		return nil
	}

	var envPtr mem.Ptr
	if member.Env().IsNull() {
		newEnv := value.NewEnv()
		newEnv.Add(0, parent.Outlet()) // stackIndex can be 0, because the parent will be dropped when the function returns.

		envPtr, err = newEnv.Allocate(vm.memory)
		if err != nil {
			return err
		}
	} else {
		// Clone the environment, with the parent at env 0.
		oldEnvCell, err := vm.getMemCell(member.Env(), mem.CellKindEnv, false)
		if err != nil {
			return err
		}

		oldEnv := oldEnvCell.(value.EnvCell).Get()
		newEnv, err := oldEnv.Clone(vm.memory)
		if err != nil {
			return err
		}

		newEnv.Prepend(0, parent.Outlet())

		envPtr, err = newEnv.Allocate(vm.memory)
		if err != nil {
			return err
		}
	}

	member = member.WithEnv(envPtr)

	val := value.NewFunction(0, 0).WithData(member)
	vm.stack.Push(val)

	parent.Drop(vm.memory)

	return nil
}

func (vm *VM) execSetMember(name string) error {
	v, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	parent, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	if parent.Type() != value.ObjectType {
		return ErrInvalidOperandType{
			Op: code.OpSetMember,
			X:  parent.Type(),
		}
	}

	ptr := parent.Data().(value.ObjectData).Ptr()

	cell, err := vm.getMemCell(ptr, mem.CellKindObject, false)
	if err != nil {
		return err
	}

	objCell := cell.(value.ObjectCell)
	obj := objCell.Get()

	obj[name] = v

	if err := vm.memory.Set(ptr, value.ObjectCell{Members: obj}); err != nil {
		return err
	}

	parent.Drop(vm.memory)
	// v was moved, no need to drop it.

	return nil
}

func (vm *VM) execGetEnv(envIndex int) error {
	envCell, err := vm.getMemCell(vm.stack.CurrentCallEnv(), mem.CellKindEnv, false)
	if err != nil {
		return err
	}
	env := envCell.(value.EnvCell).Get()

	// First check outlet.
	outletPtr := env.GetOutlet(envIndex)
	outletCell, err := vm.getMemCell(outletPtr, mem.CellKindOutlet, true)
	if err != nil {
		return err
	}

	if outletCell != nil {
		// Outlet is not null, so value escaped.
		outlet := outletCell.(value.OutletCell).Get()
		val := outlet.Clone(vm.memory)
		vm.stack.Push(val)
		return nil
	}

	// Outlet is null, so value has not escaped.
	stackIndex := env.GetStackIndex(envIndex)
	val, err := vm.stack.Get(stackIndex)
	if err != nil {
		return err
	}

	val = val.Clone(vm.memory)
	vm.stack.Push(val)

	return nil
}

func (vm *VM) execSetEnv(envIndex int) error {
	new, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	envCell, err := vm.getMemCell(vm.stack.CurrentCallEnv(), mem.CellKindEnv, false)
	if err != nil {
		return err
	}
	env := envCell.(value.EnvCell).Get()

	outletPtr := env.GetOutlet(envIndex)
	outletCell, err := vm.getMemCell(outletPtr, mem.CellKindOutlet, true)
	if err != nil {
		return err
	}

	if outletCell != nil {
		// Outlet is not null, so value escaped.
		outlet := outletCell.(value.OutletCell).Get()
		outlet.Drop(vm.memory)
		vm.memory.Set(outletPtr, value.OutletCell(new))
		return nil
	}

	stackIndex := env.GetStackIndex(envIndex)
	old, err := vm.stack.Get(stackIndex)
	if err != nil {
		return err
	}

	old.Drop(vm.memory)

	return vm.stack.Set(stackIndex, new)
}

func (vm *VM) execAddToEnv(localIndex int) error {
	f, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	if f.Type() != value.FunctionType {
		return ErrInvalidOperandType{
			Op: code.OpAddToEnv,
			X:  f.Type(),
		}
	}

	fn := f.Data().(value.FunctionData)

	var envPtr mem.Ptr
	if fn.Env().IsNull() {
		// Allocate new Env.
		envPtr, err = vm.memory.Allocate(mem.CellKindEnv)
		if err != nil {
			return err
		}

		vm.memory.Set(envPtr, value.EnvCell(value.NewEnv()))
		fn = fn.WithEnv(envPtr)
	} else {
		envPtr = fn.Env()
	}

	// Create outlet for referenced local, or use existing one.
	local, err := vm.stack.Local(int(localIndex))
	if err != nil {
		return err
	}

	if local.Outlet().IsNull() {
		outletPtr, err := vm.memory.Allocate(mem.CellKindOutlet)
		if err != nil {
			return err
		}
		local = local.WithOutlet(outletPtr)
	}

	// Add local to env.
	stackIndex := vm.stack.LocalToStackIndex(localIndex)

	envCell, err := vm.getMemCell(envPtr, mem.CellKindEnv, false)
	if err != nil {
		return err
	}

	env := envCell.(value.EnvCell).Get()
	env.Add(stackIndex, local.Outlet())
	vm.memory.Set(envPtr, value.EnvCell(env))

	// Push function back onto stack.
	f = f.WithData(fn)
	vm.stack.Push(f)

	vm.stack.Set(stackIndex, local)

	return nil
}

func (vm *VM) execAnchorType() error {
	t, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	o, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	if o.Type() != value.ObjectType {
		return ErrInvalidOperandTypes{
			Op: code.OpAnchorType,
			X:  o.Type(),
			Y:  t.Type(),
		}
	}

	if t.Type() != value.TypeRefType {
		return ErrInvalidOperandTypes{
			Op: code.OpAnchorType,
			X:  o.Type(),
			Y:  t.Type(),
		}
	}

	if !o.TypePtr().IsNull() {
		return ErrCantReanchorType{
			Type: o.Type(),
		}
	}

	o = o.WithType(t.Data().(value.TypeRefData).TypeRef())

	vm.stack.Push(o)
	return nil
}

func (vm *VM) execSetArgCount(argCount uint) error {
	f, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	if f.Type() != value.FunctionType {
		return ErrInvalidOperandType{
			Op: code.OpSetArgCount,
			X:  f.Type(),
		}
	}

	fn := f.Data().(value.FunctionData)

	fn = fn.WithArgs(argCount)
	f = f.WithData(fn)

	vm.stack.Push(f)
	return nil
}

type binaryOperation = func(value.Value, value.Value) (value.Value, error)
type typesToBinaryOperation = map[value.TypeKind]map[value.TypeKind]binaryOperation

func (vm *VM) execGenericBinaryOperation(op code.Op, fns typesToBinaryOperation) error {
	y, err := vm.stack.Pop()
	if err != nil {
		return err
	}
	x, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	subTypes, ok := fns[x.Type()]
	if !ok {
		return ErrInvalidOperandTypes{
			Op: op,
			X:  x.Type(),
			Y:  y.Type(),
		}
	}

	binaryOp, ok := subTypes[y.Type()]
	if !ok {
		return ErrInvalidOperandTypes{
			Op: op,
			X:  x.Type(),
			Y:  y.Type(),
		}
	}

	fmt.Print(x)
	fmt.Print(y)

	res, err := binaryOp(x, y)
	if err != nil {
		return err
	}

	x.Drop(vm.memory)
	y.Drop(vm.memory)

	vm.stack.Push(res)
	return nil
}

func createGenericNumericBinaryOperationTypes(ifn func(int64, int64) int64, ffn func(float64, float64) float64) typesToBinaryOperation {
	return typesToBinaryOperation{
		value.IntType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewInt(ifn(xv, yv)), nil
			},
			value.FloatType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.FloatData).Get()
				return value.NewFloat(ffn(float64(xv), yv)), nil
			},
		},
		value.FloatType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.FloatData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewFloat(ffn(xv, float64(yv))), nil
			},
			value.FloatType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.FloatData).Get()
				yv := y.Data().(value.FloatData).Get()
				return value.NewFloat(ffn(xv, yv)), nil
			},
		},
	}
}

func createGenericComparisonBinaryOperationTypes(ifn func(int64, int64) bool, ffn func(float64, float64) bool) typesToBinaryOperation {
	return typesToBinaryOperation{
		value.IntType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewBool(ifn(xv, yv)), nil
			},
			value.FloatType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.FloatData).Get()
				return value.NewBool(ffn(float64(xv), yv)), nil
			},
		},
		value.FloatType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.FloatData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewBool(ffn(xv, float64(yv))), nil
			},
			value.FloatType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.FloatData).Get()
				yv := y.Data().(value.FloatData).Get()
				return value.NewBool(ffn(xv, yv)), nil
			},
		},
	}
}

func (vm *VM) execAdd() error {
	return vm.execGenericBinaryOperation(code.OpAdd, typesToBinaryOperation{
		value.IntType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewInt(xv + yv), nil
			},
			value.FloatType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.FloatData).Get()
				return value.NewFloat(float64(xv) + yv), nil
			},
		},
		value.FloatType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.FloatData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewFloat(xv + float64(yv)), nil
			},
			value.FloatType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.FloatData).Get()
				yv := y.Data().(value.FloatData).Get()
				return value.NewFloat(xv + yv), nil
			},
		},
		value.StringType: {
			value.StringType: func(x, y value.Value) (value.Value, error) {
				xv, err := x.Data().(value.StringData).RawString(vm.memory)
				if err != nil {
					return value.Value{}, err
				}
				yv, err := y.Data().(value.StringData).RawString(vm.memory)
				if err != nil {
					return value.Value{}, err
				}

				return value.NewString(vm.memory, xv+yv)
			},
		},
	})
}

func (vm *VM) execSub() error {
	return vm.execGenericBinaryOperation(code.OpSub, createGenericNumericBinaryOperationTypes(
		func(x, y int64) int64 { return x - y },
		func(x, y float64) float64 { return x - y },
	))
}

func (vm *VM) execMul() error {
	return vm.execGenericBinaryOperation(code.OpMul, createGenericNumericBinaryOperationTypes(
		func(x, y int64) int64 { return x * y },
		func(x, y float64) float64 { return x * y },
	))
}

func (vm *VM) execDiv() error {
	return vm.execGenericBinaryOperation(code.OpDiv, createGenericNumericBinaryOperationTypes(
		func(x, y int64) int64 { return x / y },
		func(x, y float64) float64 { return x / y },
	))
}

func (vm *VM) execMod() error {
	return vm.execGenericBinaryOperation(code.OpMod, typesToBinaryOperation{
		value.IntType: {
			value.IntType: func(x, y value.Value) (value.Value, error) {
				xv := x.Data().(value.IntData).Get()
				yv := y.Data().(value.IntData).Get()
				return value.NewInt(xv % yv), nil
			},
		},
	})
}

func (vm *VM) execEq() error {
	y, err := vm.stack.Pop()
	if err != nil {
		return err
	}
	x, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	result := func(res bool) error {
		x.Drop(vm.memory)
		y.Drop(vm.memory)

		vm.stack.Push(value.NewBool(res))
		return nil
	}

	if x.Type() != y.Type() {
		return result(false)
	}

	if x.Data() != y.Data() {
		return result(false)
	}

	return result(true)
}

func (vm *VM) execLt() error {
	return vm.execGenericBinaryOperation(code.OpLt, createGenericComparisonBinaryOperationTypes(
		func(x, y int64) bool { return x < y },
		func(x, y float64) bool { return x < y },
	))
}

func (vm *VM) execGt() error {
	return vm.execGenericBinaryOperation(code.OpGt, createGenericComparisonBinaryOperationTypes(
		func(x, y int64) bool { return x > y },
		func(x, y float64) bool { return x > y },
	))
}

func (vm *VM) execLte() error {
	return vm.execGenericBinaryOperation(code.OpLte, createGenericComparisonBinaryOperationTypes(
		func(x, y int64) bool { return x <= y },
		func(x, y float64) bool { return x <= y },
	))
}

func (vm *VM) execGte() error {
	return vm.execGenericBinaryOperation(code.OpGte, createGenericComparisonBinaryOperationTypes(
		func(x, y int64) bool { return x >= y },
		func(x, y float64) bool { return x >= y },
	))
}

func (vm *VM) execIndex() error {
	i, err := vm.popAndDrop()
	if err != nil {
		return err
	}
	v, err := vm.popAndDrop()
	if err != nil {
		return err
	}

	switch v.Type() {
	case value.ArrayType:
		arr := v.Data().(value.ArrayData)
		switch i.Type() {
		case value.IntType:
			idx := i.Data().(value.IntData).Get()
			len, err := arr.Len(vm.memory)
			if err != nil {
				return err
			}
			if idx < 0 || idx >= int64(len) {
				return ErrArrayIndexOutOfBounds{
					Index: int(idx),
					Len:   int(len),
				}
			}

			val, err := arr.At(vm.memory, int(idx))
			if err != nil {
				return err
			}
			val = val.Clone(vm.memory)
			vm.stack.Push(val)
		default:
			return ErrInvalidOperandTypes{
				Op: code.OpIndex,
				X:  i.Type(),
				Y:  v.Type(),
			}
		}
	default:
		return ErrInvalidOperandTypes{
			Op: code.OpIndex,
			X:  v.Type(),
			Y:  i.Type(),
		}
	}

	return nil
}

func (vm *VM) execCall(argCount uint) error {
	var err error

	args := make([]value.Value, argCount)
	for i := 0; i < int(argCount); i++ {
		args[i], err = vm.stack.Pop()
		if err != nil {
			return err
		}
	}

	f, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	if f.Type() != value.FunctionType {
		return ErrInvalidOperandType{
			Op: code.OpCall,
			X:  f.Type(),
		}
	}

	fn := f.Data().(value.FunctionData)

	if argCount != fn.Args() {
		return ErrWrongNumberOfArguments{
			Got:    argCount,
			Needed: fn.Args(),
		}
	}

	if err = vm.stack.PushCall(fn.Pc(), vm.pc, fn.Env()); err != nil {
		return err
	}

	if fn.Native() != nil {
		val, err := fn.Native()(args)
		if err != nil {
			return err
		}

		if _, err := vm.stack.PopCall(); err != nil {
			return err
		}

		if !val.IsEmpty() {
			vm.stack.Push(val)
		}

		return nil
	}

	for _, arg := range args {
		vm.stack.Push(arg)
	}

	vm.pc = fn.Pc()

	return nil
}

func (vm *VM) execJumpIf(pc int, cond bool) error {
	b, err := vm.popAndDrop()
	if err != nil {
		return err
	}

	switch b.Type() {
	case value.BoolType:
		bl := b.Data().(value.BoolData)
		if bl.Get() == cond {
			vm.pc = pc
		}
	default:
		var op code.Op
		if cond {
			op = code.OpJt
		} else {
			op = code.OpJf
		}

		return ErrInvalidOperandType{
			Op: op,
			X:  b.Type(),
		}
	}

	return nil
}

func (vm *VM) execRet() error {
	returned, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	pc, err := vm.popCallAndDrop()
	if err != nil {
		return err
	}

	vm.stack.Push(returned)
	vm.pc = pc
	return nil
}

func (vm *VM) execTempArrLen() error {
	a, err := vm.stack.Pop()
	if err != nil {
		return err
	}

	switch a.Type() {
	case value.ArrayType:
		arr := a.Data().(value.ArrayData)
		len, err := arr.Len(vm.memory)
		if err != nil {
			return err
		}
		res := value.NewInt(int64(len))
		vm.stack.Push(res)
		a.Drop(vm.memory)
	default:
		return ErrInvalidOperandTypes{
			Op: code.OpTempArrLen,
			X:  a.Type(),
		}
	}

	return nil
}

func (vm *VM) execTempArrPush() error {
	e, err := vm.popAndDrop()
	if err != nil {
		return err
	}
	a, err := vm.popAndDrop()
	if err != nil {
		return err
	}

	switch a.Type() {
	case value.ArrayType:
		arr := a.Data().(value.ArrayData)
		arr.Push(vm.memory, e)
	default:
		return ErrInvalidOperandType{
			Op: code.OpTempArrPush,
			X:  a.Type(),
		}
	}

	return nil
}