#!/usr/bin/env python3
"""EverQuest (Titanium, eqgame.exe 2005-10-31) matching-decompilation toolkit.

Subcommands:
  verify                 check orig/eqgame.exe is the expected build
  import-ghidra FILE     seed config/eqgame/symbols.txt from tools/ghidra/ExportFunctions.java output
  build                  compile src/ with MSVC 7.1 (Wine) via ninja
  check [NAME|ADDR...]   compare compiled functions against the original, byte for byte
  diff NAME|ADDR         side-by-side disassembly of original vs compiled
  progress               matched bytes / functions
  pair SRC --range A-B   annotate SRC's unannotated functions by unique byte match in A-B
  libscan PROBE.cpp...   find copies of library-probe functions in .text (--apply: mark them lib)
  tus [N]                estimate translation-unit boundaries from .rdata marker tables
  link                   assemble build/eqgame.exe from compiled objects + original, compare sha1
  map ADDR NAME [KIND]   name (or rename) a symbol in the map
  todo [N]               list unmatched functions, smallest first
"""
import argparse
import difflib
import fnmatch
import hashlib
import os
import re
import struct
import subprocess
import sys

import yaml

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from binfmt import (CoffObject, PEImage, REL_I386_DIR32, REL_I386_REL32,  # noqa: E402
                    IMAGE_SCN_CNT_CODE)

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CFG_DIR = os.path.join(ROOT, "config", "eqgame")
BUILD = os.path.join(ROOT, "build")
OBJ_DIR = os.path.join(BUILD, "obj")


def pe_path(cfg):
    return os.path.join(ROOT, cfg["original"])


def load_config():
    with open(os.path.join(CFG_DIR, "config.yaml")) as f:
        return yaml.safe_load(f)


# ---------------------------------------------------------------- symbol map

class SymbolMap:
    """config/eqgame/symbols.txt: `ADDR SIZE KIND NAME` per line, '#' comments.

    KIND: fn (game code), lib (CRT/third-party, not decompiled), thunk, data.
    """

    def __init__(self, path):
        self.path = path
        self.entries = {}  # addr -> [size, kind, name]
        if os.path.exists(path):
            with open(path) as f:
                for line in f:
                    line = line.split("#", 1)[0].split()
                    if len(line) >= 4:
                        self.entries[int(line[0], 16)] = [int(line[1], 16), line[2], line[3]]
        self.base = {a: list(e) for a, e in self.entries.items()}
        self._reindex()

    def _reindex(self):
        self.by_name = {e[2]: a for a, e in self.entries.items()}

    def lookup(self, key):
        """Resolve a name or hex address to an address (or None)."""
        if key in self.by_name:
            return self.by_name[key]
        try:
            a = int(key, 16)
        except ValueError:
            return None
        return a if a in self.entries else None

    def set(self, addr, name, kind=None, size=None):
        old = self.by_name.get(name)
        if old is not None and old != addr:
            raise ValueError(f"{name} already mapped at 0x{old:08x}")
        e = self.entries.setdefault(addr, [0, kind or "fn", name])
        e[2] = name
        if kind:
            e[1] = kind
        if size is not None:
            e[0] = size
        self._reindex()

    def save(self):
        """Write the map. Names that come from source annotations stay in the source."""
        for a, (name, kind) in getattr(self, "annotated", {}).items():
            if a in self.base:
                self.entries[a] = list(self.base[a])
            else:
                del self.entries[a]
        with open(self.path, "w") as f:
            f.write("# eqgame.exe symbol map: ADDR SIZE KIND NAME\n")
            f.write("# KIND = fn | lib | thunk | data. Names are MSVC-mangled once decompiled.\n")
            for a in sorted(self.entries):
                size, kind, name = self.entries[a]
                f.write(f"0x{a:08x} 0x{size:x} {kind} {name}\n")


# ---------------------------------------------------------------- commands

def cmd_verify(args):
    cfg = load_config()
    path = os.path.join(ROOT, cfg["original"])
    if not os.path.exists(path):
        sys.exit(f"missing {cfg['original']}: copy or symlink your eqgame.exe there")
    with open(path, "rb") as f:
        h = hashlib.sha1(f.read()).hexdigest()
    if h != cfg["sha1"]:
        sys.exit(f"{cfg['original']}: sha1 {h} != expected {cfg['sha1']} (wrong client build)")
    print(f"{cfg['original']}: OK ({h})")


def cmd_import_ghidra(args):
    m = SymbolMap(os.path.join(CFG_DIR, "symbols.txt"))
    added = 0
    with open(args.file) as f:
        for line in f:
            addr, size, name, kind = line.split()
            addr, size = int(addr, 16), int(size, 16)
            if addr in m.entries:
                m.entries[addr][0] = m.entries[addr][0] or size
                continue
            if name.startswith("FUN_"):
                name = f"FUN_{addr:08x}"
            elif kind != "lib":
                kind = "fn"
            name = re.sub(r"[^\w@?$<>:.~-]", "_", name)
            if name in m.by_name:
                name = f"{name}_{addr:08x}"
            m.entries[addr] = [size, kind, name]
            m.by_name[name] = addr
            added += 1
    m.save()
    print(f"imported {added} functions ({len(m.entries)} total)")


ANNOT = re.compile(r"^\s*//\s*(FUNCTION|GLOBAL|LIBRARY):\s*EQGAME\s+(0x[0-9a-fA-F]+)\s*(\S+)?\s*(NONMATCHING)?")


def mangled_prefixes(qualname):
    """Possible MSVC mangled-name prefixes for `A::B::f` / `X::X` / `X::~X` / `f` / `g_var`."""
    parts = [p for p in qualname.split("::") if p]
    name, scopes = parts[-1], parts[:-1][::-1]
    scope = "".join(p + "@" for p in scopes)
    if scopes and name == scopes[0]:
        return [f"??0{scope}@"]
    if scopes and name == "~" + scopes[0]:
        return [f"??1{scope}@"]
    if name.startswith("~"):  # destructor declared inside its class body
        return [f"??1{name[1:]}@@"]
    # unqualified: a free function/global, or a member declared inside its class body
    return [f"?{name}@{scope}@"] if scopes else [f"?{name}@", f"_{name}"]


def parse_annotations(path):
    """[(kind, addr, explicit_mangled_or_None, qualname_or_None, nonmatching)] in file order."""
    with open(path, encoding="latin-1") as f:
        lines = f.read().splitlines()
    out = []
    for i, line in enumerate(lines):
        m = ANNOT.match(line)
        if not m:
            continue
        kind, addr, explicit, nonmatch = m.group(1), int(m.group(2), 16), m.group(3), bool(m.group(4))
        if explicit == "NONMATCHING":
            explicit, nonmatch = None, True
        qual = None
        if not explicit and kind != "LIBRARY":
            # the declaration follows (possibly after more annotations/comments)
            for nxt in lines[i + 1:i + 8]:
                t = nxt.strip()
                if not t or t.startswith("//"):
                    continue
                if kind == "FUNCTION":
                    mm = re.search(r"([~\w:]+)\s*\(", t)
                else:
                    mm = re.search(r"([\w:]+)\s*(?:\[[^\]]*\]\s*)*(?:=|;|$)", t.split("=")[0] + ";")
                if mm:
                    qual = mm.group(1)
                break
        out.append((kind, addr, explicit, qual, nonmatch))
    return out


def all_header_syms(objs):
    seen, out = set(), []
    for src in source_files():
        op = os.path.join(ROOT, obj_path(src))
        if os.path.exists(op):
            obj = objs.setdefault(op, CoffObject(op))
            for x in obj.symbols.values():
                if x.is_external and x.name not in seen:  # definitions and references
                    seen.add(x.name)
                    out.append(x)
    return out


def apply_annotations(smap):
    """Override map names with the annotations in src/ and include/, resolved against build/obj."""
    smap.annotated, smap.nonmatching, problems = {}, set(), []
    objs = {}
    for src in source_files() + header_files():
        anns = parse_annotations(os.path.join(ROOT, src))
        if not anns:
            continue
        syms = []
        if src in source_files():
            op = os.path.join(ROOT, obj_path(src))
            if os.path.exists(op):
                obj = objs.setdefault(op, CoffObject(op))
                syms = [x for x in obj.symbols.values() if x.is_external]  # definitions and references
        else:
            # header: inline functions/templates may be emitted by any object
            syms = all_header_syms(objs)
        used = set()
        for kind, addr, explicit, qual, nonmatch in anns:
            name = explicit
            if not name and qual:
                pref = mangled_prefixes(qual)
                cands = [x.name for x in syms if x.name not in used and any(x.name.startswith(p) for p in pref)
                         and (x.is_function == (kind == "FUNCTION"))]
                if cands:
                    name = cands[0]
                elif syms and src.endswith((".c", ".cpp")):
                    problems.append(f"{src}: {kind} 0x{addr:08x} `{qual}`: no matching symbol in object")
            if not name:
                continue  # header declaration of something defined elsewhere, or not built yet
            used.add(name)
            prev = smap.by_name.get(name)
            if prev is not None and prev != addr:
                problems.append(f"{src}: {name} annotated at 0x{addr:08x} and 0x{prev:08x}")
                continue
            k = {"FUNCTION": None, "GLOBAL": "data", "LIBRARY": "lib"}[kind]
            if addr in smap.entries and smap.entries[addr][2] != name:
                smap.by_name.pop(smap.entries[addr][2], None)
            smap.set(addr, name, k)
            smap.annotated[addr] = (name, k)
            if nonmatch:
                smap.nonmatching.add(addr)
    return problems


def load_map():
    smap = SymbolMap(os.path.join(CFG_DIR, "symbols.txt"))
    for p in apply_annotations(smap):
        print(f"warning: {p}", file=sys.stderr)
    return smap


def header_files():
    out = []
    for base, _, files in os.walk(os.path.join(ROOT, "include")):
        for fn in files:
            if fn.endswith((".h", ".hpp")):
                out.append(os.path.relpath(os.path.join(base, fn), ROOT))
    return sorted(out)


def source_files():
    out = []
    for base, _, files in os.walk(os.path.join(ROOT, "src")):
        for fn in files:
            if fn.endswith((".c", ".cpp")):
                out.append(os.path.relpath(os.path.join(base, fn), ROOT))
    return sorted(out)


def obj_path(src):
    return os.path.join("build", "obj", os.path.splitext(src)[0] + ".obj")


def cmd_build(args):
    cfg = load_config()
    lines = [
        "ninja_required_version = 1.5",
        "builddir = build",
        f"cflags = {cfg['cflags']}",
        "rule cl",
        "  command = tools/cl.sh /nologo /c $cflags $extra /Iinclude /Fo$out $in > $out.log 2>&1"
        " || { cat $out.log; rm -f $out; exit 1; }",
        "  description = CL $in",
        "",
    ]
    for src in source_files():
        with open(os.path.join(ROOT, src), encoding="latin-1") as f:
            head = f.read(4096)
        extra = " ".join(re.findall(r"^//\s*CFLAGS:\s*(.*)$", head, re.M))
        # every header is an implicit input (cl's /showIncludes paths are Windows paths under Wine)
        lines.append(f"build {obj_path(src)}: cl {src} | config/eqgame/config.yaml {' '.join(header_files())}")
        if extra:
            lines.append(f"  extra = {extra}")
    os.makedirs(BUILD, exist_ok=True)
    with open(os.path.join(BUILD, "build.ninja"), "w") as f:
        f.write("\n".join(lines) + "\n")
    r = subprocess.run(["ninja", "-f", "build/build.ninja"] + (["-v"] if args.verbose else []), cwd=ROOT)
    if r.returncode:
        sys.exit(r.returncode)


def compiled_functions():
    """name -> (obj, sym, sec, start, end) for every function in build/obj."""
    out = {}
    for src in source_files():
        p = os.path.join(ROOT, obj_path(src))
        if not os.path.exists(p):
            continue
        obj = CoffObject(p)
        for sym, sec, start, end in obj.functions():
            if sym.is_external:
                out[sym.name] = (obj, sym, sec, start, end)
    return out


# compiler-generated constants: string literals, float/double/SSE constants. Placed wherever the
# linker likes, so they're verified by content rather than by address.
COMPILER_DATA = ("??_C@", "__real@", "__xmm@", "__mask@")


class Result:
    def __init__(self, name, addr):
        self.name, self.addr = name, addr
        self.status = "?"
        self.size = 0
        self.notes = []
        self.suggest = {}  # symbol name -> original address
        self.locals = []   # compiler-generated code verified alongside: (obj, sec, start, end, addr)


def local_code(obj, tsym):
    """Compiler-generated code in another section of this object (EH handlers, unwind funclets)."""
    sec = obj.sections[tsym.section]
    if not sec.chars & IMAGE_SCN_CNT_CODE:
        return None
    ends = [x.value for x in obj.symbols.values() if x.section == tsym.section and x.value > tsym.value]
    return sec, tsym.value, min(ends + [len(sec.data)])


def compare(pe, smap, obj, sym, sec, start, end, addr=None, seen=None):
    name = sym.name
    addr = smap.by_name[name] if addr is None else addr
    seen = set() if seen is None else seen
    seen.add((obj.path, sym.name))
    res = Result(name, addr)
    cb = sec.data[start:end]
    res.size = len(cb)
    ob = pe.read(addr, len(cb))
    mask = bytearray(len(cb))
    relocs = [(va - start, si, t) for va, si, t in sec.relocs if start <= va < end]
    for off, _, _ in relocs:
        mask[off:off + 4] = b"\1\1\1\1"

    diffs = [i for i in range(len(cb)) if not mask[i] and cb[i] != ob[i]]
    mapsize = smap.entries[addr][0] if addr in smap.entries else 0
    if mapsize > len(cb):
        tail = pe.read(addr + len(cb), mapsize - len(cb))
        if tail.strip(b"\xcc\x90"):
            res.notes.append(f"original is longer (map size 0x{mapsize:x} vs 0x{len(cb):x})")
            diffs.append(len(cb))
    if diffs:
        res.status = "DIFF"
        res.notes.insert(0, f"{len(diffs)} byte(s) differ, first at +0x{diffs[0]:x}")
        return res

    unknown = bad = 0
    for off, si, typ in relocs:
        tsym = obj.symbols[si]
        if typ == REL_I386_DIR32:
            addend, = struct.unpack_from("<i", cb, off)
            base = struct.unpack_from("<I", ob, off)[0] - addend
        elif typ == REL_I386_REL32:
            addend, = struct.unpack_from("<i", cb, off)
            base = addr + off + 4 + struct.unpack_from("<i", ob, off)[0] - addend
        else:
            res.notes.append(f"+0x{off:x}: unhandled reloc type {typ:#x} to {tsym.name}")
            unknown += 1
            continue
        base &= 0xFFFFFFFF

        if tsym.section == sec.index:
            expected = addr - start + tsym.value
        elif tsym.section == -1:  # absolute symbol
            expected = tsym.value
        elif tsym.name == "__except_list":  # CRT absolute symbol: fs:[0]
            expected = 0
        elif tsym.section > 0 and not tsym.is_external and tsym.name not in smap.by_name \
                and local_code(obj, tsym):
            if (obj.path, tsym.name) in seen:
                continue
            lsec, lstart, lend = local_code(obj, tsym)
            sub = compare(pe, smap, obj, tsym, lsec, lstart, lend, base, seen)
            res.locals.append((obj, lsec, lstart, lend, base))
            res.locals += sub.locals
            res.suggest.update(sub.suggest)
            if sub.status in ("DIFF", "BADREF"):
                res.notes.append(f"+0x{off:x}: {tsym.name} at 0x{base:08x}: {sub.status} {'; '.join(sub.notes)}")
                bad += 1
            elif sub.status == "MATCH?":
                unknown += 1
            continue
        elif tsym.name in smap.by_name:
            expected = smap.by_name[tsym.name]
        elif tsym.section > 0 and (not tsym.is_external or tsym.name.startswith(COMPILER_DATA)) \
                and not (obj.sections[tsym.section].chars & IMAGE_SCN_CNT_CODE):
            # local data (string literal, float constant, static table): verify by content
            tsec = obj.sections[tsym.section]
            blob = tsec.data[tsym.value:]
            rmask = {o - tsym.value + k for o, _, _ in tsec.relocs for k in range(4)}
            got = pe.read(base, len(blob)) if pe.contains(base) else b""
            if got and all(i in rmask or got[i] == blob[i] for i in range(len(blob))):
                continue
            res.notes.append(f"+0x{off:x}: {tsym.name} content differs at 0x{base:08x}")
            bad += 1
            continue
        else:
            res.suggest[tsym.name] = base
            unknown += 1
            continue
        if expected != base:
            res.notes.append(f"+0x{off:x}: {tsym.name} -> 0x{base:08x}, map says 0x{expected:08x}")
            bad += 1
    if bad:
        res.status = "BADREF"
    elif unknown:
        res.status = "MATCH?"
        for n, a in res.suggest.items():
            cur = smap.entries.get(a)
            res.notes.append(f"unmapped {n} = 0x{a:08x}" + (f" (currently {cur[2]})" if cur else ""))
    else:
        res.status = "MATCH"
    return res


def run_checks(filters=(), inline=False):
    cfg = load_config()
    pe = PEImage(os.path.join(ROOT, cfg["original"]))
    smap = load_map()
    results, unmapped = [], []
    for name, (obj, sym, sec, start, end) in sorted(compiled_functions().items()):
        if name not in smap.by_name:
            if inline or not obj.is_inline(sec.index):  # inline copies are usually folded away
                unmapped.append((name, obj.path))
            continue
        if filters and not any(fnmatch.fnmatchcase(name, f) or smap.lookup(f) == smap.by_name[name]
                               for f in filters):
            continue
        r = compare(pe, smap, obj, sym, sec, start, end)
        if r.addr in smap.nonmatching and r.status in ("DIFF", "BADREF"):
            r.status = "NONMAT"  # known non-matching (functionally equivalent) decompilation
        results.append(r)
    return smap, results, unmapped


def cmd_check(args):
    smap, results, unmapped = run_checks(args.names, args.inline)
    for r in sorted(results, key=lambda r: r.addr):
        print(f"{r.status:7} 0x{r.addr:08x} {r.size:5d}  {r.name}")
        if args.verbose or r.status != "MATCH":
            for n in r.notes:
                print(f"          {n}")
    if not args.names:
        for name, path in unmapped:
            print(f"UNMAPPED  {name}  ({os.path.relpath(path, ROOT)}): add a // FUNCTION: EQGAME 0xADDR annotation")
    if args.learn:
        n = 0
        for r in results:
            if r.status != "MATCH?":
                continue
            for name, a in r.suggest.items():
                if name in smap.by_name:
                    continue
                kind = None if a in smap.entries else "data"
                smap.set(a, name, kind)
                n += 1
        smap.save()
        print(f"learned {n} symbol(s) from matching functions")
    counts = {}
    for r in results:
        counts[r.status] = counts.get(r.status, 0) + 1
    print("  ".join(f"{k}: {v}" for k, v in sorted(counts.items())) or "nothing to check")
    if any(r.status in ("DIFF", "BADREF") for r in results):
        sys.exit(1)


def disasm(path, extra):
    """[(address, text, reloc_target_or_None)] from objdump."""
    out = subprocess.run(["i686-w64-mingw32-objdump", "-d", "-r", "-M", "intel", "--no-show-raw-insn",
                          *extra, path], capture_output=True, text=True).stdout
    lines = []
    for line in out.splitlines():
        m = re.match(r"\s+([0-9a-f]+):\s+(.*)", line)
        if not m:
            continue
        text = m.group(2).strip()
        if re.match(r"(dir32|DISP32|IMAGE_REL|secrel)", text, re.I):
            a, t, _ = lines[-1]
            lines[-1] = (a, t, text.split()[-1])
        else:
            lines.append((int(m.group(1), 16), re.sub(r"\s+", " ", text), None))
    return lines


def disasm_blob(obj, sec, start, end):
    """Disassemble a compiled function from its raw bytes, annotating relocations by offset."""
    import tempfile
    relocs = {va - start: obj.symbols[si].name for va, si, _ in sec.relocs if start <= va < end}
    with tempfile.NamedTemporaryFile(suffix=".bin") as f:
        f.write(sec.data[start:end])
        f.flush()
        out = subprocess.run(["i686-w64-mingw32-objdump", "-D", "-b", "binary", "-m", "i386", "-M", "intel",
                              "--no-show-raw-insn", f.name], capture_output=True, text=True).stdout
    lines = []
    for line in out.splitlines():
        m = re.match(r"\s+([0-9a-f]+):\s+(.*)", line)
        if not m:
            continue
        a, text = int(m.group(1), 16), re.sub(r"\s+", " ", m.group(2).strip())
        nxt = [o for o in relocs if a <= o < a + 16]
        lines.append((a, text, None))
    # attach each relocation to the instruction containing it
    for i, (a, text, _) in enumerate(lines):
        end_a = lines[i + 1][0] if i + 1 < len(lines) else end - start
        r = [relocs[o] for o in relocs if a <= o < end_a]
        if r:
            lines[i] = (a, text, r[-1])
    return lines


BRANCH = re.compile(r"^(j\w+|call|loop\w*) (?:0x)?([0-9a-f]+)\b(?: <[^>]*>)?$")


def render(lines, base, size, names):
    """Normalise to [(offset, display, compare_key)]: intra-function branches become offsets,
    references to other symbols become their names, everything else is compared literally."""
    out = []
    for a, text, reloc in lines:
        m = BRANCH.match(text)
        if reloc:
            shown = re.sub(r"(?:0x)?[0-9a-f]+(?: <[^>]*>)?$", reloc, text) if m else \
                re.sub(r"0x0\b", reloc, text) if "0x0" in text else f"{text} ; {reloc}"
            key = f"{m.group(1)} {reloc}" if m else re.sub(r"0x[0-9a-f]+", "REF", text)
        elif m and base <= int(m.group(2), 16) < base + size:
            shown = key = f"{m.group(1)} +0x{int(m.group(2), 16) - base:x}"
        elif m:
            tgt = names.get(int(m.group(2), 16), f"0x{int(m.group(2), 16):x}")
            shown = f"{m.group(1)} {tgt}"
            key = f"{m.group(1)} {tgt}"
        else:
            shown = text
            key = text
            # absolute addresses in the original that the map can name
            for h in re.findall(r"0x[0-9a-f]{6,8}", text):
                n = names.get(int(h, 16))
                if n:
                    shown = shown.replace(h, n)
            key = re.sub(r"0x[0-9a-f]{6,8}", "REF", text)
        out.append((a - base, shown, key))
    return out


def cmd_diff(args):
    cfg = load_config()
    smap = load_map()
    addr = smap.lookup(args.name)
    if addr is None:
        sys.exit(f"{args.name}: not in symbol map")
    name = smap.entries[addr][2]
    names = {a: e[2] for a, e in smap.entries.items()}
    funcs = compiled_functions()
    size = smap.entries[addr][0]
    mine = []
    if name in funcs:
        obj, sym, sec, start, end = funcs[name]
        size = max(size, end - start)
        mine = render(disasm_blob(obj, sec, start, end), 0, end - start, {})
    else:
        print(f"{name} has not been compiled yet; showing original only")
    orig = render(disasm(os.path.join(ROOT, cfg["original"]),
                         [f"--start-address=0x{addr:x}", f"--stop-address=0x{addr + size:x}"]),
                  addr, size, names)
    # In the original, a reference to a symbol the compiled side names should compare equal.
    sm = difflib.SequenceMatcher(None, [k for _, _, k in orig], [k for _, _, k in mine], autojunk=False)
    w = 60
    print(f"{'ORIGINAL 0x%08x' % addr:<{w}} | COMPILED {name}")
    for op, i1, i2, j1, j2 in sm.get_opcodes():
        for k in range(max(i2 - i1, j2 - j1)):
            left = orig[i1 + k] if i1 + k < i2 else None
            right = mine[j1 + k] if j1 + k < j2 else None
            mark = " " if op == "equal" else ("|" if left and right else (">" if right else "<"))
            ls = f"{left[0]:5x}: {left[1]}"[:w] if left else ""
            rs = f"{right[0]:5x}: {right[1]}" if right else ""
            print(f"{ls:<{w}} {mark} {rs}")


def cmd_progress(args):
    smap, results, _ = run_checks()
    matched = {r.addr for r in results if r.status in ("MATCH", "MATCH?")}
    game = [(a, e) for a, e in smap.entries.items() if e[1] == "fn"]
    total = sum(e[0] for _, e in game) or 1
    done = sum(e[0] for a, e in game if a in matched)
    print(f"game code: {done}/{total} bytes ({100 * done / total:.3f}%), "
          f"{len(matched & {a for a, _ in game})}/{len(game)} functions")
    lib = [e for e in smap.entries.values() if e[1] == "lib"]
    print(f"library code (excluded): {len(lib)} functions, {sum(e[0] for e in lib)} bytes")


def cmd_link(args):
    """Assemble build/eqgame.exe: compiled code/data placed at their addresses with relocations
    resolved through the symbol map; everything not yet decompiled is carved from orig/."""
    cfg = load_config()
    pe = PEImage(os.path.join(ROOT, cfg["original"]))
    smap, results, _ = run_checks()
    ok = {r.addr for r in results if r.status in ("MATCH", "MATCH?")}
    out = bytearray(pe.data)
    from_src = carved_refs = 0

    def file_off(va):
        for name, sva, vsize, rawptr, rawsize, _ in pe.sections:
            if sva <= va < sva + rawsize:
                return rawptr + va - sva
        return None

    def place(obj, sec, start, end, addr):
        nonlocal from_src, carved_refs
        blob = bytearray(sec.data[start:end])
        for va, si, typ in sec.relocs:
            if not start <= va < end:
                continue
            off, tsym = va - start, obj.symbols[si]
            addend, = struct.unpack_from("<i", blob, off)
            if tsym.section == sec.index:
                target = addr - start + tsym.value
            elif tsym.name in smap.by_name:
                target = smap.by_name[tsym.name]
            else:
                blob[off:off + 4] = pe.read(addr + off, 4)  # unresolved here: keep original field
                carved_refs += 1
                continue
            if typ == REL_I386_DIR32:
                struct.pack_into("<I", blob, off, (target + addend) & 0xFFFFFFFF)
            elif typ == REL_I386_REL32:
                struct.pack_into("<I", blob, off, (target + addend - (addr + off + 4)) & 0xFFFFFFFF)
            else:
                blob[off:off + 4] = pe.read(addr + off, 4)
                carved_refs += 1
        fo = file_off(addr)
        if fo is None:
            return  # uninitialised (.bss) data: nothing in the file
        out[fo:fo + len(blob)] = blob
        from_src += len(blob)

    funcs = compiled_functions()
    for name, (obj, sym, sec, start, end) in funcs.items():
        a = smap.by_name.get(name)
        if a in ok:
            place(obj, sec, start, end, a)
    for r in results:
        if r.status in ("MATCH", "MATCH?"):
            for obj, sec, start, end, a in r.locals:
                place(obj, sec, start, end, a)
    # initialised globals annotated in source
    for a, (name, kind) in smap.annotated.items():
        if kind != "data":
            continue
        for src in source_files():
            op = os.path.join(ROOT, obj_path(src))
            if not os.path.exists(op):
                continue
            obj = CoffObject(op)
            sym = next((x for x in obj.symbols.values() if x.name == name and x.section > 0), None)
            if not sym:
                continue
            sec = obj.sections[sym.section]
            nxt = min([x.value for x in obj.symbols.values()
                       if x.section == sym.section and x.value > sym.value] + [len(sec.data)])
            if sec.chars & 0x80:  # IMAGE_SCN_CNT_UNINITIALIZED_DATA
                break
            place(obj, sec, sym.value, nxt, a)
            break

    dst = os.path.join(BUILD, "eqgame.exe")
    with open(dst, "wb") as f:
        f.write(out)
    h = hashlib.sha1(out).hexdigest()
    code = sum(s[4] for s in pe.sections if s[0] in (".text", ".rdata", ".data"))
    print(f"build/eqgame.exe sha1 {h} {'== original' if h == cfg['sha1'] else '!= ORIGINAL'}")
    print(f"from source: {from_src}/{code} bytes of .text/.rdata/.data ({100 * from_src / code:.3f}%), "
          f"{carved_refs} references still taken from the original")
    if h != cfg["sha1"]:
        bad = next(i for i in range(len(out)) if out[i] != pe.data[i])
        sys.exit(f"first difference at file offset 0x{bad:x}")


def cmd_libscan(args):
    """Compile library probes and find byte-identical copies of their functions in .text."""
    cfg = load_config()
    pe = PEImage(os.path.join(ROOT, cfg["original"]))
    name, tva, tsize, rawptr, _, _ = pe.section(".text")
    text = pe.data[rawptr:rawptr + tsize]
    smap = SymbolMap(os.path.join(CFG_DIR, "symbols.txt"))
    os.makedirs(os.path.join(BUILD, "probes"), exist_ok=True)
    found = {}
    for probe in args.probes:
        out = os.path.join("build", "probes", os.path.splitext(os.path.basename(probe))[0] + ".obj")
        r = subprocess.run(["tools/cl.sh", "/nologo", "/c", *cfg["cflags"].split(), "/Fo" + out, probe],
                           cwd=ROOT, capture_output=True, text=True)
        if r.returncode:
            sys.exit(r.stdout + r.stderr)
        obj = CoffObject(os.path.join(ROOT, out))
        for sym, sec, start, end in obj.functions():
            blob = sec.data[start:end]
            if len(blob) < args.min:
                continue
            mask = bytearray(len(blob))
            for va, _, _ in sec.relocs:
                if start <= va < end:
                    mask[va - start:va - start + 4] = b"\1\1\1\1"
            # longest unmasked run as the search key
            best, cur = (0, 0), None
            for i in range(len(blob) + 1):
                if i < len(blob) and not mask[i]:
                    cur = i if cur is None else cur
                elif cur is not None:
                    if i - cur > best[1] - best[0]:
                        best = (cur, i)
                    cur = None
            if best[1] - best[0] < 6:
                continue
            key = blob[best[0]:best[1]]
            hits = []
            pos = text.find(key)
            while pos >= 0:
                fstart = pos - best[0]
                if 0 <= fstart and fstart + len(blob) <= len(text) and all(
                        mask[i] or text[fstart + i] == blob[i] for i in range(len(blob))):
                    hits.append(tva + fstart)
                pos = text.find(key, pos + 1)
            if sym.name.startswith("$") or sym.name.startswith("__unwindfunclet"):
                continue  # EH funclets belong to whichever object owns the function
            calls = [va - start for va, _, t in sec.relocs if start <= va < end and t == REL_I386_REL32]
            for h in hits:
                targets = [(h + o + 4 + struct.unpack_from("<i", text, h - tva + o)[0]) & 0xFFFFFFFF
                           for o in calls]
                found.setdefault(h, []).append((sym.name, len(blob), len(hits), targets))
    # keep only copies whose callees are all library code (fixpoint: STL calls STL)
    crt = int(str(cfg.get("crt_start", "0x60d000")), 16)  # static CRT occupies the end of .text
    libs = {a for a, e in smap.entries.items() if e[1] == "lib" or a >= crt}
    while True:
        keep = {a for a, c in found.items() if all(t in libs or t in found or t >= crt for t in c[0][3])}
        if keep == set(found):
            break
        found = {a: c for a, c in found.items() if a in keep}
    changed = 0
    for a, cands in sorted(found.items()):
        names = sorted({c[0] for c in cands})
        size = max(c[1] for c in cands)
        if len(names) == 1 and max(c[2] for c in cands) == 1:
            n = names[0]
        else:  # identical code under several names/addresses (e.g. deleting dtors): generic name
            n = f"STL_{a:08x}"
            if args.verbose:
                print(f"  ambiguous 0x{a:08x}: {', '.join(names[:4])}")
        e = smap.entries.get(a)
        cur = e[2] if e else "-"
        print(f"0x{a:08x} {size:5d} {n}  (was {cur})")
        if args.apply and (not e or e[1] != "lib" or e[2] != n):
            if e and e[2] != n and n in smap.by_name and smap.by_name[n] != a:
                n = f"{n}@{a:08x}"  # same instantiation folded elsewhere; keep names unique
            smap.entries[a] = [size if not e else max(e[0], size), "lib", n]
            smap.by_name[n] = a
            changed += 1
    if args.apply:
        smap.save()
    print(f"{len(found)} library function copies found; {changed} map entries updated")

def undecorate(names):
    """MSVC undname.exe (under Wine): mangled -> undecorated signature."""
    out = {}
    names = list(names)
    for i in range(0, len(names), 40):
        chunk = names[i:i + 40]
        r = subprocess.run(["wine", os.path.join(ROOT, "tools", "msvc7.1", "Bin", "undname.exe"), *chunk],
                           capture_output=True, text=True, env={**os.environ, "WINEDEBUG": "-all"})
        for m in re.finditer(r'Undecoration of :- "(.*?)"\s*is :- "(.*?)"', r.stdout.replace("\r", ""), re.S):
            out[m.group(1)] = m.group(2)
    return out


def is_std(signature):
    """True for functions owned by namespace std (STL members, std:: templates, std exceptions)."""
    return re.search(r"(__thiscall|__cdecl|__stdcall|__fastcall|^public: |^protected: |^private: |^)\s*"
                     r"std::", signature) is not None


def masked_blob(sec, start, end):
    blob = sec.data[start:end]
    mask = bytearray(len(blob))
    for va, _, _ in sec.relocs:
        if start <= va < end:
            mask[va - start:va - start + 4] = b"\1\1\1\1"
    return blob, mask


def cmd_pair(args):
    """Find unannotated compiled functions of a source file in an address range and print
    (or append) // FUNCTION: annotations for the ones with exactly one byte-identical match."""
    cfg = load_config()
    pe = PEImage(os.path.join(ROOT, cfg["original"]))
    smap = load_map()
    lo, hi = (int(x, 16) for x in args.range.split("-"))
    code = pe.read(lo, hi - lo)
    obj = CoffObject(os.path.join(ROOT, obj_path(args.src)))
    taken = {a for a in smap.annotated}
    lines = []
    for sym, sec, start, end in obj.functions():
        if not sym.is_external or sym.name in smap.by_name and smap.by_name[sym.name] in taken:
            continue
        blob, mask = masked_blob(sec, start, end)
        if len(blob) < 4:
            continue
        hits = [o for o in range(0, len(code) - len(blob) + 1)
                if code[o] == blob[0] or mask[0]
                if all(mask[i] or code[o + i] == blob[i] for i in range(len(blob)))]
        hits = [lo + o for o in hits if lo + o not in taken]
        if len(hits) == 1:
            lines.append((hits[0], sym.name))
        elif args.verbose:
            print(f"# {sym.name}: {len(hits)} candidates {[hex(h) for h in hits[:5]]}")
    sigs = undecorate(n for _, n in lines)
    funcs = {sym.name: (sym, sec, start, end) for sym, sec, start, end in obj.functions()}

    def refs_ok(a, n):  # relocations must agree with the map (tells folded twins apart)
        sym, sec, start, end = funcs[n]
        return compare(pe, smap, obj, sym, sec, start, end, a).status != "BADREF"

    by_addr = {}
    for a, n in lines:  # identical-code folding: one name per address, game names first
        if not refs_ok(a, n):
            continue
        if a not in by_addr or is_std(sigs.get(by_addr[a], "")) and not is_std(sigs.get(n, "")):
            by_addr[a] = n
    out = []
    for a, n in sorted(by_addr.items()):
        kind = "LIBRARY" if is_std(sigs.get(n, "")) else "FUNCTION"
        out.append(f"// {kind}: EQGAME 0x{a:08x} {n}")
        print(out[-1] + (f"   // {sigs[n]}" if args.verbose and n in sigs else ""))
    lines = out
    if args.write and lines:
        with open(os.path.join(ROOT, args.src), "a") as f:
            f.write("\n// Compiler-generated / template instances, paired by `eq.py pair`.\n")
            for l in lines:
                f.write(l + "\n")
    print(f"{len(lines)} function(s) paired", file=sys.stderr)

def cmd_tus(args):
    """Estimate translation-unit boundaries. Every game .cpp includes a header with a static
    const int[5] {2000,2500,3000,4000,10000}, so each object's .rdata starts with that table.
    Code that references an .rdata block belongs to that object."""
    cfg = load_config()
    pe = PEImage(os.path.join(ROOT, cfg["original"]))
    smap = SymbolMap(os.path.join(CFG_DIR, "symbols.txt"))
    _, rva, rsize, rptr, rraw, _ = pe.section(".rdata")
    data = pe.data[rptr:rptr + rraw]
    pat = struct.pack("<5I", 2000, 2500, 3000, 4000, 10000)
    marks, i = [], data.find(pat)
    while i >= 0:
        marks.append(rva + i)
        i = data.find(pat, i + 4)
    marks.append(rva + rsize)
    dis = os.path.join(BUILD, "eqgame.dis")
    if not os.path.exists(dis):
        _, tva, tsize, _, _, _ = pe.section(".text")
        with open(dis, "w") as f:
            subprocess.run(["i686-w64-mingw32-objdump", "-d", "-M", "intel", "--no-show-raw-insn",
                            f"--start-address=0x{tva:x}", f"--stop-address=0x{tva + tsize:x}",
                            pe_path(cfg)], stdout=f)
    starts = sorted(a for a, e in smap.entries.items() if e[1] in ("fn", "lib", "thunk"))
    import bisect
    def func_of(a):
        k = bisect.bisect_right(starts, a) - 1
        return starts[k] if k >= 0 else a
    users = {}  # rdata target -> referencing functions
    with open(dis) as f:
        for line in f:
            m = re.match(r"\s+([0-9a-f]+):\s+(.*)", line)
            if not m or m.group(2).startswith("f"):
                continue  # x87 constants are shared COMDATs, not owned by one object
            for h in re.findall(r"0x(6[3-9][0-9a-f]{4})\b", m.group(2)):
                v = int(h, 16)
                if marks[0] <= v < marks[-1]:
                    users.setdefault(v, set()).add(func_of(int(m.group(1), 16)))
    counts = {}  # function -> {block: refs}
    for v, fns in users.items():
        if len(fns) > args.max_users:
            continue  # pooled literal / shared constant (/GF): lives in whichever object came first
        k = bisect.bisect_right(marks, v) - 1
        for fn in fns:
            c = counts.setdefault(fn, {})
            c[k] = c.get(k, 0) + 1
    owned = {}
    for fn, c in counts.items():
        k = max(sorted(c), key=lambda b: c[b])
        owned.setdefault(k, []).append(fn)
    out = ["# Estimated translation units from .rdata marker tables (a header's static int[5]).",
           "# RDATA_START RDATA_END CODE_LO CODE_HI NFUNCS   (CODE_* = 10th-90th percentile of the",
           "# functions whose string/vtable references point mostly into the block)"]
    for k in range(len(marks) - 1):
        fs = sorted(owned.get(k, ()))
        if fs:
            lo, hi = fs[len(fs) // 10], fs[(len(fs) * 9) // 10 if len(fs) > 1 else 0]
            out.append(f"0x{marks[k]:08x} 0x{marks[k + 1]:08x} 0x{lo:08x} 0x{hi:08x} {len(fs)}")
        else:
            out.append(f"0x{marks[k]:08x} 0x{marks[k + 1]:08x} - - 0")
    with open(os.path.join(CFG_DIR, "tus.txt"), "w") as f:
        f.write("\n".join(out) + "\n")
    print("\n".join(out[:args.n + 2]))
    print(f"... {len(marks) - 1} blocks written to config/eqgame/tus.txt")

def cmd_map(args):
    smap = SymbolMap(os.path.join(CFG_DIR, "symbols.txt"))
    smap.set(int(args.addr, 16), args.name, args.kind)
    smap.save()


def cmd_todo(args):
    smap, results, _ = run_checks()
    done = {r.addr for r in results if r.status in ("MATCH", "MATCH?")}
    todo = sorted((e[0], a, e[2]) for a, e in smap.entries.items() if e[1] == "fn" and a not in done)
    for size, a, name in todo[:args.n]:
        print(f"0x{a:08x} {size:6d}  {name}")


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("verify")
    s = sub.add_parser("import-ghidra"); s.add_argument("file")
    s = sub.add_parser("build"); s.add_argument("-v", "--verbose", action="store_true")
    s = sub.add_parser("check"); s.add_argument("names", nargs="*")
    s.add_argument("-v", "--verbose", action="store_true")
    s.add_argument("--learn", action="store_true", help="add symbols resolved by matching functions to the map")
    s.add_argument("--inline", action="store_true",
                   help="also list unannotated inline/template functions (pick-any COMDATs)")
    s = sub.add_parser("diff"); s.add_argument("name")
    sub.add_parser("progress")
    sub.add_parser("link")
    s = sub.add_parser("tus"); s.add_argument("n", nargs="?", type=int, default=30)
    s.add_argument("--max-users", type=int, default=4, help="ignore rdata items used by more functions")
    s = sub.add_parser("pair"); s.add_argument("src"); s.add_argument("--range", required=True)
    s.add_argument("--write", action="store_true", help="append the annotations to SRC")
    s.add_argument("-v", "--verbose", action="store_true")
    s = sub.add_parser("libscan"); s.add_argument("probes", nargs="+")
    s.add_argument("--apply", action="store_true", help="mark the copies as kind lib in symbols.txt")
    s.add_argument("--min", type=int, default=16, help="ignore functions shorter than this")
    s.add_argument("-v", "--verbose", action="store_true")
    s = sub.add_parser("map"); s.add_argument("addr"); s.add_argument("name"); s.add_argument("kind", nargs="?")
    s = sub.add_parser("todo"); s.add_argument("n", nargs="?", type=int, default=40)
    args = p.parse_args()
    globals()["cmd_" + args.cmd.replace("-", "_")](args)


if __name__ == "__main__":
    main()