package asm

import (
	"encoding/binary"
	"fmt"

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

// 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 AND:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.AndRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.AndRegisterRegister(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)

			pointer := &Pointer{
				Position: Address(len(code) - size),
				OpSize:   1,
				Size:     uint8(size),
			}

			pointer.Resolve = func() Address {
				destination, exists := labels[label.Name]

				if !exists {
					panic("unknown jump label")
				}

				distance := destination - (pointer.Position + Address(pointer.Size))
				return Address(distance)
			}

			pointers = append(pointers, pointer)

		case COMMENT:
			continue

		case COMPARE:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.CompareRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.CompareRegisterRegister(code, operands.Destination, operands.Source)
			}

		case JE, JNE, JG, JGE, JL, JLE, JUMP:
			switch x.Mnemonic {
			case JE:
				code = x64.Jump8IfEqual(code, 0x00)
			case JNE:
				code = x64.Jump8IfNotEqual(code, 0x00)
			case JG:
				code = x64.Jump8IfGreater(code, 0x00)
			case JGE:
				code = x64.Jump8IfGreaterOrEqual(code, 0x00)
			case JL:
				code = x64.Jump8IfLess(code, 0x00)
			case JLE:
				code = x64.Jump8IfLessOrEqual(code, 0x00)
			case JUMP:
				code = x64.Jump8(code, 0x00)
			}

			size := 1
			label := x.Data.(*Label)

			pointer := &Pointer{
				Position: Address(len(code) - size),
				OpSize:   1,
				Size:     uint8(size),
			}

			pointer.Resolve = func() Address {
				destination, exists := labels[label.Name]

				if !exists {
					panic("unknown jump label")
				}

				distance := destination - (pointer.Position + Address(pointer.Size))
				return Address(distance)
			}

			pointers = append(pointers, pointer)

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

		case MODULO:
			code = modulo(code, x.Data)

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

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

			case *RegisterLabel:
				start := len(code)
				code = x64.MoveRegisterNumber(code, operands.Register, 0x00_00_00_00)
				size := 4
				opSize := len(code) - size - start
				regLabel := x.Data.(*RegisterLabel)

				pointers = append(pointers, &Pointer{
					Position: Address(len(code) - size),
					OpSize:   uint8(opSize),
					Size:     uint8(size),
					Resolve: func() Address {
						destination, exists := labels[regLabel.Label]

						if !exists {
							panic("unknown label")
						}

						return Address(destination)
					},
				})
			}

		case OR:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.OrRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.OrRegisterRegister(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 SHIFTL:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.ShiftLeftNumber(code, operands.Register, byte(operands.Number)&0b111111)
			}

		case SHIFTRS:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.ShiftRightSignedNumber(code, operands.Register, byte(operands.Number)&0b111111)
			}

		case STORE:
			switch operands := x.Data.(type) {
			case *MemoryNumber:
				code = x64.StoreNumber(code, operands.Address.Base, operands.Address.Offset, operands.Address.Length, operands.Number)
			case *MemoryRegister:
				code = x64.StoreRegister(code, operands.Address.Base, operands.Address.Offset, operands.Address.Length, operands.Register)
			}

		case SYSCALL:
			code = x64.Syscall(code)

		case XOR:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x64.XorRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x64.XorRegisterRegister(code, operands.Destination, operands.Source)
			}

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

	dataStart := config.BaseAddress + config.CodeOffset + Address(len(code))
	dataStart += int32(elf.Padding(int64(dataStart), config.Align))

	for label, slice := range a.Data {
		labels[label] = dataStart + Address(len(data))
		data = append(data, slice...)
	}

restart:
	for i, pointer := range pointers {
		address := pointer.Resolve()

		if x64.SizeOf(int64(address)) > int(pointer.Size) {
			left := code[:pointer.Position-Address(pointer.OpSize)]
			right := code[pointer.Position+Address(pointer.Size):]
			size := pointer.Size + pointer.OpSize
			opCode := code[pointer.Position-Address(pointer.OpSize)]

			var jump []byte

			switch opCode {
			case 0x74: // JE
				jump = []byte{0x0F, 0x84}
			case 0x75: // JNE
				jump = []byte{0x0F, 0x85}
			case 0x7C: // JL
				jump = []byte{0x0F, 0x8C}
			case 0x7D: // JGE
				jump = []byte{0x0F, 0x8D}
			case 0x7E: // JLE
				jump = []byte{0x0F, 0x8E}
			case 0x7F: // JG
				jump = []byte{0x0F, 0x8F}
			case 0xEB: // JMP
				jump = []byte{0xE9}
			default:
				panic(fmt.Errorf("failed to increase pointer size for instruction 0x%x", opCode))
			}

			pointer.Position += Address(len(jump) - int(pointer.OpSize))
			pointer.OpSize = uint8(len(jump))
			pointer.Size = 4
			jump = binary.LittleEndian.AppendUint32(jump, uint32(address))
			offset := Address(len(jump)) - Address(size)

			for _, following := range pointers[i+1:] {
				following.Position += offset
			}

			for key, address := range labels {
				if address > pointer.Position {
					labels[key] += offset
				}
			}

			code = make([]byte, len(left)+len(jump)+len(right))
			copy(code, left)
			copy(code[len(left):], jump)
			copy(code[len(left)+len(jump):], right)
			goto restart
		}

		slice := code[pointer.Position : pointer.Position+Address(pointer.Size)]

		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
}