// The compiler
package main

import (
	"bytes"
	"embed"
	"flag"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

const ERROR_LOG_LINES = 5

//go:embed stdlib/*
var stdlib embed.FS

func countTabs(line string) int {
	tabs := 0
	for _, rune := range line {
		if rune == '\t' {
			tabs++
		}
	}
	return tabs
}

func printCompilerError(sources map[string]string, err CompilerError) {
	source, ok := sources[err.Position.SourceFile]
	if !ok {
		log.Println(err)
		return
	}

	sourceRunes := []rune(source)
	lines := strings.Split(source, "\n")
	line := 0
	col := 0
	var i uint64
	for i = 0; i < err.Position.Position; i++ {
		col++
		if sourceRunes[i] == '\n' {
			line++
			col = 0
		}
	}

	log.Println("Failed to compile: " + err.Message + " (at " + err.Position.SourceFile + ":" + strconv.Itoa(line+1) + ":" + strconv.Itoa(col+1) + ")")

	linesStart := max(0, line-ERROR_LOG_LINES)
	linesEnd := min(len(lines), line+ERROR_LOG_LINES+1)

	for _, line := range lines[linesStart:line] {
		println(strings.Replace(line, "\t", "    ", -1))
	}

	tabs := countTabs(lines[line])
	println(strings.Repeat(" ", col+tabs*3) + "v--- error occurs here ---")

	for _, line := range lines[line:linesEnd] {
		println(strings.Replace(line, "\t", "    ", -1))
	}
}

func readEmbedDir(name string, files map[string]string) {
	entries, err := stdlib.ReadDir(name)
	if err != nil {
		log.Fatalln(err)
	}

	for _, entry := range entries {
		fullName := name + "/" + entry.Name()
		if entry.IsDir() {
			readEmbedDir(fullName, files)
		} else {
			bytes, err := stdlib.ReadFile(fullName)
			if err != nil {
				log.Fatalln(err)
			}

			files[fullName] = string(bytes)
		}
	}
}

func main() {
	outputFile := flag.String("o", "a.out", "Output file")
	generateWAT := flag.Bool("wat", false, "Generate WAT instead of WASM")
	wasm64 := flag.Bool("wasm64", false, "Use 64-bit memory (may not be supported in all browsers)")
	includeStdlib := flag.Bool("stdlib", true, "Include the standard library")
	flag.Parse()

	if len(os.Args) < 2 {
		log.Fatalln("Usage: " + os.Args[0] + " <files...>")
	}

	files := flag.Args()

	fileSources := make(map[string]string)
	for _, file := range files {
		content, err := os.ReadFile(file)
		if err != nil {
			log.Fatalln("Cannot open input file.", err)
		}

		fileSources[file] = string(content)
	}

	if *includeStdlib {
		stdlibFiles := make(map[string]string)
		readEmbedDir("stdlib", stdlibFiles)
		for path, file := range stdlibFiles {
			fileSources["[embedded]/"+path] = file
		}
	}

	fileTokens := make(map[string][]LexToken)
	for file, source := range fileSources {
		tokens, err := lexer(file, source)
		if err != nil {
			if c, ok := err.(CompilerError); ok {
				printCompilerError(fileSources, c)
			} else {
				log.Println(err)
			}

			return
		}

		// log.Printf("Tokens:\n%+#v\n\n", tokens)
		fileTokens[file] = tokens
	}

	var parsedFiles []*ParsedFile
	for _, tokens := range fileTokens {
		parser := Parser{Tokens: tokens}
		parsed, err := parser.parseFile()
		if err != nil {
			if c, ok := err.(CompilerError); ok {
				printCompilerError(fileSources, c)
			} else {
				log.Println(err)
			}

			return
		}

		// log.Printf("Parsed:\n%+#v\n\n", parsed)
		parsedFiles = append(parsedFiles, parsed)
	}

	validator := Validator{Files: parsedFiles, Wasm64: *wasm64}
	errors := validator.validate()
	if len(errors) != 0 {
		for _, err := range errors {
			if c, ok := err.(CompilerError); ok {
				printCompilerError(fileSources, c)
			} else {
				log.Println(err)
			}
		}

		if len(errors) != 0 {
			return
		}
	}

	// log.Printf("Validated:\n%+#v\n\n", parsedFiles)

	compiler := Compiler{Files: parsedFiles, Wasm64: *wasm64}
	wat, err := compiler.compile()
	if err != nil {
		if c, ok := err.(CompilerError); ok {
			printCompilerError(fileSources, c)
		} else {
			log.Println(err)
		}

		return
	}

	// log.Println("WAT: " + wat)

	if *generateWAT {
		os.WriteFile(*outputFile, []byte(wat), 0o644)
		return
	}

	var command []string
	if *wasm64 {
		// TODO: wasm64 support is currently broken because x.load doesn't accept 64-bit addresses
		command = append(command, "--enable-memory64")
	}
	command = append(command, "-o", *outputFile, "-")

	cmd := exec.Command("wat2wasm", command...)

	var input bytes.Buffer
	input.Write([]byte(wat))

	cmd.Stdin = &input

	output, err := cmd.CombinedOutput()
	if err != nil {
		log.Fatalln(err, string(output))
	}
}