#!/usr/bin/env python3
"""Map obfuscated Il2CppInterop names between two game builds.

Il2CppInterop emits, for every il2cpp member, a stub line

    NativeMethodInfoPtr_<Name>_<Signature>_0 = IL2CPP.GetIl2CppMethodByToken(
        Il2CppClassPointerStore<<DeclaringType>>.NativeClassPtr, <token>);

The token encodes the global method index (0x06000000 + index), which shifts by
only a few units between builds, and <Signature> is a structural type spelling
(parameter kinds, arity, primitives) that survives renames -- only obfuscated
type identifiers inside it change. So: (normalized signature, token within a
window) identifies a member across builds without knowing the obfuscated names.

Usage:
    python tools/sept-map.py show   <aug|sep> <Type> [MemberFilter]
    python tools/sept-map.py find   <aug|sep> <token>
    python tools/sept-map.py map    <AugType> <AugMember> [window]
    python tools/sept-map.py shape  <aug|sep> <Type>
    python tools/sept-map.py types  <aug|sep> <regex>
"""

import os
import re
import sys
from collections import defaultdict

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ROOTS = {
    "aug": [
        os.path.join(ROOT, ".external", "aug-decomp", "acs"),
        os.path.join(ROOT, ".external", "aug-decomp", "rcn"),
        os.path.join(ROOT, ".external", "aug-decomp", "pfdoldleiba.cs"),
        os.path.join(ROOT, ".external", "aug-decomp", "aug-bootsequence.cs"),
        os.path.join(ROOT, ".external", "interop-decomp"),
    ],
    "sep": [
        os.path.join(ROOT, ".external", "sep-decomp", "acs"),
        os.path.join(ROOT, ".external", "sep-decomp", "rcn"),
        os.path.join(ROOT, ".external", "sep-decomp", "rp"),
        os.path.join(ROOT, ".external", "sep-decomp", "recr"),
        os.path.join(ROOT, ".external", "sep-decomp", "easy"),
        os.path.join(ROOT, ".external", "sep-decomp", "code"),
    ],
}

METHOD_RE = re.compile(
    r"NativeMethodInfoPtr_([A-Za-z0-9_]+?)_(.*?)_0 = IL2CPP\.GetIl2CppMethodByToken\("
    r"Il2CppClassPointerStore<([A-Za-z0-9_]+)>\.NativeClassPtr, (\d+)\)"
)
FIELD_RE = re.compile(r"NativeFieldInfoPtr_([A-Za-z0-9_]+?);")
TYPE_RE = re.compile(r"^\t*(?:\[[^\]]*\]\s*)*(?:public|internal|private|protected)?\s*"
                     r"(?:sealed\s+|abstract\s+|static\s+|unsafe\s+|partial\s+)*"
                     r"(class|struct|enum)\s+([A-Za-z0-9_]+)")
IDENT_RE = re.compile(r"[A-Z]{8,17}")


def norm(sig):
    """Blank out obfuscated identifiers, keep structural spelling."""
    out = IDENT_RE.sub("#", sig)
    return out


class TypeInfo:
    __slots__ = ("name", "file", "methods", "fields", "line")

    def __init__(self, name, file, line):
        self.name = name
        self.file = file
        self.line = line
        self.methods = []  # (memberName, sig, token, lineno)
        self.fields = []   # (fieldName, lineno)

    def shape(self):
        return tuple(norm(s) for _, s, _, _ in self.methods)


def load(root):
    types = {}
    by_token = defaultdict(list)
    files_to_parse = []
    for base in ROOTS[root]:
        if os.path.isfile(base) and base.endswith(".cs"):
            files_to_parse.append(base)
            continue
        if not os.path.isdir(base):
            continue
        for dirpath, _dirs, files in os.walk(base):
            for fn in files:
                if fn.endswith(".cs"):
                    files_to_parse.append(os.path.join(dirpath, fn))
    if True:
        for path in files_to_parse:
            if True:
                current = None
                with open(path, "r", encoding="utf-8", errors="replace") as fh:
                    for lineno, line in enumerate(fh, 1):
                        if current is None or "\tclass " in line or "\tstruct " in line or \
                                " enum " in line or line.startswith(("public ", "internal ")):
                            m = TYPE_RE.match(line)
                            if m:
                                name = m.group(2)
                                current = types.get(name)
                                if current is None:
                                    current = TypeInfo(name, os.path.relpath(path, ROOT), lineno)
                                    types[name] = current
                        m = METHOD_RE.search(line)
                        if m:
                            owner = m.group(3)
                            if owner not in types:
                                types[owner] = TypeInfo(owner, os.path.relpath(path, ROOT), 0)
                            entry = (m.group(1), m.group(2), int(m.group(4)), lineno)
                            types[owner].methods.append(entry)
                            by_token[int(m.group(4))].append((owner, entry))
                            continue
                        m = FIELD_RE.search(line)
                        if m and current is not None:
                            current.fields.append((m.group(1), lineno))
    return types, by_token


def show(types, name, filt=""):
    t = types.get(name)
    if t is None:
        print(f"type {name}: NOT FOUND")
        return
    print(f"=== {name}  ({t.file}) methods={len(t.methods)} fields={len(t.fields)}")
    for member, sig, token, line in t.methods:
        if filt and filt not in member:
            continue
        print(f"  {token}  {token - 0x6000000:>7}  {member}  [{norm(sig)}]  {sig}")
    for field, line in t.fields:
        if filt and filt not in field:
            continue
        print(f"  field {field}")


def map_member(types, sep_tokens, aug_type, aug_member, window=60):
    t = types.get(aug_type)
    if t is None:
        print(f"august type {aug_type} not found")
        return
    targets = [(m, s, k) for m, s, k, _ in t.methods if m == aug_member]
    if not targets:
        targets = [(m, s, k) for m, s, k, _ in t.methods if aug_member in m]
    if not targets:
        print(f"member {aug_member} not found on {aug_type}")
        return
    for member, sig, token in targets:
        print(f"\nAUG {aug_type}.{member} token={token} (idx {token - 0x6000000}) sig=[{norm(sig)}]")
        hits = []
        for tk in range(token - window, token + window + 1):
            for owner, entry in sep_tokens.get(tk, []):
                m2, s2, tok, line = entry
                if norm(s2) == norm(sig):
                    hits.append((abs(tk - token), tk, owner, m2, s2))
        hits.sort()
        for dist, tk, owner, m2, s2 in hits[:8]:
            print(f"  SEP d={dist:+3d} {tk} {owner}.{m2}  [{norm(s2)}]  {s2}")
        if not hits:
            print("  no signature match in window")


def shape(types, name, verbose=True):
    t = types.get(name)
    if t is None:
        print(f"type {name}: NOT FOUND")
        return
    if verbose:
        print(f"=== shape {name} ({t.file}) methods={len(t.methods)} fields={len(t.fields)}")
        for member, sig, token, _ in t.methods:
            print(f"  {token - 0x6000000:>7}  {member}  [{norm(sig)}]")
        for field, _ in t.fields:
            print(f"  field:{field}")
    return t


def similarity(a, b):
    """Match two ordered (sig,) sequences; return (score, token pairs)."""
    import difflib
    sa = [x[0] for x in a]
    sb = [x[0] for x in b]
    sm = difflib.SequenceMatcher(None, sa, sb, autojunk=False)
    pairs = []
    for block in sm.get_matching_blocks():
        for i in range(block.size):
            pairs.append((a[block.a + i][1], b[block.b + i][1]))
    score = 2.0 * len(pairs) / max(1, len(sa) + len(sb))
    return score, pairs


def match_type(aug_types, sep_types, aug_name, top=6):
    t = aug_types.get(aug_name)
    if t is None:
        print(f"august type {aug_name}: NOT FOUND")
        return
    a = [(norm(s), tk) for _, s, tk, _ in t.methods]
    scored = []
    for name, st in sep_types.items():
        b = [(norm(s), tk) for _, s, tk, _ in st.methods]
        if not b and not a:
            continue
        score, pairs = similarity(a, b)
        if not pairs:
            continue
        deltas = [tk2 - tk1 for tk1, tk2 in pairs]
        median = sorted(deltas)[len(deltas) // 2]
        spread = max(deltas) - min(deltas)
        scored.append((score, -abs(median), name, st, median, spread, len(pairs)))
    scored.sort(key=lambda x: (-x[0], x[1]))
    print(f"=== {aug_name} ({t.file}) {len(t.methods)} methods; {len(t.fields)} fields")
    for score, _n, name, st, median, spread, pairs in scored[:top]:
        print(f"  {score:5.2f}  {name:<16} d={median:+5d} spread={spread:<5d} pairs={pairs:<4d} "
              f"m={len(st.methods)} f={len(st.fields)}  {st.file}")


def main():
    if len(sys.argv) < 3:
        print(__doc__)
        return 1
    cmd = sys.argv[1]
    if cmd == "map":
        aug_types, _ = load("aug")
        _sep_types, sep_tokens = _load_tokens("sep")
        map_member(aug_types, sep_tokens, sys.argv[2], sys.argv[3],
                   int(sys.argv[4]) if len(sys.argv) > 4 else 60)
        return 0
    if cmd in ("match", "verify"):
        aug_types, _ = load("aug")
        sep_types, _ = load("sep")
        names = [sys.argv[2]] if cmd == "match" else sys.argv[2].split(",")
        top = int(sys.argv[3]) if len(sys.argv) > 3 else 6
        for name in names:
            match_type(aug_types, sep_types, name, top)
        return 0
    root = sys.argv[2]
    types, by_token = load(root)
    if cmd == "show":
        show(types, sys.argv[3], sys.argv[4] if len(sys.argv) > 4 else "")
    elif cmd == "shape":
        shape(types, sys.argv[3])
    elif cmd in ("find", "findidx", "window"):
        token = int(sys.argv[3], 0)
        if token < 0x6000000:
            token += 0x6000000
        span = int(sys.argv[4]) if len(sys.argv) > 4 else 0
        for tk in range(token - span, token + span + 1):
            for owner, entry in sorted(by_token.get(tk, [])):
                member, sig, tok, line = entry
                mark = "<<<" if tk == token else "   "
                print(f"  {mark} {tk}  {owner}.{member}  [{norm(sig)}]")
    elif cmd == "types":
        rx = re.compile(sys.argv[3])
        for name in sorted(t for t in types if rx.search(t)):
            print(f"  {name}  ({types[name].file})")
    else:
        print(__doc__)
        return 1
    return 0


_CACHE = {}


def _load_tokens(root):
    if root not in _CACHE:
        _CACHE[root] = load(root)
    return _CACHE[root]


if __name__ == "__main__":
    sys.exit(main())