package bot.inker.security.asm;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

public final class WithClassLoaderClassWriter extends ClassWriter {
  private final ClassLoader classLoader;
  private final List<String> missingClasses = new ArrayList<>();

  public WithClassLoaderClassWriter(int flags) {
    super(flags);
    this.classLoader = ClassLoader.getSystemClassLoader();
  }

  public WithClassLoaderClassWriter(ClassLoader classLoader, int flags) {
    super(flags);
    this.classLoader = classLoader;
  }

  public WithClassLoaderClassWriter(ClassLoader classLoader, ClassReader classReader, int flags) {
    super(classReader, flags);
    this.classLoader = classLoader;
  }

  @Override
  protected String getCommonSuperClass(String type1, String type2) {
    try {
      ClassInfo info1 = typeInfo(type1);
      ClassInfo info2 = typeInfo(type2);
      if ((info1.access() & Opcodes.ACC_INTERFACE) != 0) {
        if (typeImplements(type2, info2, type1)) {
          return type1;
        } else {
          return "java/lang/Object";
        }
      }
      if ((info2.access() & Opcodes.ACC_INTERFACE) != 0) {
        if (typeImplements(type1, info1, type2)) {
          return type2;
        } else {
          return "java/lang/Object";
        }
      }
      StringBuilder b1 = typeAncestors(type1, info1);
      StringBuilder b2 = typeAncestors(type2, info2);
      String result = "java/lang/Object";
      int end1 = b1.length();
      int end2 = b2.length();
      while (true) {
        int start1 = b1.lastIndexOf(";", end1 - 1);
        int start2 = b2.lastIndexOf(";", end2 - 1);
        if (start1 != -1 && start2 != -1
          && end1 - start1 == end2 - start2) {
          String p1 = b1.substring(start1 + 1, end1);
          String p2 = b2.substring(start2 + 1, end2);
          if (p1.equals(p2)) {
            result = p1;
            end1 = start1;
            end2 = start2;
          } else {
            return result;
          }
        } else {
          return result;
        }
      }
    } catch (IOException e) {
      throw new RuntimeException(e.toString());
    }
  }

  private StringBuilder typeAncestors(String type, ClassInfo info)
    throws IOException {
    StringBuilder b = new StringBuilder();
    while (!"java/lang/Object".equals(type)) {
      b.append(';').append(type);
      type = info.superName();
      info = typeInfo(type);
    }
    return b;
  }

  private boolean typeImplements(String type, ClassInfo info, String itf)
    throws IOException {
    while (!"java/lang/Object".equals(type)) {
      String[] itfs = info.interfaces();
      for (String it : itfs) {
        if (it.equals(itf)) {
          return true;
        }
      }
      for (String it : itfs) {
        if (typeImplements(it, typeInfo(it), itf)) {
          return true;
        }
      }
      type = info.superName();
      info = typeInfo(type);
    }
    return false;
  }

  private ClassInfo typeInfo(final String type) throws IOException {
    try (InputStream in = getClassLoader().getResourceAsStream(type + ".class")) {
      if (in == null) {
        try {
          Class<?> clazz = Class.forName(type.replace('/', '.'), false, getClassLoader());
          return new ClassInfoClassInstance(clazz);
        } catch (ClassNotFoundException e) {
          throw new TypeNotPresentException(type, new IOException("Resource not found: " + type + ".class"));
        }
      }
      return new ClassInfoClassReader(new ClassReader(in));
    }
  }

  @Override
  protected ClassLoader getClassLoader() {
    return classLoader;
  }

  public interface ClassInfo {
    int access();

    String name();

    String superName();

    String[] interfaces();
  }

  public static class ClassInfoClassReader implements ClassInfo {
    private final ClassReader classReader;

    public ClassInfoClassReader(ClassReader classReader) {
      this.classReader = classReader;
    }

    @Override
    public int access() {
      return classReader.getAccess();
    }

    @Override
    public String name() {
      return classReader.getClassName();
    }

    @Override
    public String superName() {
      return classReader.getSuperName();
    }

    @Override
    public String[] interfaces() {
      return classReader.getInterfaces();
    }
  }

  public static class ClassInfoClassInstance implements ClassInfo {
    private final Class<?> clazz;

    public ClassInfoClassInstance(Class<?> clazz) {
      this.clazz = clazz;
    }

    @Override
    public int access() {
      return clazz.getModifiers();
    }

    @Override
    public String name() {
      return clazz.getName().replace('.', '/');
    }

    @Override
    public String superName() {
      return clazz == Object.class ? null : clazz.getSuperclass().getName().replace('.', '/');
    }

    public String[] interfaces() {
      Class<?>[] interfaces = clazz.getInterfaces();
      String[] result = new String[interfaces.length];
      for (int i = 0; i < interfaces.length; i++) {
        result[i] = interfaces[i].getName().replace('.', '/');
      }
      return result;
    }
  }
}