package com.spineview.dec;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.lang.invoke.CallSite;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.reflect.Constructor;
import java.lang.reflect.Method;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipFile;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.Handle;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.InvokeDynamicInsnNode;
import org.objectweb.asm.tree.MethodNode;

/** 把 (类,方法) 里的所有 invokedynamic 站点解出来，打印 站点描述符 + 目标方法签名。 */
public final class ResolveSites {
  private static final String BSM_DESC =
      "(Ljava/lang/Object;Ljava/lang/Object;Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;";

  public static void main(String[] args) throws Exception {
    String payload = args[0];
    List<String> filters = new ArrayList<String>();
    for (int i = 1; i < args.length; i++) filters.add(args[i]);
    List<Site> sites = new ArrayList<Site>();
    try (ZipFile zip = new ZipFile(payload)) {
      Enumeration<? extends ZipEntry> entries = zip.entries();
      while (entries.hasMoreElements()) {
        ZipEntry entry = entries.nextElement();
        if (!entry.getName().endsWith(".class")) continue;
        ClassNode node;
        try { node = parse(readAll(zip.getInputStream(entry))); } catch (Throwable t) { continue; }
        if (!matches(node.name, filters)) continue;
        System.out.println("== class " + node.name + " extends " + node.superName);
        for (MethodNode method : node.methods) {
          for (AbstractInsnNode insn = method.instructions.getFirst(); insn != null;
              insn = insn.getNext()) {
            if (!(insn instanceof InvokeDynamicInsnNode)) continue;
            InvokeDynamicInsnNode indy = (InvokeDynamicInsnNode) insn;
            Handle bsm = indy.bsm;
            if (!BSM_DESC.equals(bsm.getDesc())) continue;
            sites.add(new Site(node.name, method.name, method.desc, bsm.getOwner(), bsm.getName(),
                indy.name, indy.desc, indy.bsmArgs.length > 0 ? indy.bsmArgs[0] : null));
          }
        }
      }
    }
    Map<String, Method> bootstrapCache = new HashMap<String, Method>();
    MethodHandles.Lookup[] lookupCache = new MethodHandles.Lookup[1];
    Map<String, String> resolved = new LinkedHashMap<String, String>();
    for (Site site : sites) {
      String triple = site.bsmOwner + "." + site.bsmName + "#" + site.indyName;
      String target = resolved.get(triple);
      if (target == null) {
        try { target = resolve(site, bootstrapCache, lookupCache); }
        catch (Throwable error) { target = "<失败> " + error; }
        resolved.put(triple, target);
      }
      System.out.println("  " + site.owner + "." + site.method + site.methodDesc
          + "\n    indy " + site.bsmOwner + "." + site.bsmName + " name=" + site.indyName
          + " siteType=" + site.desc + "\n    -> " + target);
    }
    System.out.println("站点 " + sites.size());
  }

  private static boolean matches(String name, List<String> filters) {
    if (filters.isEmpty()) return true;
    for (String f : filters) if (name.equals(f)) return true;
    return false;
  }

  private static String resolve(Site site, Map<String, Method> bootstrapCache,
      MethodHandles.Lookup[] lookupCache) throws Throwable {
    Class<?> owner = Class.forName(site.bsmOwner.replace('/', '.'));
    String key = site.bsmOwner + "." + site.bsmName;
    Method bootstrap = bootstrapCache.get(key);
    if (bootstrap == null) {
      for (Method candidate : owner.getDeclaredMethods()) {
        if (candidate.getName().equals(site.bsmName)
            && candidate.getParameterTypes().length == 4) {
          bootstrap = candidate;
          break;
        }
      }
      if (bootstrap == null) throw new NoSuchMethodException(site.bsmName);
      bootstrap.setAccessible(true);
      bootstrapCache.put(key, bootstrap);
    }
    if (lookupCache[0] == null) lookupCache[0] = trustedLookup();
    MethodType type = MethodType.fromMethodDescriptorString(site.desc,
        ResolveSites.class.getClassLoader());
    Object callSite = bootstrap.invoke(null, lookupCache[0], site.indyName, type, site.arg);
    if (!(callSite instanceof CallSite)) return String.valueOf(callSite);
    MethodHandle target = ((CallSite) callSite).getTarget();
    return describe(target) + "  [type=" + target.type() + "]";
  }

  private static String describe(MethodHandle target) {
    String direct = memberOf(target);
    if (direct != null) return direct;
    for (Class<?> type = target.getClass(); type != null; type = type.getSuperclass()) {
      for (java.lang.reflect.Field field : type.getDeclaredFields()) {
        if (!field.getName().startsWith("argL") && !field.getName().startsWith("argJ")
            && !field.getName().startsWith("argI")) continue;
        try {
          field.setAccessible(true);
          Object value = field.get(target);
          if (!(value instanceof MethodHandle)) continue;
          String nested = memberOf((MethodHandle) value);
          if (nested != null) return nested + " (via " + target.getClass().getSimpleName() + ")";
        } catch (Throwable ignored) { }
      }
    }
    return "?" + target.type();
  }

  private static String memberOf(MethodHandle handle) {
    try {
      Method internal = MethodHandle.class.getDeclaredMethod("internalMemberName");
      internal.setAccessible(true);
      Object member = internal.invoke(handle);
      if (member == null) return null;
      Class<?> memberClass = member.getClass();
      Method getDeclaring = memberClass.getMethod("getDeclaringClass");
      Method getName = memberClass.getMethod("getName");
      Method isStatic = memberClass.getMethod("isStatic");
      getDeclaring.setAccessible(true);
      getName.setAccessible(true);
      isStatic.setAccessible(true);
      Object declaring = getDeclaring.invoke(member);
      String owner = String.valueOf(declaring).replace("class ", "").replace("interface ", "");
      return owner + "." + getName.invoke(member)
          + (Boolean.TRUE.equals(isStatic.invoke(member)) ? " (static)" : "");
    } catch (Throwable error) {
      return null;
    }
  }

  private static MethodHandles.Lookup trustedLookup() throws Exception {
    Constructor<MethodHandles.Lookup> constructor =
        MethodHandles.Lookup.class.getDeclaredConstructor(Class.class, int.class);
    constructor.setAccessible(true);
    return constructor.newInstance(Object.class, Integer.valueOf(-1));
  }

  private static final class Site {
    final String owner, method, methodDesc, bsmOwner, bsmName, indyName, desc;
    final Object arg;
    Site(String owner, String method, String methodDesc, String bsmOwner, String bsmName,
        String indyName, String desc, Object arg) {
      this.owner = owner; this.method = method; this.methodDesc = methodDesc;
      this.bsmOwner = bsmOwner; this.bsmName = bsmName; this.indyName = indyName;
      this.desc = desc; this.arg = arg;
    }
  }

  private static ClassNode parse(byte[] raw) {
    ClassNode node = new ClassNode();
    new ClassReader(raw).accept(node, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
    return node;
  }

  private static byte[] readAll(InputStream in) throws java.io.IOException {
    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
    byte[] chunk = new byte[65536];
    int n;
    while ((n = in.read(chunk)) > 0) buffer.write(chunk, 0, n);
    return buffer.toByteArray();
  }
}