#!/usr/bin/env python3
"""pzxref: índice cruzado del bytecode de Project Zomboid (curso, módulo 08b).

Crear el índice (una vez por versión del juego, tarda unos minutos):
    python3 pzxref.py build ~/pz-java/jar            # carpeta con el .jar descomprimido
Consultas:
    python3 pzxref.py str Rotten                     # qué métodos usan una cadena (busca dentro)
    python3 pzxref.py callers LuaHookManager.TriggerHook
    python3 pzxref.py calls zombie/inventory/ItemPickerJava.getActualSpawnChance
    python3 pzxref.py fields zombieDensityCap        # quién lee (get) o escribe (put) un campo
    python3 pzxref.py events [OnFillContainer]       # eventos de Lua que dispara Java y desde dónde
El índice queda en pzxref.db (SQLite) junto a este archivo, o en la ruta de PZXREF_DB.
"""
import os, re, sqlite3, subprocess, sys

DB = os.environ.get('PZXREF_DB') or os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pzxref.db')
RE_CLASS = re.compile(r'^(?:[\w ]+ )?(?:class|interface|enum) ([\w.$]+)')
RE_METHOD = re.compile(r'^  (?=\S)(?!Code:)(.*\(.*\).*);$|^  static \{\};$')
RE_INSTR = re.compile(r'^\s+\d+: (\w+)')
RE_STR = re.compile(r'// String (.*)$')
RE_CALL = re.compile(r'// (?:Interface)?Method ([\w/$]+\.)?([\w$<>"]+):')
RE_FIELD = re.compile(r'// Field ([\w/$]+\.)?(\w+):')
KNOWN = set()
TRIGGERS = ('LuaEventManager.triggerEvent', 'LuaEventManager.triggerEventGarbage',
            'LuaEventManager.triggerEventUnique', 'LuaHookManager.TriggerHook')


def build(jardir, prefix='zombie'):
    classes = []
    for root, _, files in os.walk(os.path.join(jardir, prefix)):
        for f in files:
            if f.endswith('.class'):
                rel = os.path.relpath(os.path.join(root, f), jardir)[:-6]
                classes.append(rel.replace(os.sep, '.'))
    classes.sort()
    if os.path.exists(DB):
        os.remove(DB)
    db = sqlite3.connect(DB)
    db.executescript('''
        CREATE TABLE str(cls, meth, val);
        CREATE TABLE call(cls, meth, owner, name);
        CREATE TABLE field(cls, meth, op, owner, name);
        CREATE TABLE event(cls, meth, kind, name);
    ''')
    env = dict(os.environ, JAVA_TOOL_OPTIONS='')
    # nombres de eventos y ganchos registrados (para reconocer el nombre en cada llamada)
    reg = subprocess.run(['javap', '-p', '-c', '-cp', jardir, 'zombie.Lua.LuaEventManager',
                          'zombie.Lua.LuaHookManager'], capture_output=True, text=True, env=env).stdout
    KNOWN.update(re.findall(r'// String (\S+)', reg))
    for i in range(0, len(classes), 200):
        batch = classes[i:i + 200]
        out = subprocess.run(['javap', '-p', '-c', '-cp', jardir] + batch,
                             capture_output=True, text=True, env=env).stdout
        parse(out, db)
        print(f'\r{min(i + 200, len(classes))}/{len(classes)} clases', end='', file=sys.stderr)
    for t, cols in (('str', 'val'), ('call', 'name'), ('field', 'name'), ('event', 'name')):
        db.execute(f'CREATE INDEX i_{t} ON {t}({cols})')
    db.commit()
    print(f'\nlisto: {DB}', file=sys.stderr)


def parse(out, db):
    cls = meth = None
    last_str = None
    rows_s, rows_c, rows_f, rows_e = [], [], [], []
    for line in out.splitlines():
        m = RE_CLASS.match(line)
        if m and not line.startswith(' '):
            cls = m.group(1)
            continue
        m = RE_METHOD.match(line)
        if m:
            meth = (m.group(1) or 'static {}').strip()
            meth = re.sub(r'^.*? (\S+\()', r'\1', meth)          # sin modificadores ni tipo de retorno
            last_str = None
            continue
        if not RE_INSTR.match(line):
            continue
        op = RE_INSTR.match(line).group(1)
        m = RE_STR.search(line)
        if m:
            if m.group(1) in KNOWN:
                last_str = m.group(1)          # el nombre del evento es el primer argumento
            rows_s.append((cls, meth, m.group(1)))
            continue
        m = RE_CALL.search(line)
        if m and op.startswith('invoke'):
            owner = (m.group(1) or cls.replace('.', '/') + '.')[:-1].replace('/', '.')
            name = m.group(2)
            rows_c.append((cls, meth, owner, name))
            short = owner.rsplit('.', 1)[-1] + '.' + name
            if short in TRIGGERS and last_str:
                rows_e.append((cls, meth, 'hook' if 'Hook' in short else 'event', last_str))
                last_str = None
            continue
        m = RE_FIELD.search(line)
        if m and op in ('getfield', 'putfield', 'getstatic', 'putstatic'):
            owner = (m.group(1) or cls.replace('.', '/') + '.')[:-1].replace('/', '.')
            rows_f.append((cls, meth, op[:3], owner, m.group(2)))
    db.executemany('INSERT INTO str VALUES (?,?,?)', rows_s)
    db.executemany('INSERT INTO call VALUES (?,?,?,?)', rows_c)
    db.executemany('INSERT INTO field VALUES (?,?,?,?,?)', rows_f)
    db.executemany('INSERT INTO event VALUES (?,?,?,?)', rows_e)


def q(sql, *args):
    db = sqlite3.connect(DB)
    for row in db.execute(sql, args):
        print('  '.join(str(x) for x in row))


def split_target(t):
    t = t.replace('/', '.')
    owner, _, name = t.rpartition('.')
    return owner, name


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return
    cmd, args = argv[1], argv[2:]
    if cmd == 'build':
        build(os.path.expanduser(args[0]))
    elif cmd == 'str':
        q('SELECT DISTINCT cls, meth, val FROM str WHERE val LIKE ? ORDER BY cls', f'%{args[0]}%')
    elif cmd == 'callers':
        owner, name = split_target(args[0])
        q('SELECT DISTINCT cls, meth FROM call WHERE name = ? AND owner LIKE ? ORDER BY cls',
          name, f'%{owner}')
    elif cmd == 'calls':
        owner, name = split_target(args[0])
        q('SELECT owner, name, count(*) FROM call WHERE cls LIKE ? AND meth LIKE ? '
          'GROUP BY owner, name ORDER BY 3 DESC', f'%{owner}', f'{name}(%')
    elif cmd == 'fields':
        q('SELECT DISTINCT op, owner, name, cls, meth FROM field WHERE name = ? ORDER BY op, cls', args[0])
    elif cmd == 'events':
        if args:
            q('SELECT DISTINCT kind, name, cls, meth FROM event WHERE name = ? ORDER BY cls', args[0])
        else:
            q('SELECT kind, name, count(*) FROM event GROUP BY kind, name ORDER BY kind, name')
    else:
        print(__doc__)


if __name__ == '__main__':
    try:
        main(sys.argv)
    except BrokenPipeError:            # pzxref.py ... | head
        pass