"""Shared reverse-engineering helpers for the MAG160C SDK binaries."""
import io
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / ".tools" / "python-revlibs"))

from elftools.elf.elffile import ELFFile
from capstone import Cs, CS_ARCH_ARM64, CS_MODE_ARM, CS_ARCH_X86, CS_MODE_64

LIBCORESDK = r"C:\Project\MAG160C\IR_Camera_SDK-1.0.1\android\android\lib\arm64-v8a\libcoresdk.so"
LIBMAGCORE = r"C:\Project\MAG160C\IR_Camera_SDK-1.0.1\linux\1.2.2\1.2.2\lib\libmagcore.so.2.1.1"


class Loader:
    def __init__(self, path):
        self.path = Path(path)
        self._data = path_bytes = Path(path).read_bytes()
        self.elf = ELFFile(io.BytesIO(path_bytes))
        self.arch = self.elf["e_machine"]
        self.isa64 = self.elf.elfclass == 64
        if self.arch == "EM_AARCH64":
            self.md = Cs(CS_ARCH_ARM64, CS_MODE_ARM)
            self.md.detail = True
        else:
            self.md = Cs(CS_ARCH_X86, CS_MODE_64)
            self.md.detail = True
        self._syms = None

    @property
    def syms(self):
        if self._syms is None:
            self._syms = []
            for sec in self.elf.iter_sections():
                if sec["sh_type"] in ("SHT_SYMTAB", "SHT_DYNSYM"):
                    for s in sec.iter_symbols():
                        self._syms.append(s)
        return self._syms

    def find_sym(self, name, demangled=True):
        out = []
        for s in self.syms:
            n = s.name
            if demangled:
                d = demangle(n)
                if d is not None:
                    n = d
            if n == name:
                out.append(s)
        return out

    def vaddr_to_offset(self, va):
        for sec in self.elf.iter_sections():
            addr, size = sec["sh_addr"], sec["sh_size"]
            if size and addr <= va < addr + size:
                return sec["sh_offset"] + (va - addr)
        return None

    def bytes_at(self, va, n):
        off = self.vaddr_to_offset(va)
        if off is None:
            return None
        return self._data[off:off + n]

    def read_u32(self, va):
        b = self.bytes_at(va, 4)
        return int.from_bytes(b, "little") if b and len(b) == 4 else None

    def read_ptr(self, va):
        n = 8 if self.isa64 else 4
        b = self.bytes_at(va, n)
        return int.from_bytes(b, "little") if b and len(b) == n else None

    def cstr_at(self, va, maxlen=256):
        b = self.bytes_at(va, maxlen)
        if b is None:
            return None
        end = b.find(b"\x00")
        if end < 0:
            return None
        return b[:end].decode("latin-1", "replace")

    def disasm(self, va, count=None, stop_after_ret=8):
        """Linear-sweep disassemble starting at va; stop after N rets or count insns."""
        lines = []
        rets = 0
        i = 0
        cur = va
        while count is None or i < count:
            b = self.bytes_at(cur, 16)
            if not b:
                break
            ins = next(self.md.disasm(b, cur), None)
            if ins is None:
                lines.append(f"{cur:012x}: <unresolvable>")
                cur += 4
                continue
            lines.append(f"{ins.address:012x}: {ins.mnemonic:<8s} {ins.op_str}")
            if ins.mnemonic.startswith("ret") or ins.mnemonic == "br" or (self.arch == "EM_AARCH64" and ins.mnemonic == "ret"):
                rets += 1
                if stop_after_ret and rets >= stop_after_ret:
                    break
            cur += ins.size
            i += 1
        return lines


def demangle(name):
    """Small C++ demangler covering the symbols we care about."""
    if not name.startswith("_Z"):
        return None
    m = re.match(r"^_Z(T|V|N)?(.*)$", name)
    if not m:
        return None
    rest = m.group(2)
    if m.group(1) == "N":
        # N...E nested
        parts = []
        i = 0
        while i < len(rest) and rest[i] != "E":
            n = 0
            j = i
            while rest[j].isdigit():
                n = n * 10 + int(rest[j])
                j += 1
            if n == 0:
                return name
            parts.append(rest[j:j + n])
            i = j + n
        i += 1
        args = demangle_args(rest[i:])
        return "::".join(parts) + ("(" + ", ".join(args) + ")" if args else "")
    # plain function _Z<N><name><args>
    m2 = re.match(r"^(\d+)([A-Za-z_][A-Za-z0-9_]*)(.*)$", rest)
    if not m2:
        return name
    n = int(m2.group(1))
    fn = m2.group(2)[:n]
    args = demangle_args(m2.group(3))
    return fn + ("(" + ", ".join(args) + ")" if args else "")


_ARG_TYPES = {
    "i": "int", "j": "uint", "t": "ushort", "k": "uint", "l": "ulong",
    "P": "void*", "PK": "const void*", "E": "struct", "S": "size_t",
}


def demangle_args(s):
    args = []
    i = 0
    while i < len(s) and s[i] != "E":
        c = s[i]
        if c in "ijtl":
            args.append({"i": "int", "j": "uint", "t": "ushort", "l": "ulong"}[c])
            i += 1
        elif c == "P":
            args.append("void*")
            i += 1
        elif c == "K":
            args.append("const void*")
            i += 1
        else:
            m = re.match(r"^(\d+)([A-Za-z_][A-Za-z0-9_]*)(.*)$", s[i:])
            if not m:
                args.append("?")
                i += 1
                continue
            n = int(m.group(1))
            args.append(m.group(2)[:n])
            i += len(m.group(1)) + n
    return args


def resolve_plt(ld, va, depth=8):
    """Follow ARM64 PLT stub chains to final function address."""
    seen = set()
    cur = va
    for _ in range(depth):
        if cur in seen:
            break
        seen.add(cur)
        b = ld.bytes_at(cur, 16)
        if not b:
            return None
        # PLT stub: adrp x16, #p; ldr x17, [x16, #o]; add x16, x16, #o; br x17
        if len(b) >= 16:
            w0 = int.from_bytes(b[0:4], "little")
            w1 = int.from_bytes(b[4:8], "little")
            w2 = int.from_bytes(b[8:12], "little")
            w3 = int.from_bytes(b[12:16], "little")
            if (w0 & 0xFF000000) == 0x90000000 and (w1 & 0xFFC00000) == 0xF9400000:
                imm = ((w0 >> 29) & 0x3) | ((w0 >> 3) & 0xFFFFC)
                if w0 & 0x400000:
                    imm |= ~0xFFFFF & 0xFFFFFFFF
                imm <<= 12
                page = (cur + imm) & ~0xFFF
                off = ((w1 >> 10) & 0xFFF) * 8
                got = page + off
                target = ld.read_ptr(got)
                if target is None:
                    return None
                cur = target
                continue
        # direct branch (b/bl)
        if (b[0] & 0xFC) == 0x14:
            w = int.from_bytes(b[0:4], "little")
            imm = (w & 0x3FFFFFF) << 2
            if w & 0x2000000:
                imm |= ~0x7FFFFFF & 0xFFFFFFFF
            return (cur + imm) & 0xFFFFFFFF
        return cur
    return None


def disasm_func(path, symbol, count=None, stop_after_ret=8):
    ld = Loader(path)
    syms = ld.find_sym(symbol)
    if not syms:
        print(f"NOT FOUND: {symbol}")
        return
    for s in syms:
        if s["st_info"]["type"] != "STT_FUNC":
            continue
        print(f"== {symbol} @ {s['st_value']:x} size={s['st_size']} ==")
        for line in ld.disasm(s["st_value"], count=count, stop_after_ret=stop_after_ret):
            print(line)


def find_xrefs(path, target_va):
    """Find direct branches/calls to target_va in .text via capstone linear sweep."""
    ld = Loader(path)
    hits = []
    for sec in ld.elf.iter_sections():
        if sec["sh_type"] != "SHT_PROGBITS" or not (sec["sh_flags"] & 0x4):
            continue
        va = sec["sh_addr"]
        data = sec.data()
        cur = va
        while cur < va + len(data):
            ins = next(ld.md.disasm(data[cur - va:], cur), None)
            if ins is None:
                break
            if "call" in ins.mnemonic or "b" == ins.mnemonic or ins.mnemonic.startswith("b."):
                for op in ins.operands:
                    if op.type == 2 and op.imm == target_va:  # CS_OP_IMM
                        hits.append((ins.address, ins.mnemonic, ins.op_str))
            cur += ins.size
    return hits


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--so", default=LIBCORESDK)
    ap.add_argument("--sym")
    ap.add_argument("--addr", type=lambda x: int(x, 0))
    ap.add_argument("--count", type=int)
    ap.add_argument("--ret", type=int, default=8)
    ap.add_argument("--xrefs", action="store_true", help="find xrefs to --addr")
    args = ap.parse_args()
    if args.xrefs:
        for a, m, o in find_xrefs(args.so, args.addr):
            print(f"{a:012x}: {m} {o}")
    elif args.addr is not None:
        for line in Loader(args.so).disasm(args.addr, count=args.count, stop_after_ret=args.ret):
            print(line)
    else:
        disasm_func(args.so, args.sym, count=args.count, stop_after_ret=args.ret)