package disass

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"maps"
	"slices"
	"strings"

	"github.com/apex/log"
	"github.com/blacktop/arm64-cgo/disassemble"
	"github.com/blacktop/go-macho"
	"github.com/blacktop/go-macho/types"
	"github.com/blacktop/ipsw/internal/utils"
	"github.com/blacktop/ipsw/pkg/symbols"
)

type Disass interface {
	Triage() error
	IsFunctionStart(uint64) (bool, string)
	IsLocation(uint64) bool
	IsBranchLocation(uint64) (bool, uint64)
	IsData(uint64) (bool, *AddrDetails)
	IsPointer(uint64) (bool, *AddrDetails)
	FindSymbol(uint64) (string, bool)
	FindSwiftString(uint64) (string, bool)
	GetCString(uint64) (string, error)
	// getters
	Demangle() bool
	Quite() bool
	Color() bool
	AsJSON() bool
	Data() []byte
	StartAddr() uint64
	Middle() uint64
	ReadAddr(uint64) (uint64, error)
}

type opName uint32

const (
	AMXLDX opName = iota
	AMXLDY
	AMXSTX
	AMXSTY
	AMXLDZ
	AMXSTZ
	AMXLDZI
	AMXSTZI
	AMXEXTRX // amxextrx?
	AMXEXTRY // amxextry?
	AMXFMA64
	AMXFMS64
	AMXFMA32
	AMXFMS32
	AMXMAC16
	AMXFMA16
	AMXFMS16
	AMX17 // amxset / amxclr
	AMXVECINT
	AMXVECFP
	AMXMATINT
	AMXMATFP
	AMXGENLUT
)

func (o opName) String() string {
	switch o {
	case AMXLDX:
		return "amx_ldx"
	case AMXLDY:
		return "amx_ldy"
	case AMXSTX:
		return "amx_stx"
	case AMXSTY:
		return "amx_sty"
	case AMXLDZ:
		return "amx_ldz"
	case AMXSTZ:
		return "amx_stz"
	case AMXLDZI:
		return "amx_ldzi"
	case AMXSTZI:
		return "amx_stzi"
	case AMXEXTRX:
		return "amx_extrx"
	case AMXEXTRY:
		return "amx_extry"
	case AMXFMA64:
		return "amx_fma64"
	case AMXFMS64:
		return "amx_fms64"
	case AMXFMA32:
		return "amx_fma32"
	case AMXFMS32:
		return "amx_fms32"
	case AMXMAC16:
		return "amx_mac16"
	case AMXFMA16:
		return "amx_fma16"
	case AMXFMS16:
		return "amx_fms16"
	case AMX17:
		return "amx_op17"
	case AMXVECINT:
		return "amx_vecint"
	case AMXVECFP:
		return "amx_vecfp"
	case AMXMATINT:
		return "amx_matint"
	case AMXMATFP:
		return "amx_matfp"
	case AMXGENLUT:
		return "amx_genlut"
	default:
		return "unk"
	}
}

type Config struct {
	Image        string
	Data         []byte
	StartAddress uint64
	Middle       uint64
	AsJSON       bool
	Demangle     bool
	Quiet        bool
	Color        bool
}
type AddrDetails struct {
	Image   string
	Segment string
	Section string
	Pointer uint64
	Flags   types.SectionFlag
}

func (d AddrDetails) String() string {
	if len(d.Image) > 0 {
		return fmt.Sprintf("%s/%s.%s", d.Image, d.Segment, d.Section)
	}
	return fmt.Sprintf("%s.%s", d.Segment, d.Section)
}

type Triage struct {
	Details   map[uint64]AddrDetails
	Function  *types.Function
	Addresses map[uint64]uint64
	Locations map[uint64][]uint64
}

// IsBranchOp returns true if the operation is a branch with a PC-relative target
func IsBranchOp(op disassemble.Operation) bool {
	switch op {
	case disassemble.ARM64_B, disassemble.ARM64_BL,
		disassemble.ARM64_CBZ, disassemble.ARM64_CBNZ,
		disassemble.ARM64_TBZ, disassemble.ARM64_TBNZ,
		disassemble.ARM64_B_EQ, disassemble.ARM64_B_NE,
		disassemble.ARM64_B_CS, disassemble.ARM64_B_CC,
		disassemble.ARM64_B_MI, disassemble.ARM64_B_PL,
		disassemble.ARM64_B_VS, disassemble.ARM64_B_VC,
		disassemble.ARM64_B_HI, disassemble.ARM64_B_LS,
		disassemble.ARM64_B_GE, disassemble.ARM64_B_LT,
		disassemble.ARM64_B_GT, disassemble.ARM64_B_LE,
		disassemble.ARM64_B_AL, disassemble.ARM64_B_NV:
		return true
	}
	return false
}

// IsLoadLiteral returns true if the instruction is a PC-relative literal load
func IsLoadLiteral(instr *disassemble.Inst) bool {
	switch instr.Operation {
	case disassemble.ARM64_LDR, disassemble.ARM64_LDRSW, disassemble.ARM64_PRFM:
		return instr != nil && instr.NumOps > 1 && instr.Operands[1].Class == disassemble.LABEL
	}
	return false
}

func operandRegister(instr *disassemble.Inst, idx int) (disassemble.Register, bool) {
	if instr == nil || int(instr.NumOps) <= idx || instr.Operands[idx].NumRegisters == 0 {
		return disassemble.REG_NONE, false
	}
	return instr.Operands[idx].Registers[0], true
}

func operandImmediate(instr *disassemble.Inst, idx int) (uint64, bool) {
	if instr == nil || int(instr.NumOps) <= idx {
		return 0, false
	}
	return instr.Operands[idx].Immediate, true
}

func operandShiftValue(instr *disassemble.Inst, idx int) (uint64, bool) {
	if instr == nil || int(instr.NumOps) <= idx {
		return 0, false
	}
	return uint64(instr.Operands[idx].ShiftValue), true
}

// operandLeftShift returns operand idx's shift amount only when it is an explicit
// logical-left shift (LSL #n). It returns ok=false for unshifted operands or any
// other shift type (LSR/ASR/ROR), so callers cannot mistake them for a left shift.
func operandLeftShift(instr *disassemble.Inst, idx int) (uint64, bool) {
	if instr == nil || int(instr.NumOps) <= idx {
		return 0, false
	}
	op := instr.Operands[idx]
	if !op.ShiftValueUsed || op.ShiftType != disassemble.SHIFT_TYPE_LSL {
		return 0, false
	}
	return uint64(op.ShiftValue), true
}

func implSpecificSysReg(op *disassemble.Op) (string, bool) {
	if op == nil || !op.HasImplSpec {
		return "", false
	}

	sysReg := disassemble.SystemReg(op.ImplSpec[0])<<14 |
		disassemble.SystemReg(op.ImplSpec[1])<<11 |
		disassemble.SystemReg(op.ImplSpec[2])<<7 |
		disassemble.SystemReg(op.ImplSpec[3])<<3 |
		disassemble.SystemReg(op.ImplSpec[4])

	name := sysReg.String()
	if len(name) == 0 {
		return "", false
	}

	return name, true
}

func invalidInstruction(addr uint64, raw uint32, err error) disassemble.Instruction {
	return disassemble.Instruction{
		Address:     addr,
		Raw:         raw,
		Encoding:    0,
		Operation:   0,
		Operands:    nil,
		SetFlags:    0,
		Disassembly: fmt.Sprintf(".long\t%#x ; (%s)\n", raw, err.Error()),
	}
}

func logStubDecodeContext(
	instruction *disassemble.Inst,
	queue *[3]disassemble.Inst,
	queueValid *[3]bool,
) {
	if instruction != nil {
		if text, err := instruction.Disassemble(); err == nil {
			log.Debugf("%#08x:  %s\t%s", instruction.Address, disassemble.GetOpCodeByteString(instruction.Raw), text)
		}
	}

	for idx := range queue {
		if queueValid[idx] {
			if text, err := queue[idx].Disassemble(); err == nil {
				log.Debugf("%#08x:  %s\t%s", queue[idx].Address, disassemble.GetOpCodeByteString(queue[idx].Raw), text)
			}
		}
	}
}

const maxCStringCommentLen = 200

func cstringComment(d Disass, addr uint64) string {
	if ok, detail := d.IsData(addr); !ok || detail == nil || !detail.Flags.IsCstringLiterals() {
		return ""
	}
	cstr, err := d.GetCString(addr)
	if err != nil || len(cstr) == 0 {
		return ""
	}
	if !utils.IsPrintable(cstr) {
		return ""
	}
	if len(cstr) > maxCStringCommentLen {
		return fmt.Sprintf(" ; %#v...", cstr[:maxCStringCommentLen])
	}
	if len(cstr) > 1 {
		return fmt.Sprintf(" ; %#v", cstr)
	}
	return ""
}

func appendComment(base, extra string) string {
	if extra == "" {
		return base
	}
	if base == "" {
		return extra
	}
	return base + extra
}

func Disassemble(d Disass) string {
	var sb strings.Builder
	var instrStr string
	var instrValue uint32
	var decoder disassemble.Decoder
	var prevInstr disassemble.Inst
	var hasPrev bool
	var instructions []disassemble.Instruction

	r := bytes.NewReader(d.Data())

	startAddr := d.StartAddr()

	for {
		err := binary.Read(r, binary.LittleEndian, &instrValue)

		if err == io.EOF {
			break
		}

		if !d.AsJSON() {
			var comment string
			var instruction disassemble.Inst
			if err := decoder.DecomposeInto(startAddr, instrValue, &instruction); err != nil {
				var op string
				var oprs string
				if instrValue == 0xfeedfacf {
					op = ".long"
					oprs = fmt.Sprintf("%#x", instrValue)
					comment = " ; (possible embedded MachO)"
				} else if instrValue == 0x201420 {
					op = "genter"
				} else if instrValue == 0x00201400 {
					op = "gexit"
				} else if instrValue == 0xe7ffdefe || instrValue == 0xe7ffdeff {
					op = "trap"
				} else if instrValue > 0xffff0000 {
					op = ".long"
					oprs = fmt.Sprintf("%#x", instrValue)
					comment = " ; (probably a jump-table)"
				} else if hasPrev && strings.Contains(prevInstr.Operation.String(), "braa") {
					break // TODO: why did I do this again?
				} else if (instrValue & 0xfffffC00) == 0x00201000 {
					Xr := disassemble.Register((instrValue & 0x1F) + 34)
					m := (instrValue >> 5) & 0x1F
					if m == 17 {
						if instrValue&0x1F == 0 {
							op = "amxset"
						} else {
							op = "amxclr"
						}
					} else {
						op = opName(m).String()
						oprs = Xr.String()
					}
				} else if instrValue>>21 == 1 {
					op = ".long"
					oprs = fmt.Sprintf("%#x", instrValue)
					comment = " ; (possible unknown Apple instruction)"
				} else if cstr, err := d.GetCString(startAddr); err == nil {
					op = "DCB"
					if utils.IsASCII(cstr) {
						if len(cstr) > 200 {
							comment = fmt.Sprintf("%#v", cstr[:200])
						} else if len(cstr) > 1 {
							comment = fmt.Sprintf("%#v", cstr)
						}
					}
					// TODO: should I advance startAddr past the end of the cstring ?
					// Otherwise it'll try and disass the rest of the string (that we already printed)
				} else {
					op = ".long"
					oprs = fmt.Sprintf("%#x", instrValue)
					comment = fmt.Sprintf(" ; (%s)", err.Error())
				}

				if d.Color() {
					fmt.Fprintf(&sb, "%s:  %s   %s %s%s\n",
						colorAddr("%#08x", uint64(startAddr)),
						colorOpCodes(disassemble.GetOpCodeByteString(instrValue)),
						colorOp("%-7s", op),
						ColorOperands(" "+oprs),
						colorComment(comment),
					)
				} else {
					fmt.Fprintf(&sb, "%#08x:  %s   %s\t%s%s\n", uint64(startAddr), disassemble.GetOpCodeByteString(instrValue), op, oprs, comment)
				}

				goto INCR_ADDR
			}

			if instrStr, err = instruction.Disassemble(); err != nil {
				if d.Color() {
					fmt.Fprintf(&sb, "%s:  %s   %s %s%s\n",
						colorAddr("%#08x", uint64(startAddr)),
						colorOpCodes(disassemble.GetOpCodeByteString(instrValue)),
						colorOp("%-7s", ".long"),
						ColorOperands(fmt.Sprintf(" %#x", instrValue)),
						colorComment(fmt.Sprintf(" ; (%s)", err.Error())),
					)
				} else {
					fmt.Fprintf(&sb, "%#08x:  %s   .long\t%#x ; (%s)\n", uint64(startAddr), disassemble.GetOpCodeByteString(instrValue), instrValue, err.Error())
				}
				goto INCR_ADDR
			}

			if !d.Quite() {
				// check for start of a new function
				if ok, fname := d.IsFunctionStart(instruction.Address); ok {
					if d.Color() {
						sb.WriteString(colorOp("\n%s:\n", fname))
					} else {
						fmt.Fprintf(&sb, "\n%s:\n", fname)
					}
				} else if name, ok := d.FindSymbol(uint64(instruction.Address)); ok {
					display := symbols.FormatSymbol(name, d.Demangle())
					if d.Color() {
						sb.WriteString(colorOp("\n%s\n", display))
					} else {
						fmt.Fprintf(&sb, "\n%s\n", display)
					}
				}

				if d.IsLocation(instruction.Address) {
					if d.Color() {
						fmt.Fprintf(&sb, "%s\n", colorLocation("loc_%x", instruction.Address))
					} else {
						fmt.Fprintf(&sb, "%#08x:  ; loc_%x\n", instruction.Address, instruction.Address)
					}
				}

				if instruction.Address != startAddr {
					if swiftstr, ok := d.FindSwiftString(instruction.Address); ok {
						comment = fmt.Sprintf(" ; %#v", swiftstr)
					}
				}

				// if ok, imm := triage.HasLoc(i.Instruction.Address()); ok {
				// 	if detail, ok := triage.Details[imm]; ok {
				// 		if triage.IsData(imm) {
				// 			opStr += fmt.Sprintf(" ; %s", detail)
				// 		} else {
				// 			opStr += fmt.Sprintf(" ; %s", detail)
				// 		}
				// 	}
				// }

				if instruction.Operation == disassemble.ARM64_MRS || instruction.Operation == disassemble.ARM64_MSR {
					var ops []string
					replaced := false
					for idx := 0; idx < int(instruction.NumOps); idx++ {
						op := instruction.Operands[idx]
						if op.Class == disassemble.REG {
							if reg, ok := operandRegister(&instruction, idx); ok {
								ops = append(ops, reg.String())
							}
						} else if op.Class == disassemble.IMPLEMENTATION_SPECIFIC {
							if sysRegFix, ok := implSpecificSysReg(&op); ok {
								ops = append(ops, sysRegFix)
								replaced = true
							}
						}
						if replaced {
							instrStr = fmt.Sprintf("%s\t%s", instruction.Operation, strings.Join(ops, ", "))
						}
					}
				} else if ok, loc := d.IsBranchLocation(instruction.Address); ok {
					opStr := strings.TrimPrefix(instrStr, fmt.Sprintf("%s\t", instruction.Operation))
					for idx := 0; idx < int(instruction.NumOps); idx++ {
						operand := instruction.Operands[idx]
						if operand.Class == disassemble.LABEL {
							if name, ok := d.FindSymbol(uint64(operand.Immediate)); ok {
								display := symbols.FormatSymbol(name, d.Demangle())
								opStr = display
								comment = appendComment(comment, cstringComment(d, uint64(operand.Immediate)))
							} else {
								direction := ""
								delta := int(loc) - int(instruction.Address)
								if delta > 0 {
									direction = fmt.Sprintf(" ; ⤵ %#x", delta)
								} else if delta == 0 {
									direction = " ; ∞ loop" // TODO: I should break these out into a comment var like in errors (might speed up colorization)
								} else {
									direction = fmt.Sprintf(" ; ⤴ %#x", delta)
								}
								opStr = strings.Replace(opStr, fmt.Sprintf("%#x", loc), fmt.Sprintf("loc_%x%s", loc, direction), 1)
							}
						}
					}
					instrStr = fmt.Sprintf("%s\t%s", instruction.Operation, opStr)
				} else if instruction.Operation == disassemble.ARM64_BL || instruction.Operation == disassemble.ARM64_B {
					if target, ok := operandImmediate(&instruction, 0); ok {
						if name, ok := d.FindSymbol(target); ok {
							display := symbols.FormatSymbol(name, d.Demangle())
							instrStr = fmt.Sprintf("%s\t%s", instruction.Operation, display)
							comment = appendComment(comment, cstringComment(d, target))
						}
					}
				} else if IsLoadLiteral(&instruction) {
					if name, ok := d.FindSymbol(uint64(instruction.Operands[1].Immediate)); ok {
						display := symbols.FormatSymbol(name, d.Demangle())
						comment = fmt.Sprintf(" ; %s", display)
						comment = appendComment(comment, cstringComment(d, uint64(instruction.Operands[1].Immediate)))
					}
				} else if instruction.Operation == disassemble.ARM64_CBZ || instruction.Operation == disassemble.ARM64_CBNZ {
					if target, ok := operandImmediate(&instruction, 1); ok {
						if name, ok := d.FindSymbol(target); ok {
							display := symbols.FormatSymbol(name, d.Demangle())
							comment = fmt.Sprintf(" ; %s", display)
							comment = appendComment(comment, cstringComment(d, target))
						}
					}
				} else if instruction.Operation == disassemble.ARM64_ADR {
					opStr := strings.TrimPrefix(instrStr, fmt.Sprintf("%s\t", instruction.Operation))
					for idx := 0; idx < int(instruction.NumOps); idx++ {
						operand := instruction.Operands[idx]
						if operand.Class == disassemble.LABEL {
							if name, ok := d.FindSymbol(uint64(operand.Immediate)); ok {
								display := symbols.FormatSymbol(name, d.Demangle())
								opStr = strings.Replace(opStr, fmt.Sprintf("%#x", operand.Immediate), display, 1)
								comment = appendComment(comment, cstringComment(d, uint64(operand.Immediate)))
							} else if cstr, err := d.GetCString(uint64(operand.Immediate)); err == nil {
								if utils.IsPrintable(cstr) {
									if len(cstr) > 200 {
										comment = fmt.Sprintf(" ; %#v...", cstr[:200])
									} else if len(cstr) > 1 {
										comment = fmt.Sprintf(" ; %#v", cstr)
									}
								}
							}
						}
					}
					instrStr = fmt.Sprintf("%s\t%s", instruction.Operation, opStr)
				} else if hasPrev && prevInstr.Operation == disassemble.ARM64_ADRP &&
					(instruction.Operation == disassemble.ARM64_ADD ||
						instruction.Operation == disassemble.ARM64_LDR ||
						instruction.Operation == disassemble.ARM64_LDRB ||
						instruction.Operation == disassemble.ARM64_LDRSW ||
						instruction.Operation == disassemble.ARM64_STRB) {
					adrpRegister, ok := operandRegister(&prevInstr, 0)
					if ok {
						adrpImm, ok := operandImmediate(&prevInstr, 1)
						if ok {
							srcRegister, ok := operandRegister(&instruction, 1)
							if ok {
								validPattern := true
								if adrpRegister == srcRegister {
									switch instruction.Operation {
									case disassemble.ARM64_LDR, disassemble.ARM64_LDRB, disassemble.ARM64_LDRSW, disassemble.ARM64_STRB:
										if imm, ok := operandImmediate(&instruction, 1); ok {
											adrpImm += imm
										} else {
											validPattern = false
										}
									case disassemble.ARM64_ADD:
										if imm, ok := operandImmediate(&instruction, 2); ok {
											adrpImm += imm
										} else {
											validPattern = false
										}
									}
								}
								if validPattern {
									if name, ok := d.FindSymbol(uint64(adrpImm)); ok {
										display := symbols.FormatSymbol(name, d.Demangle())
										comment = fmt.Sprintf(" ; %s", display)
										comment = appendComment(comment, cstringComment(d, uint64(adrpImm)))
									} else if ok, detail := d.IsPointer(adrpImm); ok {
										if name, ok := d.FindSymbol(uint64(detail.Pointer)); ok {
											display := symbols.FormatSymbol(symbols.PrefixPointer+name, d.Demangle())
											comment = fmt.Sprintf(" ; %s", display)
											comment = appendComment(comment, cstringComment(d, uint64(detail.Pointer)))
										} else {
											comment = fmt.Sprintf(" ; %s%x (%s)", symbols.PrefixPointer, detail.Pointer, detail)
										}
									} else if cstr, err := d.GetCString(adrpImm); err == nil && len(cstr) > 0 {
										if utils.IsPrintable(cstr) {
											if len(cstr) > 200 {
												comment = fmt.Sprintf(" ; %#v...", cstr[:200])
											} else if len(cstr) > 1 {
												comment = fmt.Sprintf(" ; %#v", cstr)
											}
										} else { // try again with immediate as pointer
											if ptr, err := d.ReadAddr(adrpImm); err == nil {
												if name, ok := d.FindSymbol(ptr); ok {
													display := symbols.FormatSymbol(symbols.PrefixPointer+name, d.Demangle())
													comment = fmt.Sprintf(" ; %s", display)
													comment = appendComment(comment, cstringComment(d, ptr))
												}
											}
										}
									} else if ok, detail := d.IsData(adrpImm); ok {
										instrStr += fmt.Sprintf(" ; data_%x (%s)", adrpImm, detail)
									}
								}
							}
						}
					}
				}

				if instruction.Encoding == disassemble.ENC_LDR_B_LDST_IMMPRE {
					fmt.Fprintf(&sb, "%s\n", instrStr)
				}
			}

			if d.Middle() != 0 && d.Middle() == startAddr {
				if d.Color() {
					opStr := strings.TrimSpace(strings.TrimPrefix(instrStr, instruction.Operation.String()))
					fmt.Fprintf(&sb, "=>%08x:  %s   %-7s %s%s\n", uint64(startAddr), disassemble.GetOpCodeByteString(instrValue), instruction.Operation, opStr, comment)
				} else {
					fmt.Fprintf(&sb, "=>%08x:  %s\t%s%s\n", uint64(startAddr), disassemble.GetOpCodeByteString(instrValue), instrStr, comment)
				}
			} else {
				if d.Color() {
					opStr := strings.TrimSpace(strings.TrimPrefix(instrStr, instruction.Operation.String()))
					fmt.Fprintf(&sb, "%s:  %s   %s %s%s\n",
						colorAddr("%#08x", uint64(startAddr)),
						colorOpCodes(disassemble.GetOpCodeByteString(instrValue)),
						colorOp("%-7s", instruction.Operation),
						ColorOperands(" "+opStr),
						colorComment(comment),
					)
				} else {
					fmt.Fprintf(&sb, "%#08x:  %s   %s%s\n", uint64(startAddr), disassemble.GetOpCodeByteString(instrValue), instrStr, comment)
				}
			}

			prevInstr = instruction
			hasPrev = true
		} else { // output as JSON
			var instruction disassemble.Inst
			if err := decoder.DecomposeInto(startAddr, instrValue, &instruction); err != nil { // TODO: it would probably be valuable to add a "comment" field and capture the analysis above for JSON peeps
				instructions = append(instructions, invalidInstruction(startAddr, instrValue, err)) // TODO: same with error enhancements above
				goto INCR_ADDR
			}
			converted := instruction.ToInstruction()
			if converted.Disassembly, err = instruction.Disassemble(); err != nil {
				instructions = append(instructions, invalidInstruction(startAddr, instrValue, err))
				goto INCR_ADDR
			}
			instructions = append(instructions, *converted)
		}
	INCR_ADDR:
		startAddr += uint64(binary.Size(uint32(0)))
	}

	if d.AsJSON() {
		var curFunc string
		funcsJSON := make(map[string][]disassemble.Instruction)
		for _, inst := range instructions {
			if ok, fname := d.IsFunctionStart(inst.Address); ok {
				curFunc = fname
				funcsJSON[curFunc] = append(funcsJSON[curFunc], inst)
			} else {
				if len(curFunc) > 0 {
					funcsJSON[curFunc] = append(funcsJSON[curFunc], inst)
				}
			}
		}
		if len(funcsJSON) > 0 {
			if dat, err := json.Marshal(funcsJSON); err == nil {
				fmt.Fprintf(&sb, "%s\n", string(dat))
			}
		} else {
			if dat, err := json.Marshal(instructions); err == nil {
				fmt.Fprintf(&sb, "%s\n", string(dat))
			}
		}
	}

	return sb.String()
}

func ParseGotPtrs(m *macho.File) (map[uint64]uint64, error) {

	gots := make(map[uint64]uint64)

	for _, sec := range m.Sections {
		if slices.Contains([]string{"__AUTH_CONST", "__DATA_CONST"}, sec.Seg) &&
			slices.Contains([]string{"__got", "__auth_got", "__auth_ptr"}, sec.Name) ||
			(sec.Flags.IsNonLazySymbolPointers() || sec.Flags.IsLazySymbolPointers()) {
			dat, err := sec.Data()
			if err != nil {
				off, err := m.GetOffset(sec.Addr)
				if err != nil {
					return nil, fmt.Errorf("failed to get offset for %s.%s section: %v", sec.Seg, sec.Name, err)
				}
				dat = make([]byte, sec.Size)
				if n, err := m.ReadAt(dat, int64(off)); err != nil || n != len(dat) {
					return nil, fmt.Errorf("failed to get %s.%s section data: %v", sec.Seg, sec.Name, err)
				}
			}

			ptrs := make([]uint64, sec.Size/8)
			if err := binary.Read(bytes.NewReader(dat), binary.LittleEndian, &ptrs); err != nil {
				return nil, fmt.Errorf("failed to read %s.%s NonLazySymbol pointers; %v", sec.Seg, sec.Name, err)
			}

			for idx, ptr := range ptrs {
				gots[sec.Addr+uint64(idx*8)] = m.SlidePointer(ptr)
			}
		}
	}

	return gots, nil
}

func ParseStubsForMachO(m *macho.File) (map[uint64]uint64, error) {

	stubs := make(map[uint64]uint64)

	for _, sec := range m.Sections {
		if sec.Flags.IsSymbolStubs() {

			dat, err := sec.Data()
			if err != nil {
				off, err := m.GetOffset(sec.Addr)
				if err != nil {
					return nil, fmt.Errorf("failed to get offset for %s.%s section: %v", sec.Seg, sec.Name, err)
				}
				dat := make([]byte, sec.Size)
				if n, err := m.ReadAt(dat, int64(off)); err != nil || n != len(dat) {
					return nil, fmt.Errorf("failed to get %s.%s section data: %v", sec.Seg, sec.Name, err)
				}
			}

			sb, err := ParseStubsASM(dat, sec.Addr, func(u uint64) (uint64, error) {
				return m.GetPointerAtAddress(m.SlidePointer(u))
			})
			if err != nil {
				return nil, fmt.Errorf("failed to parse %s.%s section stubs data: %v", sec.Seg, sec.Name, err)
			}

			maps.Copy(stubs, sb)
		}
	}

	return stubs, nil
}

// arm64InstrSize is the fixed width of an AArch64 instruction in bytes.
const arm64InstrSize uint64 = 4

// pushStubQueue advances the fixed-depth FIFO of recently decoded instructions
// used by ParseStubsASM: it shifts slot 0 -> 1 -> 2 and inserts inst at slot 0.
// valid marks whether inst decoded successfully (false for the decode-error path,
// where the slot must never be matched against).
func pushStubQueue(queue *[3]disassemble.Inst, queueValid *[3]bool, inst disassemble.Inst, valid bool) {
	queue[2] = queue[1]
	queue[1] = queue[0]
	queue[0] = inst
	queueValid[2] = queueValid[1]
	queueValid[1] = queueValid[0]
	queueValid[0] = valid
}

func ParseStubsASM(data []byte, begin uint64, readPtr func(uint64) (uint64, error)) (map[uint64]uint64, error) {
	var instrValue uint32
	var decoder disassemble.Decoder
	var queue [3]disassemble.Inst
	var queueValid [3]bool

	stubs := make(map[uint64]uint64)

	r := bytes.NewReader(data)

	startAddr := begin

	for {
		err := binary.Read(r, binary.LittleEndian, &instrValue)

		if err == io.EOF {
			break
		}

		var instruction disassemble.Inst
		if err := decoder.DecomposeInto(startAddr, instrValue, &instruction); err != nil {
			startAddr += arm64InstrSize
			pushStubQueue(&queue, &queueValid, instruction, false)
			continue
		}

		if queueValid[1] && queue[1].Operation == disassemble.ARM64_ADRP &&
			queueValid[0] && queue[0].Operation == disassemble.ARM64_ADD &&
			instruction.Operation == disassemble.ARM64_LDR {
			adrpRegister, ok := operandRegister(&queue[1], 0)
			if ok {
				if addRegister, ok := operandRegister(&queue[0], 0); ok {
					if adrpBase, ok := operandImmediate(&queue[1], 1); ok {
						if addImm, ok := operandImmediate(&queue[0], 2); ok {
							if adrpRegister == addRegister {
								// adrp     x17, 0x221baf000
								stubTarget := adrpBase
								// add      x17, x17, #0xbe0
								stubTarget += addImm
								stubAddr := queue[1].Address // address of beginning of stub (adrp)
								// ldr      x16, [x17]
								if ldrRegister, ok := operandRegister(&instruction, 1); ok &&
									(addRegister == ldrRegister) { // check add reg and ldr reg are the same
									if ldrImm, ok := operandImmediate(&instruction, 1); ok && ldrImm != 0 { // check for immediate
										stubTarget += ldrImm
									}
									addr, err := readPtr(stubTarget)
									if err != nil {
										logStubDecodeContext(&instruction, &queue, &queueValid)
										return nil, fmt.Errorf("failed to read stub target pointer at %#x: %v", stubTarget, err)
									}
									// braa     x16, x17
									stubs[stubAddr] = addr
								}
							}
						}
					}
				}
			}
		} else if queueValid[1] && queue[1].Operation == disassemble.ARM64_ADRP &&
			queueValid[0] && queue[0].Operation == disassemble.ARM64_ADD &&
			(IsBranchOp(instruction.Operation) || instruction.Operation == disassemble.ARM64_BR) {
			adrpRegister, ok := operandRegister(&queue[1], 0)
			if ok {
				if addRegister, ok := operandRegister(&queue[0], 0); ok {
					if adrpBase, ok := operandImmediate(&queue[1], 1); ok {
						if addImm, ok := operandImmediate(&queue[0], 2); ok {
							// op2 must be an immediate (`add Xd, Xn, #imm`); a register-form add
							// (`add Xd, Xn, Xm, lsl #n`) reports Immediate=0 and would otherwise
							// record a bogus page-base stub.
							_, addOp2IsReg := operandRegister(&queue[0], 2)
							if !addOp2IsReg && adrpRegister == addRegister {
								// adrp     x16, 0x19089b000
								stubTarget := adrpBase
								// add      x16, x16, #0x94c ; ___stack_chk_fail
								stubTarget += addImm
								stubAddr := queue[1].Address // address of beginning of stub (adrp)
								// br       x16
								stubs[stubAddr] = stubTarget
							}
						}
					}
				}
			}
			// brk      #0x1
		} else if queueValid[0] && queue[0].Operation == disassemble.ARM64_ADRP &&
			instruction.Operation == disassemble.ARM64_LDR {
			if adrpRegister, ok := operandRegister(&queue[0], 0); ok {
				if adrpBase, ok := operandImmediate(&queue[0], 1); ok {
					// adrp	x16, #0x1e3be9000
					stubTarget := adrpBase // #0x1e3be9000
					// ldr	x16, [x16, #0x560]
					if ldrRegister, ok := operandRegister(&instruction, 0); ok && adrpRegister == ldrRegister {
						if ldrImm, ok := operandImmediate(&instruction, 1); ok {
							stubTarget += ldrImm
							stubAddr := queue[0].Address
							addr, err := readPtr(stubTarget)
							if err != nil {
								logStubDecodeContext(&instruction, &queue, &queueValid)
								return nil, fmt.Errorf("failed to read stub target pointer at %#x: %v", stubTarget, err)
							}
							stubs[stubAddr] = addr
						}
					}
				}
			}
		} else if queueValid[2] && queue[2].Operation == disassemble.ARM64_ADR &&
			queueValid[1] && queue[1].Operation == disassemble.ARM64_MOV &&
			queueValid[0] && queue[0].Operation == disassemble.ARM64_ADD &&
			instruction.Operation == disassemble.ARM64_BR {
			// iOS 27+ far-target arithmetic stub (no GOT load, no PAC):
			//   adr  x16, <base>            ; byte-exact base (ADR, not page-granular ADRP)
			//   mov  x17, #<imm>            ; immediate, not a register move
			//   add  x16, x16, x17, lsl #<shift>
			//   br   x16                    ; unconditional indirect tail-call only
			// target = base + (imm << shift), reaching far/forward beyond ADRP+ADD's ±4 GiB.
			baseReg, baseOK := operandRegister(&queue[2], 0)
			movReg, movOK := operandRegister(&queue[1], 0)
			_, movSrcIsReg := operandRegister(&queue[1], 1) // reject `mov Xd, Xn` (ORR alias); we need `mov Xd, #imm`
			addReg, addOK := operandRegister(&queue[0], 0)
			addSrcReg, addSrcOK := operandRegister(&queue[0], 1)
			addShiftReg, addShiftOK := operandRegister(&queue[0], 2) // op2 must be a register (the shifted reg)
			brReg, brOK := operandRegister(&instruction, 0)
			if baseOK && movOK && !movSrcIsReg && addOK && addSrcOK && addShiftOK && brOK &&
				baseReg == addReg && addReg == addSrcReg && addReg == brReg && // base/add/br share the destination reg
				movReg == addShiftReg { // add shifts the reg loaded by mov
				if base, ok := operandImmediate(&queue[2], 1); ok {
					if imm, ok := operandImmediate(&queue[1], 1); ok {
						if shift, ok := operandLeftShift(&queue[0], 2); ok { // LSL only
							// adr x16, <base>  (queue[2] is the start of the stub)
							stubs[queue[2].Address] = base + (imm << shift)
						}
					}
				}
			}
		}

		// fmt.Printf("%#08x:  %s\t%s\n", instruction.Address, disassemble.GetOpCodeByteString(instrValue), instruction)

		pushStubQueue(&queue, &queueValid, instruction, true)
		startAddr += arm64InstrSize
	}

	return stubs, nil
}

func ParseHelpersASM(m *macho.File) (map[uint64]uint64, error) {
	var instrValue uint32
	var decoder disassemble.Decoder

	helpers := make(map[uint64]uint64)

	if sec := m.Section("__TEXT", "__stub_helper"); sec != nil {
		dat, err := sec.Data()
		if err != nil {
			off, err := m.GetOffset(sec.Addr)
			if err != nil {
				return nil, fmt.Errorf("failed to get offset for %s.%s section: %v", sec.Seg, sec.Name, err)
			}
			dat := make([]byte, sec.Size)
			if n, err := m.ReadAt(dat, int64(off)); err != nil || n != len(dat) {
				return nil, fmt.Errorf("failed to get %s.%s section data: %v", sec.Seg, sec.Name, err)
			}
		}

		startAddr := sec.Addr
		stubHelperFnStart := startAddr

		r := bytes.NewReader(dat)

		for {
			err := binary.Read(r, binary.LittleEndian, &instrValue)

			if err == io.EOF {
				break
			}

			var instruction disassemble.Inst
			if err := decoder.DecomposeInto(startAddr, instrValue, &instruction); err != nil {
				startAddr += uint64(binary.Size(uint32(0)))
				continue
			}

			if instruction.Encoding == disassemble.ENC_BR_64_BRANCH_REG { // check if branch location is a function
				startAddr += uint64(binary.Size(uint32(0)))
				stubHelperFnStart = startAddr
				// fmt.Printf("%#08x:  %s\t%s\n", instruction.Address, disassemble.GetOpCodeByteString(instrValue), instruction)
				// fmt.Println()
				continue
			}
			if instruction.Encoding == disassemble.ENC_BL_ONLY_BRANCH_IMM {
				if target, ok := operandImmediate(&instruction, 0); ok {
					helpers[stubHelperFnStart] = target
				}
			}
			// fmt.Printf("%#08x:  %s\t%s\n", instruction.Address, disassemble.GetOpCodeByteString(instrValue), instruction)
			startAddr += uint64(binary.Size(uint32(0)))
		}

		return helpers, nil
	}

	return nil, fmt.Errorf("dylib does NOT contain __TEXT.__stub_helper section: %w", macho.ErrMachOSectionNotFound)
}
