package main

import "bytes"
import "encoding/hex"
import "fmt"
import "strings"

import "github.com/lunixbochs/struc"
import "github.com/pkg/errors"

type OperandType uint8

const (
	Literal             OperandType = 0
	FunctionMemoryIndex             = 1
	GlobalMemoryIndex               = 2
)

const InstructionBytes = 32

type Operand struct {
	Type     OperandType
	Reserved [3]byte
	Value    uint64
}

type Instruction struct {
	OpCode   uint32
	Reserved [4]byte
	Operand1 Operand
	Operand2 Operand
}

type State struct {
	Error error

	FunctionByteCode [][]byte
	FunctionIndex    int64
	InstructionIndex int64

	ComparisonResult bool
	GlobalMemory     [16]uint64
	Stack            []*StackFrame
	StackFrame       *StackFrame
}

type StackFrame struct {
	PreviousFunctionIndex    int64
	PreviousInstructionIndex int64
	FunctionMemory           [16]uint64
}

type OpHandler func(*State, *Instruction)

var OpHandlers = map[uint32]OpHandler{
	0x00000000: (*State).NoOp,
	0x00000001: (*State).Call,
	0x00000002: (*State).Return,

	0x00000100: (*State).Move,

	0x00000200: (*State).Add,
	0x00000201: (*State).Subtract,
	0x00000202: (*State).Multiply,
	0x00000203: (*State).DivideUnsigned,
	0x00000204: (*State).DivideSigned,

	0x00000300: (*State).IsEqual,
	0x00000301: (*State).IsLessThanUnsigned,
	0x00000302: (*State).IsLessThanSigned,
	0x00000303: (*State).IsGreaterThanUnsigned,
	0x00000304: (*State).IsGreaterThanSigned,
	0x00000305: (*State).IsLessThanOrEqualUnsigned,
	0x00000306: (*State).IsLessThanOrEqualSigned,
	0x00000307: (*State).IsGreaterThanOrEqualUnsigned,
	0x00000308: (*State).IsGreaterThanOrEqualSigned,

	0x00000400: (*State).Jump,
	0x00000401: (*State).JumpIfTrue,
	0x00000402: (*State).JumpIfFalse,
}

func NewState(functionByteCode [][]byte) *State {
	return &State{
		FunctionByteCode: functionByteCode,
	}
}

func (state *State) Execute() {
	state.call(0)

	for len(state.Stack) > 0 && state.Error == nil {
		fnc := state.FunctionByteCode[state.FunctionIndex]
		start := state.InstructionIndex * InstructionBytes
		chunk := fnc[start : start+InstructionBytes]
		state.InstructionIndex += 1
		state.ProcessInstruction(chunk)
	}

	if state.Error != nil {
		state.InstructionIndex -= 1
	}
}

func (state *State) ProcessInstruction(byteCode []byte) {
	reader := bytes.NewReader(byteCode)

	instr := &Instruction{}

	err := struc.Unpack(reader, instr)
	if err != nil {
		state.Error = errors.Wrap(err, "Error decoding instruction")
		return
	}

	fmt.Printf("%+v\n", instr)

	handler, found := OpHandlers[instr.OpCode]
	if !found {
		state.Error = fmt.Errorf("Invalid OpCode: 0x%08x", instr.OpCode)
		return
	}

	handler(state, instr)
}

func (state *State) ReadUnsigned(op *Operand) uint64 {
	switch op.Type {
	case Literal:
		return op.Value
	case FunctionMemoryIndex:
		return state.StackFrame.FunctionMemory[op.Value]
	case GlobalMemoryIndex:
		return state.GlobalMemory[op.Value]
	default:
		state.Error = fmt.Errorf("Unknown operand type: 0x%02x", op.Type)
		return 0
	}
}

func (state *State) ReadSigned(op *Operand) int64 {
	return int64(state.ReadUnsigned(op))
}

func (state *State) WriteUnsigned(op *Operand, value uint64) {
	switch op.Type {
	case Literal:
		state.Error = fmt.Errorf("Write to literal operand")
	case FunctionMemoryIndex:
		state.StackFrame.FunctionMemory[op.Value] = value
	case GlobalMemoryIndex:
		state.GlobalMemory[op.Value] = value
	default:
		state.Error = fmt.Errorf("Unknown operand type: 0x%02x", op.Type)
	}
}

func (state *State) WriteSigned(op *Operand, value int64) {
	state.WriteUnsigned(op, uint64(value))
}

func (state *State) NoOp(instr *Instruction) {
	fmt.Printf("NoOp\n")
}

func (state *State) Call(instr *Instruction) {
	fmt.Printf("Call\n")
	in := state.ReadSigned(&instr.Operand1)
	state.call(in)
}

func (state *State) call(functionOffset int64) {
	stackFrame := &StackFrame{
		PreviousFunctionIndex:    state.FunctionIndex,
		PreviousInstructionIndex: state.InstructionIndex,
	}
	state.Stack = append(state.Stack, stackFrame)
	state.StackFrame = stackFrame
	state.FunctionIndex += functionOffset
	state.InstructionIndex = 0
}

func (state *State) Return(instr *Instruction) {
	fmt.Printf("Return\n")
	state.FunctionIndex = state.StackFrame.PreviousFunctionIndex
	state.InstructionIndex = state.StackFrame.PreviousInstructionIndex
	state.Stack = state.Stack[:len(state.Stack)-1]
	if len(state.Stack) > 0 {
		state.StackFrame = state.Stack[len(state.Stack)-1]
	}
}

func (state *State) Move(instr *Instruction) {
	fmt.Printf("Move\n")
	in := state.ReadUnsigned(&instr.Operand2)
	state.WriteUnsigned(&instr.Operand1, in)
}

func (state *State) Add(instr *Instruction) {
	fmt.Printf("Add\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.WriteUnsigned(&instr.Operand1, in1+in2)
}

func (state *State) Subtract(instr *Instruction) {
	fmt.Printf("Subtract\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.WriteUnsigned(&instr.Operand1, in1-in2)
}

func (state *State) Multiply(instr *Instruction) {
	fmt.Printf("Multiply\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.WriteUnsigned(&instr.Operand1, in1*in2)
}

func (state *State) DivideUnsigned(instr *Instruction) {
	fmt.Printf("DivideUnsigned\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.WriteUnsigned(&instr.Operand1, in1/in2)
}

func (state *State) DivideSigned(instr *Instruction) {
	fmt.Printf("DivideSigned\n")
	in1 := state.ReadSigned(&instr.Operand1)
	in2 := state.ReadSigned(&instr.Operand2)
	state.WriteSigned(&instr.Operand1, in1/in2)
}

func (state *State) IsEqual(instr *Instruction) {
	fmt.Printf("IsEqual\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.ComparisonResult = (in1 == in2)
}

func (state *State) IsLessThanUnsigned(instr *Instruction) {
	fmt.Printf("IsLessThanUnsigned\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.ComparisonResult = (in1 < in2)
}

func (state *State) IsLessThanSigned(instr *Instruction) {
	fmt.Printf("IsLessThanSigned\n")
	in1 := state.ReadSigned(&instr.Operand1)
	in2 := state.ReadSigned(&instr.Operand2)
	state.ComparisonResult = (in1 < in2)
}

func (state *State) IsGreaterThanUnsigned(instr *Instruction) {
	fmt.Printf("IsGreaterThanUnsigned\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.ComparisonResult = (in1 > in2)
}

func (state *State) IsGreaterThanSigned(instr *Instruction) {
	fmt.Printf("IsGreaterThanSigned\n")
	in1 := state.ReadSigned(&instr.Operand1)
	in2 := state.ReadSigned(&instr.Operand2)
	state.ComparisonResult = (in1 > in2)
}

func (state *State) IsLessThanOrEqualUnsigned(instr *Instruction) {
	fmt.Printf("IsLessThanOrEqualUnsigned\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.ComparisonResult = (in1 <= in2)
}

func (state *State) IsLessThanOrEqualSigned(instr *Instruction) {
	fmt.Printf("IsLessThanOrEqualSigned\n")
	in1 := state.ReadSigned(&instr.Operand1)
	in2 := state.ReadSigned(&instr.Operand2)
	state.ComparisonResult = (in1 <= in2)
}

func (state *State) IsGreaterThanOrEqualUnsigned(instr *Instruction) {
	fmt.Printf("IsGreaterThanOrEqualUnsigned\n")
	in1 := state.ReadUnsigned(&instr.Operand1)
	in2 := state.ReadUnsigned(&instr.Operand2)
	state.ComparisonResult = (in1 >= in2)
}

func (state *State) IsGreaterThanOrEqualSigned(instr *Instruction) {
	fmt.Printf("IsGreaterThanOrEqualSigned\n")
	in1 := state.ReadSigned(&instr.Operand1)
	in2 := state.ReadSigned(&instr.Operand2)
	state.ComparisonResult = (in1 >= in2)
}

func (state *State) Jump(instr *Instruction) {
	fmt.Printf("Jump\n")
	in := state.ReadSigned(&instr.Operand1)
	state.InstructionIndex += in - 1
}

func (state *State) JumpIfTrue(instr *Instruction) {
	fmt.Printf("JumpIfTrue\n")
	if state.ComparisonResult == true {
		state.Jump(instr)
	}
}

func (state *State) JumpIfFalse(instr *Instruction) {
	fmt.Printf("JumpIfFalse\n")
	if state.ComparisonResult == false {
		state.Jump(instr)
	}
}

func main() {
	asm := [][]string{
		[]string{
			"0000020000000000010000000000000000000000000000000000000000000001",
			"0000000100000000000000000000000000000001000000000000000000000000",
			"0000030100000000010000000000000000000000000000000000000000000003",
			"000004010000000000000000fffffffffffffffd000000000000000000000000",
			"0000000200000000000000000000000000000000000000000000000000000000",
		},
		[]string{
			"0000020000000000020000000000000000000000000000000000000000000001",
			"0000000200000000000000000000000000000000000000000000000000000000",
		},
	}

	functionByteCode := [][]byte{}

	for _, fnc := range asm {
		fncString := strings.Join(fnc, "")

		byteCode, err := hex.DecodeString(fncString)
		if err != nil {
			panic(err)
		}

		functionByteCode = append(functionByteCode, byteCode)
	}

	state := NewState(functionByteCode)
	state.Execute()

	if state.Error != nil {
		fmt.Printf("ERROR: %s\n", state.Error)
		fmt.Printf("\tat function index 0x%016x, instruction index 0x%016x\n", state.FunctionIndex, state.InstructionIndex)
	}

	fmt.Printf("\n")
	fmt.Printf("Comparison Result: %t\n", state.ComparisonResult)

	fmt.Printf("\n")
	fmt.Printf("Global memory:\n")
	for i, v := range state.GlobalMemory {
		fmt.Printf("\t0x%08x: 0x%016x %d %d\n", i, v, v, int64(v))
	}
}