package main

import (
	"bufio"
	"strings"
	"fmt"
	"os"
	"errors"
	"encoding/binary"
	"strconv"
)

const STACK_SIZE = 1024

type Operation int32

const (
	HALT Operation = 0
	PUSH Operation = 1
	POP Operation = 2
	DUP Operation = 3
	ADD Operation = 4
	SUB Operation = 5
	MUL Operation = 6
	DIV Operation = 7
	EQ Operation = 8
	JMP Operation = 9
	JMPIF Operation = 10
	CALL Operation = 11
	RET Operation = 12
)

type Inst struct {
	Operation Operation
	Operand int32
}

type Frame struct {
	Locals map[int]int32
	ReturnAddress int
}

func CreateFrame(returnAddress int32) *Frame {
	f := Frame{}
	f.Locals = make(map[int]int32)
	f.ReturnAddress = int(returnAddress)

	return &f
}

func (f *Frame) StoreVariable(address int, value int32) {
	f.Locals[address] = value
}

func (f *Frame) GetVariable(address int) int32 {
	return f.Locals[address]
}
 
type Machine struct {
	program []Inst
	stack []int32
	stackFrame []*Frame
	ip int32
	sp int32
	isHalted bool;
}

func Constructor() *Machine {
	m := Machine{}
	m.program = make([]Inst, 0)
	m.stack = make([]int32, STACK_SIZE)
	m.stackFrame = make([]*Frame, 0)
	m.stackFrame = append(m.stackFrame, CreateFrame(0))
	m.ip = 0
	m.sp = 0
	m.isHalted = false

	return &m
}

func (m* Machine) Push(i Inst) error {
	m.program = append(m.program, i)
	return nil
}

func (m* Machine) Execute() error {
	instr := m.program[m.ip]
	switch op := instr.Operation; op {
	case HALT:
		m.isHalted = true
	case PUSH:
		if (m.sp >= STACK_SIZE) {
            return errors.New("Stack Overflow");
        }

		m.stack[m.sp] = instr.Operand
		m.sp++
		m.ip++
	case POP:
		if (m.sp <= 0) {
			return errors.New("Empty Stack")
		}

		m.stack[m.sp - 1] = 0
		m.sp--
		m.ip++
	case DUP:
		if (m.sp - instr.Operand <= 0) {
			return errors.New("Stack Underflow")
		}

		if (m.sp >= STACK_SIZE) {
            return errors.New("Stack Overflow");
        }

		if (instr.Operand < 0) {
			return errors.New("Illegal access")
		}

		m.stack[m.sp] = m.stack[m.sp - 1 - instr.Operand]
		m.sp++
		m.ip++
	case ADD:
		if (m.sp < 2) {
			return errors.New("Stack size is less than required to execute binary operation")
		}
		
		m.stack[m.sp - 2] += m.stack[m.sp - 1];
		m.sp--
		m.ip++
	case SUB:
		if (m.sp < 2) {
			return errors.New("Stack size is less than required to execute binary operation")
		}

		m.stack[m.sp - 2] -= m.stack[m.sp - 1];
		m.sp--
		m.ip++
	case MUL:
		if (m.sp < 2) {
			return errors.New("Stack size is less than required to execute binary operation")
		}

		m.stack[m.sp - 2] *= m.stack[m.sp - 1];
		m.sp--
		m.ip++
	case DIV:
		if (m.sp < 2) {
			return errors.New("Stack size is less than required to execute binary operation")
		}

		if (m.stack[m.sp - 2] == 0 || m.stack[m.sp - 1] == 0) {
			return errors.New("Divide by zero exception")
		}

		m.stack[m.sp - 2] /= m.stack[m.sp - 1];
		m.sp--
		m.ip++
	case EQ:
		if (m.sp < 2) {
			return errors.New("Stack size is less than required to execute binary operation")
		}

		eq := m.stack[m.sp - 1] == m.stack[m.sp - 2]
		if (eq) {
			m.stack[m.sp - 2] = 1
		} else {
			m.stack[m.sp - 2] = 0
		}

		m.sp--
		m.ip++
	case JMP:
		if (int32(len(m.program) - 1) < instr.Operand || instr.Operand < 0) {
			return errors.New("Illegal access")
		}

		m.ip = instr.Operand
	case JMPIF:
		if (m.sp < 1) {
			return errors.New("Stack Underflow")
		}

		if (int32(len(m.program) - 1) < instr.Operand || instr.Operand < 0) {
			return errors.New("Illegal access")
		}
		
		if (m.stack[m.sp - 1] >= 1) {
			m.sp--
			m.ip = instr.Operand
		} else {
			m.ip++
		}
	case CALL:
		if (int32(len(m.program) - 1) < instr.Operand || instr.Operand < 0) {
			return errors.New("Illegal access")
		}

		m.stackFrame = append(m.stackFrame, CreateFrame(m.ip + 1))
		m.ip = instr.Operand
	case RET:
		if len(m.stackFrame) < 2 {
			return errors.New("Stackframe underflow")
		}

		currentFrame := m.stackFrame[len(m.stackFrame) - 1]
		m.stackFrame = m.stackFrame[:len(m.stackFrame) - 1]
		m.ip = int32(currentFrame.ReturnAddress)
	}
	return nil
}

func (m* Machine) Run() error {
	for !m.isHalted {
		err := m.Execute()
		if err != nil {
			return err
		}
	}
	return nil
}

func (m* Machine) Print() {
	fmt.Println("Stack:");
	for i := 0; int32(i) < m.sp; i++ {
		fmt.Println(m.stack[i])
	}
}

func (m* Machine) DumpProgramBinary(filename string) error {
	f, err := os.Create(filename)
    if err != nil {
        return errors.New("Couldn't open file")
    }
    defer f.Close()

	err = binary.Write(f, binary.LittleEndian, m.program)
    if err != nil {
		return err
        //return errors.New("Couldn't write to file")
    }

	return nil
}

func (m* Machine) LoadProgramFromBinary(filename string) error {
	f, err := os.Open(filename)
	if err != nil {
		return errors.New("Couldn't open file")
	}
	defer f.Close()

	program := make([]Inst, STACK_SIZE)
	err = binary.Read(f, binary.LittleEndian, &program)
	if err != nil {
		return err
	}

	m.program = program
	return nil
}

func (m* Machine) InterpretProgramFromFile(filename string) error {
	f, err := os.Open(filename)
	if err != nil {
		return errors.New("Couldn't open file")
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)

    for scanner.Scan() {
        line := scanner.Text()
		
		if (strings.HasPrefix(line, "HALT")){
			m.Push(Inst{Operation: HALT})
		} else if (strings.HasPrefix(line, "PUSH")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: PUSH, Operand: int32(op)})
		} else if (strings.HasPrefix(line, "POP")) {
			m.Push(Inst{Operation: POP})
		} else if (strings.HasPrefix(line, "DUP")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: DUP, Operand: int32(op)})
		} else if (strings.HasPrefix(line, "ADD")) {
			m.Push(Inst{Operation: ADD})
		} else if (strings.HasPrefix(line, "SUB")) {
			m.Push(Inst{Operation: SUB})
		} else if (strings.HasPrefix(line, "DIV")) {
			m.Push(Inst{Operation: DIV})
		} else if (strings.HasPrefix(line, "MUL")) {
			m.Push(Inst{Operation: MUL})
		} else if (strings.HasPrefix(line, "JMPIF")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: JMPIF, Operand: int32(op)})
		} else if (strings.HasPrefix(line, "JMP")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: JMP, Operand: int32(op)})
		} else if (strings.HasPrefix(line, "EQ")) {
			m.Push(Inst{Operation: EQ})
		} else if (strings.HasPrefix(line, "CALL")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: CALL, Operand: int32(op)})
		} else if (strings.HasPrefix(line, "RET")) {
			m.Push(Inst{Operation: RET})
		}
    }

    if err := scanner.Err(); err != nil {
        return err
    }

	return nil
}

func main() {
	m := Constructor()

	
	err := m.InterpretProgramFromFile("program.lil")
	if err != nil {
		fmt.Fprintf(os.Stderr, "[ERR] %s\n", err);
		os.Exit(1)
	}

	stackCounter := 0
	for !m.isHalted && stackCounter < 100 {
		err = m.Execute()
		m.Print()
		if err != nil {
			fmt.Fprintf(os.Stderr, "[ERR] %s\n", err);
			os.Exit(1)
		}
		
		stackCounter++
	}
	m.Print()
}