package org.glavo.javah;

import org.objectweb.asm.*;

import java.io.*;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

public final class HeaderGenerator {
    private static String escape(String source) {
        StringBuilder builder = new StringBuilder();
        char ch;
        for (int i = 0; i < source.length(); i++) {
            switch (ch = source.charAt(i)) {
                case '_':
                    builder.append("_1");
                    break;
                case ';':
                    builder.append("_2");
                    break;
                case '[':
                    builder.append("_3");
                    break;
                case '/':
                    builder.append('.');
                    break;
                default:
                    if ((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')) {
                        builder.append(ch);
                    } else {
                        builder.append("_0").append(String.format("%04x", (int) ch));
                    }
            }
        }
        return builder.toString();
    }

    private static String typeToNative(Type tpe) {
        if (tpe == Type.BOOLEAN_TYPE) {
            return "jboolean";
        } else if (tpe == Type.BYTE_TYPE) {
            return "jbyte";
        } else if (tpe == Type.CHAR_TYPE) {
            return "jchar";
        } else if (tpe == Type.SHORT_TYPE) {
            return "jshort";
        } else if (tpe == Type.INT_TYPE) {
            return "jint";
        } else if (tpe == Type.LONG_TYPE) {
            return "jlong";
        } else if (tpe == Type.FLOAT_TYPE) {
            return "jfloat";
        } else if (tpe == Type.DOUBLE_TYPE) {
            return "jdouble";
        } else if (tpe == Type.VOID_TYPE) {
            return "void";
        } else {
            String str = tpe.toString();
            if (str.startsWith("[")) {
                str = typeToNative(tpe.getElementType());
                if ("jclass".equals(str)) {
                    return "jobjectArray";
                } else {
                    return str + "Array";
                }
            }
            return "jobject"; //TODO: jthrowable
        }
    }

    public static void generateHeader(ClassReader reader, PrintWriter output) {
        output.println("/* DO NOT EDIT THIS FILE - it is machine generated */");
        output.println("#include <jni.h>");

        Generator generator = new Generator();
        reader.accept(generator, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);

        StringBuilder builder = new StringBuilder();
        String className = escape(generator.getClassName());

        output.println("/* Header for class " + className + " */");

        String includeHeader = "_Include_" + className;
        output.println("#ifndef " + includeHeader);
        output.println("#define " + includeHeader);

        output.println("#ifdef __cplusplus\n" +
                "extern \"C\" {\n" +
                "#endif");

        for (Map.Entry<String, Set<MethodDesc>> entry : generator.getMethods().entrySet()) {
            boolean overload = entry.getValue().size() > 1;
            for (MethodDesc desc : entry.getValue()) {
                String methodName = escape(entry.getKey());
                output.println("/*" + "\n" +
                        " * Class:     " + className + "\n" +
                        " * Method:    " + entry.getKey() + "\n" +
                        " * Signature: " + desc.descriptor + "\n" +
                        " */"
                );

                Type[] argTypes = Type.getArgumentTypes(desc.descriptor);
                Type retType = Type.getReturnType(desc.descriptor);

                output.print(
                        "JNIEXPORT " + typeToNative(retType) + " JNICALL Java_" + className + "_" + methodName
                );

                if (overload) {
                    output.print("__");
                    for (Type tpe : argTypes) {
                        output.print(escape(tpe.toString()));
                    }
                }
                output.println();

                output.print("  (JNIEnv *, ");
                if (desc.isStatic) {
                    output.print("jclass");
                } else {
                    output.print("jobject");
                }
                for (Type tpe : argTypes) {
                    output.print(", " + typeToNative(tpe));
                }
                output.println(");\n");
            }

        }

        output.println("#ifdef __cplusplus\n" +
                "}\n" +
                "#endif\n" +
                "#endif\n");

    }

    public static void generateHeader(byte[] classFile, PrintWriter output) {
        generateHeader(new ClassReader(classFile), output);
    }

    public static void generateHeader(byte[] classFileBuffer, int classFileOffset, int classFileLength, PrintWriter output) {
        generateHeader(new ClassReader(classFileBuffer, classFileOffset, classFileLength), output);
    }

    public static void generateHeader(String className, PrintWriter output) throws IOException {
        generateHeader(new ClassReader(className), output);
    }

    public static void generateHeader(InputStream input, PrintWriter output) throws IOException {
        generateHeader(new ClassReader(input), output);
    }

    public static void main(String[] args) throws IOException {
        File f = new File(args[0]);
        PrintWriter writer = new PrintWriter(System.out);
        generateHeader(new FileInputStream(f), writer);
        writer.flush();
    }
}

class MethodDesc {
    public final boolean isStatic;
    public final String descriptor;

    public MethodDesc(boolean isStatic, String descriptor) {
        this.isStatic = isStatic;
        this.descriptor = descriptor;
    }
}

class Generator extends ClassVisitor {
    private String className;

    private Map<String, Set<MethodDesc>> methods = new LinkedHashMap<String, Set<MethodDesc>>();

    public Generator() {
        super(Opcodes.ASM5);
    }

    @Override
    public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
        className = name;
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
        if ((access & Opcodes.ACC_NATIVE) != 0) {
            if (methods.containsKey(name)) {
                methods.get(name).add(new MethodDesc((access & Opcodes.ACC_STATIC) != 0, descriptor));
            } else {
                LinkedHashSet<MethodDesc> set = new LinkedHashSet<MethodDesc>();
                set.add(new MethodDesc((access & Opcodes.ACC_STATIC) != 0, descriptor));
                methods.put(name, set);
            }
        }
        return null;
    }

    public String getClassName() {
        return className;
    }

    public Map<String, Set<MethodDesc>> getMethods() {
        return methods;
    }
}