#!/usr/bin/env python3
"""
Patch ARM64/ARM64e unconditional branches across LC_FILESET_ENTRY members in a
little-endian 64-bit Mach-O fileset.

The tool parses the outer Mach-O, locates LC_FILESET_ENTRY load commands,
parses the embedded Mach-Os, resolves locations such as:

    __TEXT_EXEC:__text+0x1000
    __TEXT_EXEC+0x200
    +0x1234
    0x1234
    symbol:_my_function
    symbol:_my_function+0x10

and overwrites the source instruction with an AArch64 unconditional branch
(`B`) or branch-with-link (`BL`) to the target virtual address.

Symbol mode uses the original standalone Mach-O for a fileset entry to find the
symbol's segment/section/offset there, then re-resolves that location inside the
embedded fileset entry. This is useful when the fileset builder strips symbols
and/or reorders segments inside the outer file.

Examples:

    # Show fileset entries
    ./macho_fileset_branch_v2.py list kernelcache

    # Resolve a symbol through its original standalone Mach-O into the fileset
    ./macho_fileset_branch_v2.py lookup-symbol kernelcache \
        --original 'com.example.file_a=/tmp/file_a.macho' \
        --lookup 'com.example.file_a symbol:_patch_here+0x10'

    # Patch one branch and write a new file
    ./macho_fileset_branch_v2.py patch kernelcache \
        --patch 'com.apple.kernel __TEXT_EXEC:__text+0x1000 com.apple.iokit.IOUSBHostFamily __TEXT_EXEC:__text+0x2000' \
        --output kernelcache.patched

    # Patch using symbols from original standalone files
    ./macho_fileset_branch_v2.py patch kernelcache \
        --original 'com.example.file_a=/tmp/file_a.macho' \
        --original 'com.example.file_b=/tmp/file_b.macho' \
        --patch 'com.example.file_a symbol:_patch_here com.example.file_b symbol:_target+0x10' \
        --output kernelcache.patched
"""

from __future__ import annotations

import argparse
import pathlib
import shlex
import struct
import sys
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence, Tuple


MH_MAGIC_64 = 0xFEEDFACF
FAT_MAGIC = 0xCAFEBABE
FAT_MAGIC_64 = 0xCAFEBABF
MH_CIGAM_64 = 0xCFFAEDFE
MH_CIGAM = 0xCEFAEDFE
FAT_CIGAM = 0xBEBAFECA
FAT_CIGAM_64 = 0xBFBAFECA

CPU_TYPE_ARM64 = 0x0100000C
LC_REQ_DYLD = 0x80000000
LC_SEGMENT_64 = 0x19
LC_SYMTAB = 0x2
LC_FILESET_ENTRY_BASE = 0x35
LC_FILESET_ENTRY = LC_FILESET_ENTRY_BASE | LC_REQ_DYLD

MACH_HEADER_64_SIZE = 32
LOAD_COMMAND_SIZE = 8
SEGMENT_COMMAND_64_SIZE = 72
SECTION_64_SIZE = 80
FILESET_ENTRY_COMMAND_SIZE = 32
SYMTAB_COMMAND_SIZE = 24
NLIST_64_SIZE = 16

N_STAB = 0xE0
N_TYPE = 0x0E
N_UNDF = 0x00
N_ABS = 0x02
N_SECT = 0x0E
N_PBUD = 0x0C
N_INDR = 0x0A


class MachOError(Exception):
    pass


@dataclass
class Section64:
    sectname: str
    segname: str
    addr: int
    size: int
    offset: int
    align: int
    reloff: int
    nreloc: int
    flags: int
    reserved1: int
    reserved2: int
    reserved3: int


@dataclass
class Segment64:
    segname: str
    vmaddr: int
    vmsize: int
    fileoff: int
    filesize: int
    maxprot: int
    initprot: int
    nsects: int
    flags: int
    sections: List[Section64] = field(default_factory=list)

    def contains_va(self, va: int) -> bool:
        return self.vmaddr <= va < (self.vmaddr + self.vmsize)

    def contains_filebacked_va(self, va: int) -> bool:
        return self.vmaddr <= va < (self.vmaddr + self.filesize)


@dataclass
class MachHeader64:
    magic: int
    cputype: int
    cpusubtype: int
    filetype: int
    ncmds: int
    sizeofcmds: int
    flags: int
    reserved: int


@dataclass
class LocationSpec:
    raw: str
    kind: str
    segment: Optional[str] = None
    section: Optional[str] = None
    addend: int = 0
    symbol: Optional[str] = None

    def describe(self) -> str:
        if self.kind == "symbol":
            base = f"symbol {self.symbol}"
        elif self.segment is None:
            base = "fileset base"
        elif self.section is None:
            base = f"segment {self.segment}"
        else:
            base = f"section {self.segment}:{self.section}"

        if self.addend:
            sign = "+" if self.addend >= 0 else "-"
            return f"{base} {sign} 0x{abs(self.addend):x}"
        return base

    def to_expr(self) -> str:
        if self.kind == "symbol":
            if self.addend:
                sign = "+" if self.addend >= 0 else "-"
                return f"symbol:{self.symbol}{sign}0x{abs(self.addend):x}"
            return f"symbol:{self.symbol}"

        if self.segment is None:
            if self.addend == 0:
                return "+0x0"
            sign = "+" if self.addend >= 0 else "-"
            return f"{sign}0x{abs(self.addend):x}"

        if self.section is None:
            if self.addend:
                sign = "+" if self.addend >= 0 else "-"
                return f"{self.segment}{sign}0x{abs(self.addend):x}"
            return self.segment

        if self.addend:
            sign = "+" if self.addend >= 0 else "-"
            return f"{self.segment}:{self.section}{sign}0x{abs(self.addend):x}"
        return f"{self.segment}:{self.section}"


@dataclass
class MachOImage:
    root_offset: int
    header: MachHeader64
    segments: List[Segment64]
    sections_by_index: List[Section64]
    fileoff_mode: str

    def find_segment(self, name: str) -> Segment64:
        for seg in self.segments:
            if seg.segname == name:
                return seg
        raise MachOError(f"segment {name!r} not found")

    def find_section(self, segname: str, sectname: str) -> Section64:
        seg = self.find_segment(segname)
        for sec in seg.sections:
            if sec.sectname == sectname:
                return sec
        raise MachOError(f"section {segname}:{sectname!r} not found")

    def find_section_by_ordinal(self, index: int) -> Section64:
        if not (1 <= index <= len(self.sections_by_index)):
            raise MachOError(
                f"section index {index} is out of range (1..{len(self.sections_by_index)})"
            )
        return self.sections_by_index[index - 1]

    def section_containing_va(self, va: int) -> Optional[Section64]:
        for sec in self.sections_by_index:
            if sec.addr <= va < (sec.addr + sec.size):
                return sec
        return None

    def segment_containing_va(self, va: int) -> Optional[Segment64]:
        for seg in self.segments:
            if seg.contains_va(va):
                return seg
        return None

    def va_to_fileoff(self, va: int) -> int:
        for seg in self.segments:
            if not seg.contains_va(va):
                continue
            delta = va - seg.vmaddr
            if delta >= seg.filesize:
                raise MachOError(
                    f"VA 0x{va:x} lands in zero-fill part of segment {seg.segname} "
                    f"(delta 0x{delta:x} >= filesize 0x{seg.filesize:x})"
                )
            base = 0 if self.fileoff_mode == "absolute" else self.root_offset
            return base + seg.fileoff + delta
        raise MachOError(f"VA 0x{va:x} does not map to any segment")


@dataclass
class FilesetEntry:
    entry_id: str
    vmaddr: int
    fileoff: int
    cmd_off: int
    macho: MachOImage

    def resolve_location(self, spec: LocationSpec) -> Tuple[int, int, str]:
        if spec.kind == "symbol":
            raise MachOError("symbol locations must be normalized before resolve_location()")

        if spec.segment is None:
            base_va = self.vmaddr
            region_size = None
            base_desc = f"fileset base of {self.entry_id}"
        else:
            seg = self.macho.find_segment(spec.segment)
            if spec.section is None:
                base_va = seg.vmaddr
                region_size = seg.vmsize
                base_desc = f"segment {spec.segment} in {self.entry_id}"
            else:
                sec = self.macho.find_section(spec.segment, spec.section)
                base_va = sec.addr
                region_size = sec.size
                base_desc = f"section {spec.segment}:{spec.section} in {self.entry_id}"

        if spec.addend < 0:
            raise MachOError("negative offsets are not supported")

        if region_size is not None and spec.addend >= region_size:
            raise MachOError(
                f"offset 0x{spec.addend:x} is outside {base_desc} "
                f"(size 0x{region_size:x})"
            )

        va = base_va + spec.addend
        fileoff = self.macho.va_to_fileoff(va)
        desc = f"{base_desc} + 0x{spec.addend:x}" if spec.addend else base_desc
        return va, fileoff, desc


@dataclass
class OuterMachO:
    path: pathlib.Path
    data: bytearray
    header: MachHeader64
    entries: List[FilesetEntry]

    def get_entry(self, entry_id: str) -> FilesetEntry:
        matches = [e for e in self.entries if e.entry_id == entry_id]
        if not matches:
            known = ", ".join(sorted(e.entry_id for e in self.entries))
            raise MachOError(f"fileset entry {entry_id!r} not found. Known entries: {known}")
        if len(matches) > 1:
            raise MachOError(f"fileset entry name {entry_id!r} is not unique")
        return matches[0]


@dataclass
class SymtabInfo:
    symoff: int
    nsyms: int
    stroff: int
    strsize: int


@dataclass
class SymbolInfo:
    name: str
    value: int
    n_type: int
    n_sect: int
    n_desc: int
    segment: Optional[Segment64]
    section: Optional[Section64]
    segment_offset: Optional[int]
    section_offset: Optional[int]

    def is_locatable(self) -> bool:
        return self.segment is not None and self.segment_offset is not None

    def location_spec(self, addend: int) -> LocationSpec:
        if not self.is_locatable():
            raise MachOError(f"symbol {self.name!r} is not locatable inside a segment")

        if self.section is not None and self.section_offset is not None:
            return LocationSpec(
                raw=f"symbol:{self.name}",
                kind="range",
                segment=self.section.segname,
                section=self.section.sectname,
                addend=self.section_offset + addend,
            )

        return LocationSpec(
            raw=f"symbol:{self.name}",
            kind="range",
            segment=self.segment.segname,
            section=None,
            addend=self.segment_offset + addend,
        )

    def location_description(self, addend: int) -> str:
        if self.section is not None and self.section_offset is not None:
            base = (
                f"section {self.section.segname}:{self.section.sectname} + "
                f"0x{self.section_offset:x}"
            )
        elif self.segment is not None and self.segment_offset is not None:
            base = f"segment {self.segment.segname} + 0x{self.segment_offset:x}"
        else:
            base = "non-locatable"

        if addend:
            base += f", then + 0x{addend:x}"
        return base


@dataclass
class OriginalMachO:
    path: pathlib.Path
    image: MachOImage
    symbols_by_name: Dict[str, List[SymbolInfo]]

    def lookup_symbol(self, name: str) -> SymbolInfo:
        attempts = [name]
        if name.startswith("_"):
            attempts.append(name[1:])
        else:
            attempts.append("_" + name)

        found_name = None
        candidates: List[SymbolInfo] = []
        for candidate_name in attempts:
            current = self.symbols_by_name.get(candidate_name)
            if current:
                found_name = candidate_name
                candidates = current
                break

        if not candidates:
            raise MachOError(f"symbol {name!r} not found in {self.path}")

        locatable = [sym for sym in candidates if sym.is_locatable()]
        if not locatable:
            raise MachOError(
                f"symbol {found_name!r} exists in {self.path}, but it does not resolve to a segment/section"
            )

        unique_by_location: Dict[Tuple[int, int, int], SymbolInfo] = {}
        for sym in locatable:
            segname = sym.segment.segname if sym.segment is not None else ""
            sectname = sym.section.sectname if sym.section is not None else ""
            key = (sym.value, segname, sectname)
            unique_by_location[key] = sym

        unique = list(unique_by_location.values())
        if len(unique) == 1:
            return unique[0]

        rendered = ", ".join(
            f"0x{sym.value:x}"
            + (
                f" ({sym.section.segname}:{sym.section.sectname}+0x{sym.section_offset:x})"
                if sym.section is not None and sym.section_offset is not None
                else f" ({sym.segment.segname}+0x{sym.segment_offset:x})"
            )
            for sym in unique[:8]
        )
        raise MachOError(
            f"symbol {found_name!r} is ambiguous in {self.path}; candidates: {rendered}"
        )


@dataclass
class ResolvedLocation:
    normalized: LocationSpec
    via: str


@dataclass
class PatchSpec:
    src_entry: str
    src_loc: LocationSpec
    dst_entry: str
    dst_loc: LocationSpec


@dataclass
class AppliedPatch:
    spec: PatchSpec
    src_va: int
    src_fileoff: int
    dst_va: int
    dst_fileoff: int
    insn: int
    original_insn: int
    src_desc: str
    dst_desc: str


@dataclass
class LookupResult:
    entry_id: str
    requested: LocationSpec
    normalized: LocationSpec
    via: str
    va: int
    fileoff: int
    desc: str
    word32: Optional[int]


def read_exact(data: bytes, off: int, size: int) -> bytes:
    end = off + size
    if off < 0 or size < 0 or end > len(data):
        raise MachOError(f"read out of bounds at 0x{off:x} size 0x{size:x}")
    return data[off:end]


def parse_cstr_from_command(cmd_data: bytes, off: int) -> str:
    if off >= len(cmd_data):
        raise MachOError("string offset outside load command")
    raw = cmd_data[off:]
    nul = raw.find(b"\x00")
    if nul == -1:
        raise MachOError("unterminated string inside load command")
    return raw[:nul].decode("utf-8", errors="replace")


def trim_name(raw: bytes) -> str:
    return raw.split(b"\x00", 1)[0].decode("ascii", errors="replace")


def parse_mach_header_64(data: bytes, off: int) -> MachHeader64:
    hdr = read_exact(data, off, MACH_HEADER_64_SIZE)
    fields = struct.unpack("<IIIIIIII", hdr)
    magic = fields[0]
    if magic == MH_CIGAM_64:
        raise MachOError("big-endian Mach-O is not supported")
    if magic == MH_CIGAM:
        raise MachOError("32-bit big-endian Mach-O is not supported")
    if magic in (FAT_MAGIC, FAT_MAGIC_64, FAT_CIGAM, FAT_CIGAM_64):
        raise MachOError("fat/universal Mach-O is not supported; pass a flat little-endian file")
    if magic != MH_MAGIC_64:
        raise MachOError(f"unsupported Mach-O magic 0x{magic:08x} at offset 0x{off:x}")
    header = MachHeader64(*fields)
    if header.cputype != CPU_TYPE_ARM64:
        raise MachOError(f"unsupported CPU type 0x{header.cputype:x}; expected ARM64/ARM64e")
    return header


def parse_segments(data: bytes, macho_off: int, header: MachHeader64) -> Tuple[List[Segment64], List[Section64]]:
    segments: List[Segment64] = []
    sections_by_index: List[Section64] = []
    cmd_off = macho_off + MACH_HEADER_64_SIZE
    end_of_cmds = cmd_off + header.sizeofcmds

    if end_of_cmds > len(data):
        raise MachOError("load commands run beyond end of file")

    for _ in range(header.ncmds):
        cmd, cmdsize = struct.unpack_from("<II", data, cmd_off)
        if cmdsize < LOAD_COMMAND_SIZE:
            raise MachOError(f"invalid load command size 0x{cmdsize:x} at 0x{cmd_off:x}")
        if cmd_off + cmdsize > len(data):
            raise MachOError("load command extends beyond end of file")
        if cmd == LC_SEGMENT_64:
            if cmdsize < SEGMENT_COMMAND_64_SIZE:
                raise MachOError(f"short LC_SEGMENT_64 at 0x{cmd_off:x}")
            seg_fields = struct.unpack_from("<16sQQQQiiII", data, cmd_off + LOAD_COMMAND_SIZE)
            segname = trim_name(seg_fields[0])
            vmaddr, vmsize, fileoff, filesize, maxprot, initprot, nsects, flags = seg_fields[1:]
            sec_off = cmd_off + SEGMENT_COMMAND_64_SIZE
            sections: List[Section64] = []
            expected_cmdsize = SEGMENT_COMMAND_64_SIZE + (nsects * SECTION_64_SIZE)
            if cmdsize < expected_cmdsize:
                raise MachOError(
                    f"LC_SEGMENT_64 {segname} too small for {nsects} sections "
                    f"(cmdsize 0x{cmdsize:x}, expected >= 0x{expected_cmdsize:x})"
                )
            for _sec_idx in range(nsects):
                sec_fields = struct.unpack_from("<16s16sQQIIIIIIII", data, sec_off)
                sec = Section64(
                    sectname=trim_name(sec_fields[0]),
                    segname=trim_name(sec_fields[1]),
                    addr=sec_fields[2],
                    size=sec_fields[3],
                    offset=sec_fields[4],
                    align=sec_fields[5],
                    reloff=sec_fields[6],
                    nreloc=sec_fields[7],
                    flags=sec_fields[8],
                    reserved1=sec_fields[9],
                    reserved2=sec_fields[10],
                    reserved3=sec_fields[11],
                )
                sections.append(sec)
                sections_by_index.append(sec)
                sec_off += SECTION_64_SIZE
            segments.append(
                Segment64(
                    segname=segname,
                    vmaddr=vmaddr,
                    vmsize=vmsize,
                    fileoff=fileoff,
                    filesize=filesize,
                    maxprot=maxprot,
                    initprot=initprot,
                    nsects=nsects,
                    flags=flags,
                    sections=sections,
                )
            )
        cmd_off += cmdsize

    return segments, sections_by_index


def find_symtab(data: bytes, macho_off: int, header: MachHeader64) -> Optional[SymtabInfo]:
    cmd_off = macho_off + MACH_HEADER_64_SIZE
    for _ in range(header.ncmds):
        cmd, cmdsize = struct.unpack_from("<II", data, cmd_off)
        if cmdsize < LOAD_COMMAND_SIZE:
            raise MachOError(f"invalid load command size 0x{cmdsize:x} at 0x{cmd_off:x}")
        if cmd == LC_SYMTAB:
            if cmdsize < SYMTAB_COMMAND_SIZE:
                raise MachOError(f"short LC_SYMTAB at 0x{cmd_off:x}")
            symoff, nsyms, stroff, strsize = struct.unpack_from("<IIII", data, cmd_off + LOAD_COMMAND_SIZE)
            return SymtabInfo(symoff=symoff, nsyms=nsyms, stroff=stroff, strsize=strsize)
        cmd_off += cmdsize
    return None


def read_cstr_from_table(data: bytes, start: int, size: int, off: int) -> str:
    if off >= size:
        raise MachOError(f"string-table offset 0x{off:x} is outside table of size 0x{size:x}")
    raw = data[start + off : start + size]
    nul = raw.find(b"\x00")
    if nul == -1:
        raise MachOError("unterminated string in Mach-O string table")
    return raw[:nul].decode("utf-8", errors="replace")


def detect_inner_fileoff_mode(entry_fileoff: int, segments: Sequence[Segment64], forced: str) -> str:
    if forced in ("absolute", "relative"):
        return forced

    filebacked = [seg for seg in segments if seg.filesize > 0]
    if not filebacked:
        return "relative"

    min_fileoff = min(seg.fileoff for seg in filebacked)
    if min_fileoff == 0:
        return "relative"
    if any(seg.fileoff < entry_fileoff for seg in filebacked):
        return "relative"
    return "absolute"


def parse_macho_image_at(data: bytes, macho_off: int, fileoff_mode: str) -> MachOImage:
    header = parse_mach_header_64(data, macho_off)
    segments, sections_by_index = parse_segments(data, macho_off, header)
    return MachOImage(
        root_offset=macho_off,
        header=header,
        segments=segments,
        sections_by_index=sections_by_index,
        fileoff_mode=fileoff_mode,
    )


def parse_outer_macho(path: pathlib.Path, fileoff_mode: str) -> OuterMachO:
    data = bytearray(path.read_bytes())
    header = parse_mach_header_64(data, 0)

    entries_meta: List[Tuple[str, int, int, int]] = []
    cmd_off = MACH_HEADER_64_SIZE
    for _ in range(header.ncmds):
        cmd, cmdsize = struct.unpack_from("<II", data, cmd_off)
        if cmdsize < LOAD_COMMAND_SIZE:
            raise MachOError(f"invalid load command size 0x{cmdsize:x} at 0x{cmd_off:x}")
        cmd_data = read_exact(data, cmd_off, cmdsize)

        if (cmd & ~LC_REQ_DYLD) == LC_FILESET_ENTRY_BASE:
            if cmdsize < FILESET_ENTRY_COMMAND_SIZE:
                raise MachOError(f"short LC_FILESET_ENTRY at 0x{cmd_off:x}")
            vmaddr, fileoff, entry_id_off, reserved = struct.unpack_from("<QQII", cmd_data, LOAD_COMMAND_SIZE)
            _ = reserved
            entry_id = parse_cstr_from_command(cmd_data, entry_id_off)
            entries_meta.append((entry_id, vmaddr, fileoff, cmd_off))

        cmd_off += cmdsize

    if not entries_meta:
        raise MachOError("no LC_FILESET_ENTRY load commands found")

    entries: List[FilesetEntry] = []
    for entry_id, vmaddr, fileoff, lc_off in entries_meta:
        inner_header = parse_mach_header_64(data, fileoff)
        inner_segments, inner_sections = parse_segments(data, fileoff, inner_header)
        mode = detect_inner_fileoff_mode(fileoff, inner_segments, fileoff_mode)
        entries.append(
            FilesetEntry(
                entry_id=entry_id,
                vmaddr=vmaddr,
                fileoff=fileoff,
                cmd_off=lc_off,
                macho=MachOImage(
                    root_offset=fileoff,
                    header=inner_header,
                    segments=inner_segments,
                    sections_by_index=inner_sections,
                    fileoff_mode=mode,
                ),
            )
        )

    return OuterMachO(path=path, data=data, header=header, entries=entries)


def parse_original_macho(path: pathlib.Path) -> OriginalMachO:
    data = path.read_bytes()
    image = parse_macho_image_at(data, 0, "relative")
    symtab = find_symtab(data, 0, image.header)
    if symtab is None:
        raise MachOError(f"{path} does not contain an LC_SYMTAB")

    if symtab.symoff + symtab.nsyms * NLIST_64_SIZE > len(data):
        raise MachOError(f"symbol table in {path} extends beyond end of file")
    if symtab.stroff + symtab.strsize > len(data):
        raise MachOError(f"string table in {path} extends beyond end of file")

    symbols_by_name: Dict[str, List[SymbolInfo]] = {}
    for idx in range(symtab.nsyms):
        off = symtab.symoff + idx * NLIST_64_SIZE
        n_strx, n_type, n_sect, n_desc, n_value = struct.unpack_from("<IBBhQ", data, off)

        if n_type & N_STAB:
            continue
        if n_strx == 0:
            continue

        name = read_cstr_from_table(data, symtab.stroff, symtab.strsize, n_strx)
        if not name:
            continue

        segment: Optional[Segment64] = None
        section: Optional[Section64] = None
        segment_offset: Optional[int] = None
        section_offset: Optional[int] = None

        if n_value != 0:
            if n_sect != 0 and (n_type & N_TYPE) == N_SECT:
                try:
                    section = image.find_section_by_ordinal(n_sect)
                    segment = image.find_segment(section.segname)
                    if section.addr <= n_value < (section.addr + section.size):
                        section_offset = n_value - section.addr
                        segment_offset = n_value - segment.vmaddr
                    elif segment.contains_va(n_value):
                        # Some symbols may still be tied to a section ordinal even if
                        # their exact value lands just outside the section's recorded
                        # byte range. In that case, keep the segment-relative location.
                        section = None
                        segment_offset = n_value - segment.vmaddr
                    else:
                        section = None
                        segment = None
                except MachOError:
                    section = None
                    segment = None

            if segment is None:
                section = image.section_containing_va(n_value)
                if section is not None:
                    segment = image.find_segment(section.segname)
                    section_offset = n_value - section.addr
                    segment_offset = n_value - segment.vmaddr
                else:
                    segment = image.segment_containing_va(n_value)
                    if segment is not None:
                        segment_offset = n_value - segment.vmaddr

        sym = SymbolInfo(
            name=name,
            value=n_value,
            n_type=n_type,
            n_sect=n_sect,
            n_desc=n_desc,
            segment=segment,
            section=section,
            segment_offset=segment_offset,
            section_offset=section_offset,
        )
        symbols_by_name.setdefault(name, []).append(sym)

    return OriginalMachO(path=path, image=image, symbols_by_name=symbols_by_name)


def parse_int(text: str) -> int:
    try:
        return int(text, 0)
    except ValueError as exc:
        raise MachOError(f"invalid integer {text!r}") from exc


def split_symbol_and_addend(text: str) -> Tuple[str, int]:
    body = text[len("symbol:") :]
    if not body:
        raise MachOError(f"invalid location spec {text!r}: empty symbol name")

    for idx in range(len(body) - 1, 0, -1):
        ch = body[idx]
        if ch not in "+-":
            continue
        suffix = body[idx + 1 :].strip()
        if not suffix:
            continue
        try:
            addend = parse_int(suffix)
        except MachOError:
            continue
        symbol = body[:idx].strip()
        if not symbol:
            raise MachOError(f"invalid location spec {text!r}: empty symbol name")
        return symbol, addend if ch == "+" else -addend

    return body.strip(), 0


def parse_location_spec(text: str) -> LocationSpec:
    s = text.strip()
    if not s:
        return LocationSpec(raw=text, kind="range", segment=None, section=None, addend=0)

    if s.startswith("symbol:"):
        symbol, addend = split_symbol_and_addend(s)
        return LocationSpec(raw=text, kind="symbol", symbol=symbol, addend=addend)

    if ":" not in s and s[0] in "+-0123456789":
        return LocationSpec(raw=text, kind="range", segment=None, section=None, addend=parse_int(s))

    addend = 0
    main = s
    if "+" in s:
        main, add = s.rsplit("+", 1)
        addend = parse_int(add.strip())
        main = main.strip()

    if not main:
        return LocationSpec(raw=text, kind="range", segment=None, section=None, addend=addend)

    if ":" in main:
        seg, sec = main.split(":", 1)
        seg = seg.strip()
        sec = sec.strip()
        if not seg:
            raise MachOError(f"invalid location spec {text!r}: empty segment name")
        if not sec:
            raise MachOError(f"invalid location spec {text!r}: empty section name")
        return LocationSpec(raw=text, kind="range", segment=seg, section=sec, addend=addend)

    return LocationSpec(raw=text, kind="range", segment=main, section=None, addend=addend)


def parse_patch_spec(text: str) -> PatchSpec:
    try:
        parts = shlex.split(text)
    except ValueError as exc:
        raise MachOError(f"failed to parse --patch spec {text!r}: {exc}") from exc
    if len(parts) != 4:
        raise MachOError(
            "--patch must have exactly four tokens: "
            "<source-entry> <source-loc> <target-entry> <target-loc>"
        )
    return PatchSpec(
        src_entry=parts[0],
        src_loc=parse_location_spec(parts[1]),
        dst_entry=parts[2],
        dst_loc=parse_location_spec(parts[3]),
    )


def parse_lookup_spec(text: str) -> Tuple[str, LocationSpec]:
    try:
        parts = shlex.split(text)
    except ValueError as exc:
        raise MachOError(f"failed to parse --lookup spec {text!r}: {exc}") from exc
    if len(parts) != 2:
        raise MachOError(
            "--lookup must have exactly two tokens: <entry> <loc>. "
            "Example: --lookup 'com.example.file symbol:_my_symbol+0x10'"
        )
    return parts[0], parse_location_spec(parts[1])


def parse_original_mapping(text: str) -> Tuple[str, pathlib.Path]:
    if "=" not in text:
        raise MachOError(
            f"invalid --original mapping {text!r}; expected '<fileset-entry>=<path-to-original-macho>'"
        )
    entry, path = text.split("=", 1)
    entry = entry.strip()
    path = path.strip()
    if not entry:
        raise MachOError(f"invalid --original mapping {text!r}: empty fileset entry name")
    if not path:
        raise MachOError(f"invalid --original mapping {text!r}: empty path")
    return entry, pathlib.Path(path)


def load_original_mappings(mappings: Sequence[str]) -> Dict[str, OriginalMachO]:
    loaded: Dict[str, OriginalMachO] = {}
    for item in mappings:
        entry, path = parse_original_mapping(item)
        loaded[entry] = parse_original_macho(path)
    return loaded


def normalize_location_spec(
    entry: FilesetEntry,
    spec: LocationSpec,
    originals: Mapping[str, OriginalMachO],
) -> ResolvedLocation:
    if spec.kind != "symbol":
        return ResolvedLocation(normalized=spec, via=spec.describe())

    original = originals.get(entry.entry_id)
    if original is None:
        raise MachOError(
            f"location {spec.raw!r} for fileset entry {entry.entry_id!r} requires an original Mach-O. "
            f"Provide --original '{entry.entry_id}=PATH'"
        )

    if spec.addend < 0:
        raise MachOError("negative symbol addends are not supported")

    sym = original.lookup_symbol(spec.symbol or "")
    normalized = sym.location_spec(spec.addend)
    via = (
        f"{spec.raw} -> {normalized.describe()} "
        f"(resolved via {sym.name} in {original.path}; {sym.location_description(spec.addend)})"
    )
    return ResolvedLocation(normalized=normalized, via=via)


def resolve_lookup(
    outer: OuterMachO,
    entry_id: str,
    loc: LocationSpec,
    originals: Mapping[str, OriginalMachO],
) -> LookupResult:
    entry = outer.get_entry(entry_id)
    resolved = normalize_location_spec(entry, loc, originals)
    va, fileoff, desc = entry.resolve_location(resolved.normalized)
    word32 = None
    if fileoff + 4 <= len(outer.data):
        word32 = struct.unpack_from("<I", outer.data, fileoff)[0]
    return LookupResult(
        entry_id=entry_id,
        requested=loc,
        normalized=resolved.normalized,
        via=resolved.via,
        va=va,
        fileoff=fileoff,
        desc=desc,
        word32=word32,
    )


def encode_aarch64_branch(src_va: int, dst_va: int, with_link: bool) -> int:
    if (src_va & 3) != 0:
        raise MachOError(f"source VA 0x{src_va:x} is not 4-byte aligned")
    if (dst_va & 3) != 0:
        raise MachOError(f"target VA 0x{dst_va:x} is not 4-byte aligned")

    delta = dst_va - src_va
    if (delta & 3) != 0:
        raise MachOError(
            f"branch delta 0x{delta:x} is not 4-byte aligned; cannot encode as AArch64 B/BL"
        )

    imm26 = delta >> 2
    if not (-(1 << 25) <= imm26 < (1 << 25)):
        raise MachOError(
            f"branch from 0x{src_va:x} to 0x{dst_va:x} is out of range for B/BL "
            f"(delta 0x{delta:x}, allowed +/-128 MiB)"
        )

    opcode = 0x94000000 if with_link else 0x14000000
    return opcode | (imm26 & 0x03FFFFFF)


def apply_patch(
    outer: OuterMachO,
    patch: PatchSpec,
    with_link: bool,
    originals: Mapping[str, OriginalMachO],
) -> AppliedPatch:
    src_entry = outer.get_entry(patch.src_entry)
    dst_entry = outer.get_entry(patch.dst_entry)

    src_norm = normalize_location_spec(src_entry, patch.src_loc, originals)
    dst_norm = normalize_location_spec(dst_entry, patch.dst_loc, originals)

    src_va, src_fileoff, src_desc_native = src_entry.resolve_location(src_norm.normalized)
    dst_va, dst_fileoff, dst_desc_native = dst_entry.resolve_location(dst_norm.normalized)

    src_desc = src_desc_native if patch.src_loc.kind != "symbol" else f"{src_norm.via}; final: {src_desc_native}"
    dst_desc = dst_desc_native if patch.dst_loc.kind != "symbol" else f"{dst_norm.via}; final: {dst_desc_native}"

    if src_fileoff + 4 > len(outer.data):
        raise MachOError(f"source file offset 0x{src_fileoff:x} is outside file")

    original_insn = struct.unpack_from("<I", outer.data, src_fileoff)[0]
    new_insn = encode_aarch64_branch(src_va, dst_va, with_link)
    struct.pack_into("<I", outer.data, src_fileoff, new_insn)

    return AppliedPatch(
        spec=patch,
        src_va=src_va,
        src_fileoff=src_fileoff,
        dst_va=dst_va,
        dst_fileoff=dst_fileoff,
        insn=new_insn,
        original_insn=original_insn,
        src_desc=src_desc,
        dst_desc=dst_desc,
    )


def format_hex(x: int) -> str:
    return f"0x{x:x}"


def cmd_list(args: argparse.Namespace) -> int:
    outer = parse_outer_macho(pathlib.Path(args.macho), args.fileoff_mode)

    print(f"Mach-O: {outer.path}")
    print(f"Fileset entries: {len(outer.entries)}")
    print()
    for entry in outer.entries:
        print(f"{entry.entry_id}")
        print(f"  vmaddr:        {format_hex(entry.vmaddr)}")
        print(f"  fileoff:       {format_hex(entry.fileoff)}")
        print(f"  inner-fileoff: {entry.macho.fileoff_mode}")
        if args.segments:
            for seg in entry.macho.segments:
                print(
                    f"    seg {seg.segname:<16} vmaddr={format_hex(seg.vmaddr):>14} "
                    f"vmsize={format_hex(seg.vmsize):>10} fileoff={format_hex(seg.fileoff):>10} "
                    f"filesize={format_hex(seg.filesize):>10}"
                )
                if args.sections:
                    for sec in seg.sections:
                        print(
                            f"      sec {sec.sectname:<16} addr={format_hex(sec.addr):>14} "
                            f"size={format_hex(sec.size):>10} offset={format_hex(sec.offset):>10}"
                        )
        print()
    return 0


def cmd_lookup_symbol(args: argparse.Namespace) -> int:
    if not args.lookup:
        raise MachOError("at least one --lookup specification is required")

    outer = parse_outer_macho(pathlib.Path(args.macho), args.fileoff_mode)
    originals = load_original_mappings(args.original or [])
    parsed = [parse_lookup_spec(item) for item in args.lookup]

    print(f"Mach-O: {outer.path}")
    print()
    for idx, (entry_id, loc) in enumerate(parsed, 1):
        result = resolve_lookup(outer, entry_id, loc, originals)
        print(f"Lookup #{idx}")
        print(f"  entry:          {result.entry_id}")
        print(f"  requested loc:  {result.requested.raw}")
        print(f"  normalized loc: {result.normalized.to_expr()}")
        print(f"  resolution:     {result.via}")
        print(f"  final desc:     {result.desc}")
        print(f"  final VA:       {format_hex(result.va)}")
        print(f"  final fileoff:  {format_hex(result.fileoff)}")
        if result.word32 is not None:
            print(f"  current word:   0x{result.word32:08x}")
        print()

    return 0


def choose_output_path(args: argparse.Namespace, in_path: pathlib.Path) -> pathlib.Path:
    if args.in_place and args.output:
        raise MachOError("use either --in-place or --output, not both")
    if args.in_place:
        return in_path
    if args.output:
        return pathlib.Path(args.output)
    return in_path.with_suffix(in_path.suffix + ".patched")


def cmd_patch(args: argparse.Namespace) -> int:
    if not args.patch:
        raise MachOError("at least one --patch specification is required")

    outer = parse_outer_macho(pathlib.Path(args.macho), args.fileoff_mode)
    originals = load_original_mappings(args.original or [])
    parsed_patches = [parse_patch_spec(p) for p in args.patch]

    applied: List[AppliedPatch] = []
    for patch in parsed_patches:
        applied.append(apply_patch(outer, patch, args.link, originals))

    out_path = choose_output_path(args, pathlib.Path(args.macho))
    out_path.write_bytes(outer.data)

    print(f"Wrote {out_path}")
    print()
    for idx, ap in enumerate(applied, 1):
        mnemonic = "BL" if args.link else "B"
        print(f"Patch #{idx}")
        print(f"  source entry:   {ap.spec.src_entry}")
        print(f"  source loc:     {ap.spec.src_loc.raw} ({ap.src_desc})")
        print(f"  source VA:      {format_hex(ap.src_va)}")
        print(f"  source fileoff: {format_hex(ap.src_fileoff)}")
        print(f"  target entry:   {ap.spec.dst_entry}")
        print(f"  target loc:     {ap.spec.dst_loc.raw} ({ap.dst_desc})")
        print(f"  target VA:      {format_hex(ap.dst_va)}")
        print(f"  target fileoff: {format_hex(ap.dst_fileoff)}")
        print(f"  original insn:  0x{ap.original_insn:08x}")
        print(f"  new insn:       0x{ap.insn:08x} ({mnemonic})")
        print()

    return 0


def build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Patch AArch64 branches across LC_FILESET_ENTRY members in a Mach-O fileset."
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    p_list = sub.add_parser("list", help="list LC_FILESET_ENTRY members")
    p_list.add_argument("macho", help="path to the little-endian flat Mach-O fileset")
    p_list.add_argument(
        "--segments",
        action="store_true",
        help="also show segments for each fileset entry",
    )
    p_list.add_argument(
        "--sections",
        action="store_true",
        help="also show sections for each segment (implies --segments)",
    )
    p_list.add_argument(
        "--fileoff-mode",
        choices=("auto", "relative", "absolute"),
        default="auto",
        help="interpret embedded segment fileoffs as relative to each entry or absolute in the outer file",
    )
    p_list.set_defaults(func=cmd_list)

    p_lookup = sub.add_parser("lookup-symbol", help="resolve symbol-based locations inside a fileset entry")
    p_lookup.add_argument("macho", help="path to the little-endian flat Mach-O fileset")
    p_lookup.add_argument(
        "--lookup",
        action="append",
        metavar="SPEC",
        help=(
            "lookup specification: '<entry> <loc>'. Usually <loc> is of the form "
            "'symbol:_my_symbol+0x10', but any supported location expression may be used."
        ),
    )
    p_lookup.add_argument(
        "--original",
        action="append",
        metavar="ENTRY=PATH",
        help=(
            "map a fileset entry name to its original standalone Mach-O for symbol-based resolution. "
            "May be repeated. Example: --original 'com.example.driver=/tmp/driver.macho'"
        ),
    )
    p_lookup.add_argument(
        "--fileoff-mode",
        choices=("auto", "relative", "absolute"),
        default="auto",
        help="interpret embedded segment fileoffs as relative to each entry or absolute in the outer file",
    )
    p_lookup.set_defaults(func=cmd_lookup_symbol)

    p_patch = sub.add_parser("patch", help="apply one or more branch patches")
    p_patch.add_argument("macho", help="path to the little-endian flat Mach-O fileset")
    p_patch.add_argument(
        "--patch",
        action="append",
        metavar="SPEC",
        help=(
            "patch specification: '<source-entry> <source-loc> <target-entry> <target-loc>'. "
            "Locations may be '__TEXT_EXEC:__text+0x1000', '__TEXT_EXEC+0x40', '+0x1000', "
            "or 'symbol:_my_func+0x10'."
        ),
    )
    p_patch.add_argument(
        "--original",
        action="append",
        metavar="ENTRY=PATH",
        help=(
            "map a fileset entry name to its original standalone Mach-O for symbol-based resolution. "
            "May be repeated. Example: --original 'com.example.driver=/tmp/driver.macho'"
        ),
    )
    p_patch.add_argument(
        "--link",
        action="store_true",
        help="emit BL instead of B",
    )
    p_patch.add_argument(
        "--in-place",
        action="store_true",
        help="overwrite the input file",
    )
    p_patch.add_argument(
        "--output",
        help="write patched Mach-O to this path (default: <input>.patched)",
    )
    p_patch.add_argument(
        "--fileoff-mode",
        choices=("auto", "relative", "absolute"),
        default="auto",
        help="interpret embedded segment fileoffs as relative to each entry or absolute in the outer file",
    )
    p_patch.set_defaults(func=cmd_patch)

    return p


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = build_argparser()
    args = parser.parse_args(argv)
    if getattr(args, "sections", False):
        args.segments = True

    try:
        return args.func(args)
    except MachOError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())