package compiler

import (
	"bufio"
	"fmt"
	"io"
	"os"

	"git.akyoto.dev/cli/q/src/arch/x64"
	"git.akyoto.dev/cli/q/src/asm"
	"git.akyoto.dev/cli/q/src/config"
	"git.akyoto.dev/cli/q/src/core"
	"git.akyoto.dev/cli/q/src/os/linux"
	"git.akyoto.dev/cli/q/src/os/linux/elf"
	"git.akyoto.dev/cli/q/src/os/mac"
	"git.akyoto.dev/cli/q/src/os/mac/macho"
)

// Result contains all the compiled functions in a build.
type Result struct {
	Main             *core.Function
	Functions        map[string]*core.Function
	InstructionCount int
	DataCount        int
}

// finalize generates the final machine code.
func (r *Result) finalize() ([]byte, []byte) {
	// This will be the entry point of the executable.
	// The only job of the entry function is to call `main` and exit cleanly.
	// The reason we call `main` instead of using `main` itself is to place
	// a return address on the stack, which allows return statements in `main`.
	final := asm.Assembler{
		Instructions: make([]asm.Instruction, 0, r.InstructionCount+8),
		Data:         make(map[string][]byte, r.DataCount),
	}

	sysExit := 0

	switch config.TargetOS {
	case "linux":
		sysExit = linux.Exit
	case "mac":
		sysExit = mac.Exit
	}

	final.Call("main.main")
	final.RegisterNumber(asm.MOVE, x64.SyscallInputRegisters[0], sysExit)
	final.RegisterNumber(asm.MOVE, x64.SyscallInputRegisters[1], 0)
	final.Syscall()

	// This will place the main function immediately after the entry point
	// and also add everything the main function calls recursively.
	r.eachFunction(r.Main, map[*core.Function]bool{}, func(f *core.Function) {
		final.Merge(f.Assembler)
	})

	final.Label(asm.LABEL, "_crash")
	final.RegisterNumber(asm.MOVE, x64.SyscallInputRegisters[0], sysExit)
	final.RegisterNumber(asm.MOVE, x64.SyscallInputRegisters[1], 1)
	final.Syscall()

	code, data := final.Finalize()
	return code, data
}

// eachFunction recursively finds all the calls to external functions.
// It avoids calling the same function twice with the help of a hashmap.
func (r *Result) eachFunction(caller *core.Function, traversed map[*core.Function]bool, call func(*core.Function)) {
	call(caller)
	traversed[caller] = true

	for _, x := range caller.Assembler.Instructions {
		if x.Mnemonic != asm.CALL {
			continue
		}

		name := x.Data.(*asm.Label).Name
		callee, exists := r.Functions[name]

		if !exists {
			continue
		}

		if traversed[callee] {
			continue
		}

		r.eachFunction(callee, traversed, call)
	}
}

// PrintInstructions prints out the generated instructions.
func (r *Result) PrintInstructions() {
	r.eachFunction(r.Main, map[*core.Function]bool{}, func(f *core.Function) {
		f.PrintInstructions()
	})
}

// Write writes the executable to the given writer.
func (r *Result) Write(writer io.Writer) error {
	code, data := r.finalize()
	return write(writer, code, data)
}

// Write writes an executable file to disk.
func (r *Result) WriteFile(path string) error {
	file, err := os.Create(path)

	if err != nil {
		return err
	}

	err = r.Write(file)

	if err != nil {
		file.Close()
		return err
	}

	err = file.Close()

	if err != nil {
		return err
	}

	return os.Chmod(path, 0755)
}

// write writes an executable file to the given writer.
func write(writer io.Writer, code []byte, data []byte) error {
	buffer := bufio.NewWriter(writer)

	switch config.TargetOS {
	case "linux":
		exe := elf.New(code, data)
		exe.Write(buffer)
	case "mac":
		exe := macho.New(code, data)
		exe.Write(buffer)
	default:
		return fmt.Errorf("unsupported platform '%s'", config.TargetOS)
	}

	return buffer.Flush()
}