package asm

import (
	"encoding/binary"
	"fmt"
	"slices"

	"git.urbach.dev/cli/q/src/cpu"
	"git.urbach.dev/cli/q/src/exe"
	"git.urbach.dev/cli/q/src/x86"
)

type compilerX86 struct {
	*compiler
}

func (c *compilerX86) Compile(instr Instruction) {
	switch instr := instr.(type) {
	case *Add:
		if instr.Destination == instr.Operand {
			panic("add destination register cannot be equal to the operand register")
		}

		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.AddRegisterRegister(c.code, instr.Destination, instr.Operand)
	case *AddNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.AddRegisterNumber(c.code, instr.Destination, instr.Number)
	case *And:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.AndRegisterRegister(c.code, instr.Destination, instr.Operand)
	case *AndNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.AndRegisterNumber(c.code, instr.Destination, instr.Number)
	case *Call:
		c.code = x86.Call(c.code, 0)
		patch := c.PatchLast4Bytes()

		patch.apply = func(code []byte) []byte {
			address, exists := c.labels[instr.Label]

			if !exists {
				panic("unknown label: " + instr.Label)
			}

			offset := address - patch.end
			binary.LittleEndian.PutUint32(code, uint32(offset))
			return code
		}
	case *CallExtern:
		c.code = x86.SubRegisterNumber(c.code, x86.SP, 32)
		c.code = x86.CallAt(c.code, 0)
		patch := c.PatchLast4Bytes()
		c.code = x86.AddRegisterNumber(c.code, x86.SP, 32)

		patch.apply = func(code []byte) []byte {
			index := c.libraries.Index(instr.Library, instr.Function)

			if index == -1 {
				panic(fmt.Sprintf("unknown extern function '%s' in library '%s'", instr.Function, instr.Library))
			}

			address := c.importsStart + index*8
			offset := address - patch.end
			binary.LittleEndian.PutUint32(code, uint32(offset))
			return code
		}
	case *CallRegister:
		c.code = x86.CallRegister(c.code, instr.Address)
	case *Compare:
		c.code = x86.CompareRegisterRegister(c.code, instr.Destination, instr.Source, 8)
	case *CompareAndSwap:
		c.code = x86.Lock(c.code)
		c.code = x86.CompareAndSwapFixedOffset(c.code, instr.Address, 0, instr.Length, instr.NewValue)
	case *CompareNumber:
		if instr.Number == 0 {
			c.code = x86.TestRegisterRegister(c.code, instr.Destination, instr.Destination)
		} else {
			c.code = x86.CompareRegisterNumber(c.code, instr.Destination, instr.Number)
		}
	case *ConditionalSet:
		switch instr.Condition {
		case Equal:
			c.code = x86.SetIfEqual(c.code, instr.Destination)
		case NotEqual:
			c.code = x86.SetIfNotEqual(c.code, instr.Destination)
		case Greater:
			c.code = x86.SetIfGreater(c.code, instr.Destination)
		case GreaterEqual:
			c.code = x86.SetIfGreaterEqual(c.code, instr.Destination)
		case Less:
			c.code = x86.SetIfLess(c.code, instr.Destination)
		case LessEqual:
			c.code = x86.SetIfLessEqual(c.code, instr.Destination)
		case UnsignedGreater:
			c.code = x86.SetIfUnsignedGreater(c.code, instr.Destination)
		case UnsignedGreaterEqual:
			c.code = x86.SetIfUnsignedGreaterEqual(c.code, instr.Destination)
		case UnsignedLess:
			c.code = x86.SetIfUnsignedLess(c.code, instr.Destination)
		case UnsignedLessEqual:
			c.code = x86.SetIfUnsignedLessEqual(c.code, instr.Destination)
		default:
			panic("unknown condition")
		}
	case *Divide:
		if instr.Operand == x86.R0 {
			panic("divisor register cannot be R0")
		}

		if instr.Operand == x86.R2 {
			panic("divisor register cannot be R2")
		}

		if instr.Source != x86.R0 {
			c.code = x86.MoveRegisterRegister(c.code, x86.R0, instr.Source)
		}

		c.code = x86.XorRegisterRegister(c.code, x86.R2, x86.R2)
		c.code = x86.DivUnsignedRegister(c.code, instr.Operand)

		if instr.Destination != x86.R0 {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, x86.R0)
		}
	case *DivideSigned:
		if instr.Operand == x86.R0 {
			panic("divisor register cannot be R0")
		}

		if instr.Operand == x86.R2 {
			panic("divisor register cannot be R2")
		}

		if instr.Source != x86.R0 {
			c.code = x86.MoveRegisterRegister(c.code, x86.R0, instr.Source)
		}

		c.code = x86.SignExtendR0ToR2(c.code)
		c.code = x86.DivSignedRegister(c.code, instr.Operand)

		if instr.Destination != x86.R0 {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, x86.R0)
		}
	case *Jump:
		switch instr.Condition {
		case Equal:
			c.code = x86.Jump8IfEqual(c.code, 0x00)
		case NotEqual:
			c.code = x86.Jump8IfNotEqual(c.code, 0x00)
		case Greater:
			c.code = x86.Jump8IfGreater(c.code, 0x00)
		case GreaterEqual:
			c.code = x86.Jump8IfGreaterEqual(c.code, 0x00)
		case Less:
			c.code = x86.Jump8IfLess(c.code, 0x00)
		case LessEqual:
			c.code = x86.Jump8IfLessEqual(c.code, 0x00)
		case UnsignedGreater:
			c.code = x86.Jump8IfUnsignedGreater(c.code, 0x00)
		case UnsignedGreaterEqual:
			c.code = x86.Jump8IfUnsignedGreaterEqual(c.code, 0x00)
		case UnsignedLess:
			c.code = x86.Jump8IfUnsignedLess(c.code, 0x00)
		case UnsignedLessEqual:
			c.code = x86.Jump8IfUnsignedLessEqual(c.code, 0x00)
		default:
			c.code = x86.Jump8(c.code, 0x00)
		}

		patch := &patch{
			start: len(c.code) - 2,
			end:   len(c.code),
		}

		patch.apply = func(code []byte) []byte {
			address, exists := c.labels[instr.Label]

			if !exists {
				panic("unknown label: " + instr.Label)
			}

			offset := address - patch.end

			switch code[0] {
			case 0x72, 0x73, 0x74, 0x75, 0x76, 0x77, 0x7C, 0x7D, 0x7E, 0x7F, 0xEB:
				if cpu.SizeInt(offset) == 1 {
					code[1] = byte(offset)
					return code
				}

				var jump []byte

				switch code[0] {
				case 0x72: // JB
					jump = []byte{0x0F, 0x82}
				case 0x73: // JAE
					jump = []byte{0x0F, 0x83}
				case 0x74: // JE
					jump = []byte{0x0F, 0x84}
				case 0x75: // JNE
					jump = []byte{0x0F, 0x85}
				case 0x76: // JBE
					jump = []byte{0x0F, 0x86}
				case 0x77: // JA
					jump = []byte{0x0F, 0x87}
				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.Sprintf("failed to increase pointer size for instruction 0x%x", code[0]))
				}

				shift := len(jump) + 2
				offset -= shift
				jump = binary.LittleEndian.AppendUint32(jump, uint32(offset))
				return jump
			case 0xE9:
				binary.LittleEndian.PutUint32(code[1:], uint32(offset))
				return code
			case 0x0F:
				binary.LittleEndian.PutUint32(code[2:], uint32(offset))
				return code
			default:
				panic(fmt.Sprintf("not a jump instruction 0x%x", code[0]))
			}
		}

		c.earlyPatches = append(c.earlyPatches, patch)
	case *Label:
		if instr.Align > 0 {
			padding := exe.Pad(len(c.code), int(instr.Align))
			c.code = x86.Nop(c.code, padding)
		}

		c.labels[instr.Name] = len(c.code)
	case *Load:
		scale := toX86Scale(instr.Scale, instr.Length)

		if instr.Signed {
			if instr.Length <= 4 {
				c.code = x86.LoadSignExtend(c.code, instr.Destination, instr.Base, instr.Index, scale, instr.Length)
			} else {
				c.code = x86.Load(c.code, instr.Destination, instr.Base, instr.Index, scale, instr.Length)
			}
		} else {
			if instr.Length <= 2 {
				c.code = x86.LoadZeroExtend(c.code, instr.Destination, instr.Base, instr.Index, scale, instr.Length)
			} else {
				c.code = x86.Load(c.code, instr.Destination, instr.Base, instr.Index, scale, instr.Length)
			}
		}
	case *LoadFixedOffset:
		if instr.Scale && instr.Index != 0 {
			panic("x86-64 does not support offset scaling")
		}

		if instr.Signed {
			if instr.Length <= 4 {
				c.code = x86.LoadFixedOffsetSignExtend(c.code, instr.Destination, instr.Base, int32(instr.Index), instr.Length)
			} else {
				c.code = x86.LoadFixedOffset(c.code, instr.Destination, instr.Base, int32(instr.Index), instr.Length)
			}
		} else {
			if instr.Length <= 2 {
				c.code = x86.LoadFixedOffsetZeroExtend(c.code, instr.Destination, instr.Base, int32(instr.Index), instr.Length)
			} else {
				c.code = x86.LoadFixedOffset(c.code, instr.Destination, instr.Base, int32(instr.Index), instr.Length)
			}
		}
	case *Modulo:
		if instr.Operand == x86.R0 {
			panic("modulo operand register cannot be R0")
		}

		if instr.Operand == x86.R2 {
			panic("modulo operand register cannot be R2")
		}

		if instr.Source != x86.R0 {
			c.code = x86.MoveRegisterRegister(c.code, x86.R0, instr.Source)
		}

		c.code = x86.XorRegisterRegister(c.code, x86.R2, x86.R2)
		c.code = x86.DivUnsignedRegister(c.code, instr.Operand)

		if instr.Destination != x86.R2 {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, x86.R2)
		}
	case *ModuloSigned:
		if instr.Operand == x86.R0 {
			panic("modulo operand register cannot be R0")
		}

		if instr.Operand == x86.R2 {
			panic("modulo operand register cannot be R2")
		}

		if instr.Source != x86.R0 {
			c.code = x86.MoveRegisterRegister(c.code, x86.R0, instr.Source)
		}

		c.code = x86.SignExtendR0ToR2(c.code)
		c.code = x86.DivSignedRegister(c.code, instr.Operand)

		if instr.Destination != x86.R2 {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, x86.R2)
		}
	case *Move:
		c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
	case *MoveLabel:
		c.code = x86.LoadAddress(c.code, instr.Destination, 0)
		patch := c.PatchLast4Bytes()

		patch.apply = func(code []byte) []byte {
			address, exists := c.labels[instr.Label]

			if !exists {
				panic("unknown label: " + instr.Label)
			}

			offset := address - patch.end
			binary.LittleEndian.PutUint32(code, uint32(offset))
			return code
		}
	case *MoveNumber:
		c.code = x86.MoveRegisterNumber(c.code, instr.Destination, instr.Number)
	case *Multiply:
		if instr.Destination == instr.Operand {
			panic("multiply destination register cannot be equal to the operand register")
		}

		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.MulRegisterRegister(c.code, instr.Destination, instr.Operand)
	case *Negate:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.NegateRegister(c.code, instr.Destination)
	case *Or:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.OrRegisterRegister(c.code, instr.Destination, instr.Operand)
	case *OrNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.OrRegisterNumber(c.code, instr.Destination, instr.Number)
	case *Pop:
		for _, register := range slices.Backward(instr.Registers) {
			c.code = x86.PopRegister(c.code, register)
		}
	case *Push:
		for _, register := range instr.Registers {
			c.code = x86.PushRegister(c.code, register)
		}
	case *ReadSystemRegister:
		switch instr.SystemRegister {
		case x86.FS:
			c.code = x86.ReadFS(c.code, instr.Destination)
		case x86.GS:
			c.code = x86.ReadGS(c.code, instr.Destination)
		default:
			panic("unknown system register")
		}
	case *Return:
		c.code = x86.Return(c.code)
	case *ShiftLeft:
		c.code = prepareShiftX86(c.code, instr.Destination, instr.Source, instr.Operand)
		c.code = x86.ShiftLeft(c.code, instr.Destination)
	case *ShiftLeftNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.ShiftLeftNumber(c.code, instr.Destination, byte(instr.Number))
	case *ShiftRight:
		c.code = prepareShiftX86(c.code, instr.Destination, instr.Source, instr.Operand)
		c.code = x86.ShiftRight(c.code, instr.Destination)
	case *ShiftRightNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.ShiftRightNumber(c.code, instr.Destination, byte(instr.Number))
	case *ShiftRightSigned:
		c.code = prepareShiftX86(c.code, instr.Destination, instr.Source, instr.Operand)
		c.code = x86.ShiftRightSigned(c.code, instr.Destination)
	case *ShiftRightSignedNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.ShiftRightSignedNumber(c.code, instr.Destination, byte(instr.Number))
	case *StackFrameStart:
		if instr.FramePointer {
			c.code = x86.PushRegister(c.code, x86.R5)
			c.code = x86.MoveRegisterRegister(c.code, x86.R5, x86.SP)
		}

		if instr.ExternCalls {
			c.code = x86.AndRegisterNumber(c.code, x86.SP, -16)
		}
	case *StackFrameEnd:
		if instr.FramePointer {
			c.code = x86.MoveRegisterRegister(c.code, x86.SP, x86.R5)
			c.code = x86.PopRegister(c.code, x86.R5)
		}
	case *Store:
		scale := toX86Scale(instr.Scale, instr.Length)
		c.code = x86.StoreRegister(c.code, instr.Base, instr.Index, scale, instr.Length, instr.Source)
	case *StoreFixedOffset:
		if instr.Scale && instr.Index != 0 {
			panic("x86-64 does not support offset scaling")
		}

		c.code = x86.StoreFixedOffsetRegister(c.code, instr.Base, int32(instr.Index), instr.Length, instr.Source)
	case *StoreFixedOffsetNumber:
		if instr.Scale && instr.Index != 0 {
			panic("x86-64 does not support offset scaling")
		}

		c.code = x86.StoreFixedOffsetNumber(c.code, instr.Base, int32(instr.Index), instr.Length, instr.Number)
	case *StoreNumber:
		scale := toX86Scale(instr.Scale, instr.Length)
		c.code = x86.StoreNumber(c.code, instr.Base, instr.Index, scale, instr.Length, instr.Number)
	case *Subtract:
		if instr.Destination == instr.Operand {
			panic("subtract destination register cannot be equal to the operand register")
		}

		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.SubRegisterRegister(c.code, instr.Destination, instr.Operand)
	case *SubtractNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.SubRegisterNumber(c.code, instr.Destination, instr.Number)
	case *Syscall:
		c.code = x86.Syscall(c.code)
	case *WriteSystemRegister:
		switch instr.SystemRegister {
		case x86.FS:
			c.code = x86.WriteFS(c.code, instr.Source)
		case x86.GS:
			c.code = x86.WriteGS(c.code, instr.Source)
		default:
			panic("unknown system register")
		}
	case *Xor:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.XorRegisterRegister(c.code, instr.Destination, instr.Operand)
	case *XorNumber:
		if instr.Destination != instr.Source {
			c.code = x86.MoveRegisterRegister(c.code, instr.Destination, instr.Source)
		}

		c.code = x86.XorRegisterNumber(c.code, instr.Destination, instr.Number)
	default:
		panic("unknown instruction")
	}
}

// prepareShiftX86 checks that the registers are correct for the shift instruction
// and also moves source to destination and the operand to R1 if needed.
func prepareShiftX86(code []byte, destination cpu.Register, source cpu.Register, operand cpu.Register) []byte {
	if destination == x86.R1 {
		panic("shift destination cannot be R1")
	}

	if destination == operand {
		panic("shift destination register cannot be equal to the operand register")
	}

	if destination != source {
		code = x86.MoveRegisterRegister(code, destination, source)
	}

	if operand != x86.R1 {
		code = x86.MoveRegisterRegister(code, x86.R1, operand)
	}

	return code
}

// toX86Scale returns the scale factor for the memory instruction.
func toX86Scale(enable bool, length byte) x86.Scale {
	if !enable {
		return x86.Scale1
	}

	switch length {
	case 1:
		return x86.Scale1
	case 2:
		return x86.Scale2
	case 4:
		return x86.Scale4
	case 8:
		return x86.Scale8
	default:
		panic("unsupported scale")
	}
}