package asm

import (
	"math"
	"strings"

	"git.akyoto.dev/cli/q/src/config"
	"git.akyoto.dev/cli/q/src/dll"
	"git.akyoto.dev/cli/q/src/x86"
)

// Finalize generates the final machine code.
func (a Assembler) Finalize(dlls dll.List) ([]byte, []byte) {
	var (
		code         = make([]byte, 0, len(a.Instructions)*8)
		data         []byte
		codeLabels   = map[string]Address{}
		dataLabels   map[string]Address
		codePointers []*pointer
		dataPointers []*pointer
		dllPointers  []*pointer
		codeStart    = CodeOffset()
	)

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

		case AND:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x86.AndRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x86.AndRegisterRegister(code, operands.Destination, operands.Source)
			}

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

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

		case DIV:
			switch operands := x.Data.(type) {
			case *RegisterRegister:
				if operands.Destination != x86.RAX {
					code = x86.MoveRegisterRegister(code, x86.RAX, operands.Destination)
				}

				code = x86.ExtendRAXToRDX(code)
				code = x86.DivRegister(code, operands.Source)

				if operands.Destination != x86.RAX {
					code = x86.MoveRegisterRegister(code, operands.Destination, x86.RAX)
				}
			}

		case MODULO:
			switch operands := x.Data.(type) {
			case *RegisterRegister:
				if operands.Destination != x86.RAX {
					code = x86.MoveRegisterRegister(code, x86.RAX, operands.Destination)
				}

				code = x86.ExtendRAXToRDX(code)
				code = x86.DivRegister(code, operands.Source)

				if operands.Destination != x86.RDX {
					code = x86.MoveRegisterRegister(code, operands.Destination, x86.RDX)
				}
			}

		case CALL:
			code = x86.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 := codeLabels[label.Name]

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

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

			codePointers = append(codePointers, pointer)

		case COMMENT:
			continue

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

		case DLLCALL:
			size := 4
			// TODO: R15 could be in use.
			code = x86.MoveRegisterRegister(code, x86.R15, x86.RSP)
			code = x86.AlignStack(code)
			code = x86.SubRegisterNumber(code, x86.RSP, 32)
			code = x86.CallAtAddress(code, 0x00_00_00_00)
			position := len(code) - size
			code = x86.MoveRegisterRegister(code, x86.RSP, x86.R15)

			label := x.Data.(*Label)
			pointer := &pointer{
				Position: Address(position),
				OpSize:   2,
				Size:     uint8(size),
			}

			pointer.Resolve = func() Address {
				dot := strings.Index(label.Name, ".")
				library := label.Name[:dot]
				funcName := label.Name[dot+1:]
				index := dlls.Index(library, funcName)

				if index == -1 {
					panic("unknown DLL function " + label.Name)
				}

				return Address(index * 8)
			}

			dllPointers = append(dllPointers, pointer)

		case JE, JNE, JG, JGE, JL, JLE, JUMP:
			switch x.Mnemonic {
			case JE:
				code = x86.Jump8IfEqual(code, 0x00)
			case JNE:
				code = x86.Jump8IfNotEqual(code, 0x00)
			case JG:
				code = x86.Jump8IfGreater(code, 0x00)
			case JGE:
				code = x86.Jump8IfGreaterOrEqual(code, 0x00)
			case JL:
				code = x86.Jump8IfLess(code, 0x00)
			case JLE:
				code = x86.Jump8IfLessOrEqual(code, 0x00)
			case JUMP:
				code = x86.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 := codeLabels[label.Name]

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

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

			codePointers = append(codePointers, pointer)

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

		case LOAD:
			switch operands := x.Data.(type) {
			case *MemoryRegister:
				code = x86.LoadRegister(code, operands.Register, operands.Address.Offset, operands.Address.Length, operands.Address.Base)
			}

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

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

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

				if strings.HasPrefix(regLabel.Label, "data_") {
					dataPointers = append(dataPointers, &pointer{
						Position: Address(len(code) - size),
						OpSize:   uint8(opSize),
						Size:     uint8(size),
						Resolve: func() Address {
							destination, exists := dataLabels[regLabel.Label]

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

							return Address(destination)
						},
					})
				} else {
					codePointers = append(codePointers, &pointer{
						Position: Address(len(code) - size),
						OpSize:   uint8(opSize),
						Size:     uint8(size),
						Resolve: func() Address {
							destination, exists := codeLabels[regLabel.Label]

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

							return config.BaseAddress + codeStart + destination
						},
					})
				}
			}

		case NEGATE:
			switch operands := x.Data.(type) {
			case *Register:
				code = x86.NegateRegister(code, operands.Register)
			}

		case OR:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x86.OrRegisterNumber(code, operands.Register, operands.Number)
			case *RegisterRegister:
				code = x86.OrRegisterRegister(code, operands.Destination, operands.Source)
			}

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

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

		case RETURN:
			code = x86.Return(code)

		case SHIFTL:
			switch operands := x.Data.(type) {
			case *RegisterNumber:
				code = x86.ShiftLeftNumber(code, operands.Register, byte(operands.Number)&0b111111)
			}

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

		case STORE:
			switch operands := x.Data.(type) {
			case *MemoryNumber:
				if operands.Address.OffsetRegister == math.MaxUint8 {
					code = x86.StoreNumber(code, operands.Address.Base, operands.Address.Offset, operands.Address.Length, operands.Number)
				} else {
					code = x86.StoreDynamicNumber(code, operands.Address.Base, operands.Address.OffsetRegister, operands.Address.Length, operands.Number)
				}
			case *MemoryLabel:
				start := len(code)

				if operands.Address.OffsetRegister == math.MaxUint8 {
					code = x86.StoreNumber(code, operands.Address.Base, operands.Address.Offset, operands.Address.Length, 0b00_00_00_00)
				} else {
					code = x86.StoreDynamicNumber(code, operands.Address.Base, operands.Address.OffsetRegister, operands.Address.Length, 0b00_00_00_00)
				}

				size := 4
				opSize := len(code) - size - start
				memLabel := x.Data.(*MemoryLabel)

				codePointers = append(codePointers, &pointer{
					Position: Address(len(code) - size),
					OpSize:   uint8(opSize),
					Size:     uint8(size),
					Resolve: func() Address {
						destination, exists := codeLabels[memLabel.Label]

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

						return config.BaseAddress + codeStart + destination
					},
				})
			case *MemoryRegister:
				if operands.Address.OffsetRegister == math.MaxUint8 {
					code = x86.StoreRegister(code, operands.Address.Base, operands.Address.Offset, operands.Address.Length, operands.Register)
				} else {
					code = x86.StoreDynamicRegister(code, operands.Address.Base, operands.Address.OffsetRegister, operands.Address.Length, operands.Register)
				}
			}

		case SYSCALL:
			code = x86.Syscall(code)

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

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

	data, dataLabels = a.Data.Finalize()

	if config.TargetOS == config.Windows && len(data) == 0 {
		data = []byte{0}
	}

	code = a.resolvePointers(code, data, codeStart, codeLabels, codePointers, dataPointers, dllPointers)
	return code, data
}