package asm

import (
	"encoding/binary"

	"git.akyoto.dev/cli/q/src/build/arch/x64"
)

// Assembler contains a list of instructions.
type Assembler struct {
	Instructions []Instruction
}

// Finalize generates the final machine code.
func (a Assembler) Finalize() ([]byte, []byte) {
	code := make([]byte, 0, len(a.Instructions)*8)
	data := make([]byte, 0, 16)
	labels := map[string]Address{}
	pointers := []Pointer{}

	for _, x := range a.Instructions {
		switch x.Mnemonic {
		case ADD:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.AddRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.AddRegisterRegister(code, operands.Destination, operands.Source)
			}

		case SUB:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.SubRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.SubRegisterRegister(code, operands.Destination, operands.Source)
			}

		case MUL:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.MulRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.MulRegisterRegister(code, operands.Destination, operands.Source)
			}

		case DIV:
			code = divide(code, x.Data)

		case CALL:
			code = x64.Call(code, 0x00_00_00_00)
			size := 4
			label := x.Data.(*Label)
			nextInstructionAddress := Address(len(code))

			pointers = append(pointers, Pointer{
				Position: Address(len(code) - size),
				Size:     uint8(size),
				Resolve: func() Address {
					destination := labels[label.Name]
					distance := destination - nextInstructionAddress
					return Address(distance)
				},
			})

		case JUMP:
			code = x64.Jump8(code, 0x00)
			size := 1
			label := x.Data.(*Label)
			nextInstructionAddress := Address(len(code))

			pointers = append(pointers, Pointer{
				Position: Address(len(code) - size),
				Size:     uint8(size),
				Resolve: func() Address {
					destination := labels[label.Name]
					distance := destination - nextInstructionAddress
					return Address(distance)
				},
			})

		case LABEL:
			labels[x.Data.(*Label).Name] = Address(len(code))

		case MOVE:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.MoveRegisterNumber32(code, operands.Register, uint32(operands.Number))

			case *RegisterRegister:
				code = x64.MoveRegisterRegister64(code, operands.Destination, operands.Source)
			}

		case POP:
			switch operands := x.Data.(type) {
			case *Register:
				code = x64.PopRegister(code, operands.Register)
			}

		case PUSH:
			switch operands := x.Data.(type) {
			case *Register:
				code = x64.PushRegister(code, operands.Register)
			}

		case RETURN:
			code = x64.Return(code)

		case SYSCALL:
			code = x64.Syscall(code)

		default:
			panic("Unknown mnemonic: " + x.Mnemonic.String())
		}
	}

	// dataStart := config.BaseAddress + config.CodeOffset + Address(len(code))

	for _, pointer := range pointers {
		slice := code[pointer.Position : pointer.Position+Address(pointer.Size)]
		address := pointer.Resolve()

		switch pointer.Size {
		case 1:
			slice[0] = uint8(address)

		case 2:
			binary.LittleEndian.PutUint16(slice, uint16(address))

		case 4:
			binary.LittleEndian.PutUint32(slice, uint32(address))

		case 8:
			binary.LittleEndian.PutUint64(slice, uint64(address))
		}
	}

	return code, data
}

// Merge combines the contents of this assembler with another one.
func (a *Assembler) Merge(b Assembler) {
	a.Instructions = append(a.Instructions, b.Instructions...)
}

// divide implements the division on x64 machines.
func divide(code []byte, data any) []byte {
	code = x64.PushRegister(code, x64.RDX)

	switch operands := data.(type) {
	case *RegisterNumber:
		if operands.Register == x64.RAX {
			code = x64.PushRegister(code, x64.RCX)
			code = x64.MoveRegisterNumber32(code, x64.RCX, uint32(operands.Number))
			code = x64.ExtendRAXToRDX(code)
			code = x64.DivRegister(code, x64.RCX)
			code = x64.PopRegister(code, x64.RCX)
		} else {
			code = x64.PushRegister(code, x64.RAX)
			code = x64.MoveRegisterRegister64(code, x64.RAX, operands.Register)
			code = x64.MoveRegisterNumber32(code, operands.Register, uint32(operands.Number))
			code = x64.ExtendRAXToRDX(code)
			code = x64.DivRegister(code, operands.Register)
			code = x64.MoveRegisterRegister64(code, operands.Register, x64.RAX)
			code = x64.PopRegister(code, x64.RAX)
		}

	case *RegisterRegister:
		if operands.Destination == x64.RAX {
			code = x64.ExtendRAXToRDX(code)
			code = x64.DivRegister(code, operands.Source)
		} else {
			code = x64.PushRegister(code, x64.RAX)
			code = x64.MoveRegisterRegister64(code, x64.RAX, operands.Destination)
			code = x64.ExtendRAXToRDX(code)
			code = x64.DivRegister(code, operands.Source)
			code = x64.MoveRegisterRegister64(code, operands.Destination, x64.RAX)
			code = x64.PopRegister(code, x64.RAX)
		}
	}

	code = x64.PopRegister(code, x64.RDX)
	return code
}