#!/usr/bin/env python3
"""Static client-perf surface scan of every mod jar in the pack.

Usage: jarscan.py <jar-dir> <repo-root> <out.json>

Per mod (including jar-in-jar children, attributed to the parent):
  - mixins: count by side, @Mixin targets, injector kinds (Overwrite/Redirect...)
  - client hot-path classification of mixin targets
  - classes referencing per-frame / per-tick client event types
  - texture pressure: block/item atlas pixels, entity textures, animated textures
  - models, blockstates, shaders (incl. vanilla core shader overrides), sounds
  - geckolib geo/animation files, EMF/OptiFine CEM files, particles
  - registered-content proxies from en_us.json keys
"""
import io, json, os, re, struct, sys, zipfile, tomllib
from collections import Counter, defaultdict

# ---------- class file parsing (constant pool + annotations) ----------

def parse_class(data):
    """Return (utf8 set, class annotations, method annotations) for a class file.
    Annotations are dicts {type: {elem: value}} for RuntimeInvisible/Visible."""
    if data[:4] != b'\xca\xfe\xba\xbe':
        return None
    pos = 10
    count = struct.unpack('>H', data[8:10])[0]
    cp = [None] * count
    i = 1
    while i < count:
        tag = data[pos]
        if tag == 1:
            ln = struct.unpack('>H', data[pos+1:pos+3])[0]
            cp[i] = ('utf8', data[pos+3:pos+3+ln].decode('utf-8', 'replace'))
            pos += 3 + ln
        elif tag in (3, 4):
            cp[i] = ('num', None); pos += 5
        elif tag in (5, 6):
            cp[i] = ('num', None); pos += 9; i += 1
        elif tag == 7:
            cp[i] = ('class', struct.unpack('>H', data[pos+1:pos+3])[0]); pos += 3
        elif tag in (8, 16, 19, 20):
            cp[i] = ('ref', struct.unpack('>H', data[pos+1:pos+3])[0]); pos += 3
        elif tag in (9, 10, 11, 12, 17, 18):
            cp[i] = ('ref2', None); pos += 5
        elif tag == 15:
            cp[i] = ('mh', None); pos += 4
        else:
            return None
        i += 1

    def u(idx):
        e = cp[idx]
        return e[1] if e and e[0] == 'utf8' else None

    utf8 = {e[1] for e in cp if e and e[0] == 'utf8'}

    def read_elem(p):
        tag = chr(data[p]); p += 1
        if tag in 'BCDFIJSZs':
            return u(struct.unpack('>H', data[p:p+2])[0]), p + 2
        if tag == 'e':
            return u(struct.unpack('>H', data[p+2:p+4])[0]), p + 4
        if tag == 'c':
            return u(struct.unpack('>H', data[p:p+2])[0]), p + 2
        if tag == '@':
            return read_ann(p)
        if tag == '[':
            n = struct.unpack('>H', data[p:p+2])[0]; p += 2
            out = []
            for _ in range(n):
                v, p = read_elem(p); out.append(v)
            return out, p
        raise ValueError(tag)

    def read_ann(p):
        t = u(struct.unpack('>H', data[p:p+2])[0])
        n = struct.unpack('>H', data[p+2:p+4])[0]; p += 4
        elems = {}
        for _ in range(n):
            name = u(struct.unpack('>H', data[p:p+2])[0]); p += 2
            v, p = read_elem(p); elems[name] = v
        return (t, elems), p

    def read_attrs(p):
        anns = []
        n = struct.unpack('>H', data[p:p+2])[0]; p += 2
        for _ in range(n):
            name = u(struct.unpack('>H', data[p:p+2])[0])
            ln = struct.unpack('>I', data[p+2:p+6])[0]
            body = p + 6
            if name in ('RuntimeInvisibleAnnotations', 'RuntimeVisibleAnnotations'):
                k = struct.unpack('>H', data[body:body+2])[0]; q = body + 2
                for _ in range(k):
                    a, q = read_ann(q); anns.append(a)
            p = body + ln
        return anns, p

    try:
        p = pos + 6
        ni = struct.unpack('>H', data[p:p+2])[0]; p += 2 + 2 * ni
        method_anns = []
        for _member in ('fields', 'methods'):
            n = struct.unpack('>H', data[p:p+2])[0]; p += 2
            for _ in range(n):
                anns, p = read_attrs(p + 6)
                if _member == 'methods':
                    method_anns.extend(anns)
        class_anns, p = read_attrs(p)
    except Exception:
        return utf8, [], []
    return utf8, class_anns, method_anns


# ---------- classification ----------

HOT = [  # (category, regex on internal class name of mixin target)
    ('blaze3d', r'^com/mojang/blaze3d/'),
    ('sodium', r'^net/caffeinemc/'),
    ('iris', r'^net/irisshaders/'),
    ('world_render', r'^net/minecraft/client/renderer/(LevelRenderer|GameRenderer|LightTexture|FogRenderer|DimensionSpecialEffects|SectionOcclusionGraph|ViewArea|chunk/|culling/|PostChain|RenderType|RenderStateShard|ShaderInstance|LevelEventHandler)'),
    ('entity_render', r'^net/minecraft/client/(renderer/entity/|model/|renderer/blockentity/|renderer/ItemInHandRenderer)'),
    ('item_render', r'^net/minecraft/client/renderer/(ItemRenderer|item/)'),
    ('block_model', r'^net/minecraft/client/(renderer/block/|resources/model/|renderer/ItemBlockRenderTypes)'),
    ('texture_load', r'^net/minecraft/client/(renderer/texture/|resources/)'),
    ('camera', r'^net/minecraft/client/Camera'),
    ('gui', r'^net/minecraft/client/gui/'),
    ('particle', r'^net/minecraft/client/particle/'),
    ('sound', r'^(net/minecraft/client/sounds/|com/mojang/blaze3d/audio/)'),
    ('client_core', r'^net/minecraft/client/(Minecraft|multiplayer/|player/|KeyboardHandler|MouseHandler|Options)'),
    ('client_other', r'^net/minecraft/client/'),
    ('entity_common', r'^net/minecraft/world/entity/'),
    ('level_common', r'^net/minecraft/world/level/'),
    ('server', r'^net/minecraft/server/'),
    ('mc_other', r'^net/minecraft/'),
    ('other_mod', r'.'),
]
HOT_RE = [(c, re.compile(r)) for c, r in HOT]
FRAME_CATS = {'blaze3d', 'sodium', 'iris', 'world_render', 'entity_render', 'item_render', 'camera', 'gui', 'particle'}

def classify(target):
    for c, r in HOT_RE:
        if r.search(target):
            return c
    return 'other_mod'

EVENTS = {
    # per-frame
    'RenderLevelStageEvent': 'frame', 'RenderFrameEvent': 'frame', 'RenderGuiEvent': 'frame',
    'RenderGuiLayerEvent': 'frame', 'CustomizeGuiOverlayEvent': 'frame', 'RenderLivingEvent': 'frame',
    'RenderPlayerEvent': 'frame', 'RenderHandEvent': 'frame', 'RenderArmEvent': 'frame',
    'RenderHighlightEvent': 'frame', 'ViewportEvent': 'frame', 'RenderNameTagEvent': 'frame',
    'RenderBlockScreenEffectEvent': 'frame', 'RenderTooltipEvent': 'frame', 'ScreenEvent': 'screen',
    'ComputeFovModifierEvent': 'frame',
    # per-tick
    'ClientTickEvent': 'tick', 'PlayerTickEvent': 'tick', 'EntityTickEvent': 'tick', 'LevelTickEvent': 'tick',
    # load
    'ModelEvent': 'load', 'RegisterShadersEvent': 'load', 'TextureAtlasStitchedEvent': 'load',
}
EVENT_RE = re.compile(r'net/neoforged/neoforge/(?:client/)?event/(?:[a-z/]+/)?(' + '|'.join(EVENTS) + r')(?:\$(\w+))?')
OTHER_FRAME_RE = re.compile(r'(dev/architectury/event/events/client/(ClientGuiEvent|ClientTickEvent)|net/fabricmc/fabric/api/client/rendering/v1/(WorldRenderEvents|HudRenderCallback|HudLayerRegistrationCallback)|net/blay09/mods/balm/api/event/client/(GuiDrawEvent|RenderHandEvent|ClientTickHandler))')

INJECTORS = {
    'Lorg/spongepowered/asm/mixin/Overwrite;': 'Overwrite',
    'Lorg/spongepowered/asm/mixin/injection/Redirect;': 'Redirect',
    'Lorg/spongepowered/asm/mixin/injection/Inject;': 'Inject',
    'Lorg/spongepowered/asm/mixin/injection/ModifyVariable;': 'ModifyVariable',
    'Lorg/spongepowered/asm/mixin/injection/ModifyArg;': 'ModifyArg',
    'Lorg/spongepowered/asm/mixin/injection/ModifyArgs;': 'ModifyArgs',
    'Lorg/spongepowered/asm/mixin/injection/ModifyConstant;': 'ModifyConstant',
    'Lcom/llamalad7/mixinextras/injector/wrapoperation/WrapOperation;': 'WrapOperation',
    'Lcom/llamalad7/mixinextras/injector/wrapmethod/WrapMethod;': 'WrapMethod',
    'Lcom/llamalad7/mixinextras/injector/ModifyExpressionValue;': 'ModifyExpressionValue',
    'Lcom/llamalad7/mixinextras/injector/ModifyReturnValue;': 'ModifyReturnValue',
    'Lcom/llamalad7/mixinextras/injector/WrapWithCondition;': 'WrapWithCondition',
    'Lcom/llamalad7/mixinextras/injector/v2/WrapWithCondition;': 'WrapWithCondition',
}


def png_size(b):
    if b[:8] == b'\x89PNG\r\n\x1a\n' and b[12:16] == b'IHDR':
        return struct.unpack('>II', b[16:24])
    return None


def load_json_lenient(b):
    s = b.decode('utf-8', 'replace')
    s = re.sub(r'^\s*//.*$', '', s, flags=re.M)
    try:
        return json.loads(s)
    except Exception:
        return None


def mod_ids(z):
    for name in ('META-INF/neoforge.mods.toml', 'META-INF/mods.toml'):
        if name in z.namelist():
            try:
                t = tomllib.loads(z.read(name).decode('utf-8', 'replace'))
            except Exception:
                return [], []
            ids = [m.get('modId') for m in t.get('mods', [])]
            mix = [m.get('config') for m in t.get('mixins', []) if m.get('config')]
            return ids, mix
    return [], []


def scan_zip(z, acc, depth=0):
    names = z.namelist()
    nameset = set(names)
    ids, mixcfgs = mod_ids(z)
    acc['mod_ids'].extend(i for i in ids if i)
    if 'META-INF/MANIFEST.MF' in nameset:
        mf = z.read('META-INF/MANIFEST.MF').decode('utf-8', 'replace').replace('\r\n ', '').replace('\n ', '')
        m = re.search(r'^MixinConfigs:\s*(.+)$', mf, re.M)
        if m:
            mixcfgs += [c.strip() for c in m.group(1).split(',')]
    if 'fabric.mod.json' in nameset:
        fj = load_json_lenient(z.read('fabric.mod.json')) or {}
        for c in fj.get('mixins', []):
            mixcfgs.append(c if isinstance(c, str) else c.get('config'))
    mixcfgs = sorted({c for c in mixcfgs if c and c in nameset})

    # mixin classes
    mixin_classes = {}  # path -> side
    for cfg in mixcfgs:
        j = load_json_lenient(z.read(cfg)) or {}
        pkg = (j.get('package') or '').replace('.', '/')
        for side in ('mixins', 'client', 'server'):
            for c in j.get(side, []) or []:
                mixin_classes[f'{pkg}/{c.replace(".", "/")}.class'] = {'mixins': 'common'}.get(side, side)
    acc['mixin_configs'] += len(mixcfgs)

    tex = acc['tex']
    for n in names:
        if n.endswith('/'):
            continue
        if n.endswith('.class'):
            try:
                data = z.read(n)
            except Exception:
                continue
            parsed = parse_class(data)
            if not parsed:
                continue
            utf8, canns, manns = parsed
            if n in mixin_classes:
                side = mixin_classes[n]
                acc['mixins'][side] += 1
                targets = []
                for t, el in canns:
                    if t == 'Lorg/spongepowered/asm/mixin/Mixin;':
                        for v in el.get('value', []) or []:
                            if v: targets.append(v[1:-1])
                        for v in el.get('targets', []) or []:
                            if v: targets.append(v.replace('.', '/'))
                kinds = Counter(INJECTORS[t] for t, _ in manns if t in INJECTORS)
                acc['injectors'].update(kinds)
                for tg in targets or ['?']:
                    cat = classify(tg)
                    acc['targets'][tg] += 1
                    acc['target_cats'][cat] += 1
                    if cat in FRAME_CATS:
                        acc['frame_mixins'] += 1
                        acc['frame_targets'][tg.rsplit('/', 1)[-1]] += 1
                        acc['frame_injectors'].update(kinds)
            evs = set()
            for s in utf8:
                for m in EVENT_RE.finditer(s):
                    evs.add(m.group(1) + ('$' + m.group(2) if m.group(2) else ''))
                m = OTHER_FRAME_RE.search(s)
                if m:
                    evs.add(m.group(1).rsplit('/', 1)[-1])
            for e in evs:
                acc['events'][e] += 1
            continue
        if n.startswith(('META-INF/jarjar/', 'META-INF/jars/')) and n.endswith('.jar') and depth < 3:
            try:
                sub = zipfile.ZipFile(io.BytesIO(z.read(n)))
                subnames = set(sub.namelist())
                if 'fabric.mod.json' in subnames and not subnames & {'META-INF/neoforge.mods.toml', 'META-INF/mods.toml'}:
                    continue  # Fabric-only payload of a multi-loader jar; never loaded on NeoForge
                acc['nested'].append(n.rsplit('/', 1)[-1])
                scan_zip(sub, acc, depth + 1)
            except Exception:
                pass
            continue
        parts = n.split('/')
        if parts[0] != 'assets' or len(parts) < 3:
            continue
        ns, kind = parts[1], parts[2]
        if kind == 'textures' and n.endswith('.png'):
            sub = parts[3] if len(parts) > 4 else ''
            info = z.getinfo(n)
            with z.open(n) as fh:
                wh = png_size(fh.read(24))
            px = wh[0] * wh[1] if wh else 0
            anim = None
            if wh and n + '.mcmeta' in nameset:
                mm = load_json_lenient(z.read(n + '.mcmeta')) or {}
                anim = mm.get('animation') if isinstance(mm, dict) else None
                if isinstance(anim, dict):
                    # only one frame lives in the atlas; the rest is CPU-side upload source
                    fw = anim.get('width', min(wh)); fh = anim.get('height', fw if 'width' in anim else min(wh))
                    frames = max(1, (wh[0] // max(1, fw)) * (wh[1] // max(1, fh)))
                    tex['anim_frames'] += frames
                    px = fw * fh
            bucket = {'block': 'atlas', 'item': 'atlas', 'blocks': 'atlas', 'items': 'atlas',
                      'entity': 'entity', 'models': 'entity', 'particle': 'particle',
                      'gui': 'gui', 'environment': 'env', 'painting': 'painting', 'mob_effect': 'gui'}.get(sub, 'other')
            tex[bucket + '_n'] += 1
            tex[bucket + '_px'] += px
            tex['bytes'] += info.file_size
            if wh and max(wh) > 512 and bucket in ('atlas', 'entity'):
                acc['big_textures'].append(f'{n} {wh[0]}x{wh[1]}' + (' (animated)' if anim else ''))
            if isinstance(anim, dict) and bucket == 'atlas':
                tex['animated_n'] += 1
        elif kind == 'models' and n.endswith('.json'):
            acc['assets']['models'] += 1
        elif kind == 'blockstates' and n.endswith('.json'):
            acc['assets']['blockstates'] += 1
        elif kind == 'shaders':
            if n.endswith(('.vsh', '.fsh', '.glsl', '.json')):
                acc['assets']['shader_files'] += 1
                if ns == 'minecraft' and parts[3] == 'core':
                    acc['core_shader_overrides'].append(n)
                if parts[3] == 'post':
                    acc['assets']['post_effects'] += 1
        elif kind == 'sounds' and n.endswith('.ogg'):
            acc['assets']['ogg_n'] += 1
            acc['assets']['ogg_bytes'] += z.getinfo(n).file_size
        elif kind in ('geo', 'geckolib') or (kind == 'animations' and n.endswith('.json')):
            acc['assets']['geckolib_files'] += 1
        elif kind == 'particles' and n.endswith('.json'):
            acc['assets']['particle_defs'] += 1
        elif kind == 'optifine' or kind == 'emf' or kind == 'etf':
            acc['assets']['cem_files'] += 1
        elif kind == 'lang' and parts[-1] == 'en_us.json':
            j = load_json_lenient(z.read(n)) or {}
            for k in j:
                p0 = k.split('.', 1)[0]
                if p0 in ('block', 'item', 'entity', 'biome', 'enchantment', 'effect'):
                    acc['lang'][p0] += 1


def new_acc():
    return {'mod_ids': [], 'nested': [], 'mixin_configs': 0, 'mixins': Counter(),
            'targets': Counter(), 'target_cats': Counter(), 'frame_mixins': 0,
            'frame_targets': Counter(), 'frame_injectors': Counter(), 'injectors': Counter(),
            'events': Counter(), 'tex': Counter(), 'big_textures': [], 'assets': Counter(),
            'core_shader_overrides': [], 'lang': Counter()}


def main():
    jar_dir, repo, out = sys.argv[1:4]
    meta = {}
    for d in ('mods', 'resourcepacks', 'shaderpacks'):
        for f in os.listdir(os.path.join(repo, d)):
            if not f.endswith('.pw.toml'):
                continue
            t = tomllib.load(open(os.path.join(repo, d, f), 'rb'))
            meta[t['filename']] = {'slug': f[:-8], 'name': t['name'], 'side': t.get('side', 'both'), 'dir': d,
                                   'optional': bool(t.get('option', {}).get('optional'))}
    results = {}
    for fn, m in sorted(meta.items()):
        path = os.path.join(jar_dir, fn)
        if not os.path.exists(path):
            print('missing', fn, file=sys.stderr)
            continue
        acc = new_acc()
        try:
            scan_zip(zipfile.ZipFile(path), acc)
        except zipfile.BadZipFile:
            print('badzip', fn, file=sys.stderr)
            continue
        acc['jar_bytes'] = os.path.getsize(path)
        acc['big_textures'] = acc['big_textures'][:15]
        acc['targets'] = dict(acc['targets'].most_common(40))
        results[m['slug']] = {**m, 'filename': fn, **{k: (dict(v) if isinstance(v, Counter) else v) for k, v in acc.items()}}
    json.dump(results, open(out, 'w'), indent=1, sort_keys=True)
    print(f'scanned {len(results)} files -> {out}')


if __name__ == '__main__':
    main()