package main

import (
	"unicode"
	"bytes"
	"bufio"
	"strings"
	"strconv"
	"os"
	"fmt"
	"errors"
	. "e1lama/elm"
)


func Any[T any](ts []T, pred func(T) bool) bool {
	for _, t := range ts {
        if pred(t) {
            return true
        }
    }
    return false
}

func parseLabelledOperand(line string, preprocessed *PreprocessorResult) (int, error) {
	op := 0;
	dest := strings.Fields(line)[1]
	containsLetter := Any([]rune(dest), unicode.IsLetter)
	if (containsLetter) {
		op = preprocessed.LabelMap[dest]
	} else {
		op, err := strconv.Atoi(dest)
		if err != nil {
			return op, err
		}
	}

	return op, nil
}


var InstructionTable map[string]Operation = map[string]Operation{
	"HALT": 0,
	"PUSH": 1,
	"POP" : 2,
	"DUP" : 3,
	"ADD" : 4,
	"SUB" : 5,
	"MUL" : 6,
	"DIV" : 7,
	"EQ" : 8,
	"JMP" : 9,
	"JMPIF" : 10,
	"CALL" : 11,
	"RET" : 12,
	"STORE" : 13,
	"GET" : 14,
}


type PreprocessorResult struct {
	LabelMap map[string]int
	Listing *bytes.Buffer
}

func PreprocessAssembly(f *os.File) (*PreprocessorResult, error) {
	labelMap := make(map[string]int)
	lineCount := 0

	buf := bytes.Buffer{}

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
        line := scanner.Text()
		
		if (strings.HasSuffix(line, ":")) {
			labelMap[strings.Split(line, ":")[0]] = lineCount
		} else {
			_, ok := InstructionTable[strings.Fields(line)[0]]
			if (ok) {
				lineCount++
				fmt.Fprintln(&buf, line)
			} else {
				return nil, errors.New("Unknown instruction: " + line)
			}
		}
		
	}

	return &PreprocessorResult{ LabelMap: labelMap, Listing: &buf }, nil
}

func parseProgramFromFile(m* Machine, f *os.File) error {
	preprocessed, err := PreprocessAssembly(f)
	if err != nil {
		return err
	}

	scanner := bufio.NewScanner(preprocessed.Listing)
    for scanner.Scan() {
        line := scanner.Text()
		if (strings.HasPrefix(line, "//")) {
			continue;
		} else 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: int64(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: int64(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 := parseLabelledOperand(line, preprocessed)
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: JMPIF, Operand: int64(op)})
		} else if (strings.HasPrefix(line, "JMP")) {
			op, err := parseLabelledOperand(line, preprocessed)
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: JMP, Operand: int64(op)})
		} else if (strings.HasPrefix(line, "EQ")) {
			m.Push(Inst{Operation: EQ})
		} else if (strings.HasPrefix(line, "CALL")) {
			op, err := parseLabelledOperand(line, preprocessed)
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: CALL, Operand: int64(op)})
		} else if (strings.HasPrefix(line, "RET")) {
			m.Push(Inst{Operation: RET})
		} else if (strings.HasPrefix(line, "STORE")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: STORE, Operand: int64(op)})
		} else if (strings.HasPrefix(line, "GET")) {
			op, err := strconv.Atoi(strings.Fields(line)[1])
			if err != nil {
				return err
			}
			m.Push(Inst{Operation: GET, Operand: int64(op)})
		} else {
			return errors.New("Unknown instruction: " + line)
		}
    }

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

	return nil
}

func main() {
	f, err := os.Open("program.lil")
	if err != nil {
		fmt.Fprintf(os.Stderr, "[ERR] Couldn't open file: %s\n", err);
		os.Exit(1)
	}
	defer f.Close()

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

	m.Run()
}