/*
 * Copyright (c) 2002-2008 LWJGL Project
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 *
 * * Neither the name of 'LWJGL' nor the names of
 *   its contributors may be used to endorse or promote products derived
 *   from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package org.lwjgl.util.generator;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.Reader;
import java.io.Writer;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.annotation.Annotation;
import java.nio.Buffer;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import javax.annotation.processing.Messager;
import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.ElementKindVisitor6;
import javax.tools.Diagnostic;
import javax.tools.Diagnostic.Kind;
import javax.tools.FileObject;
import javax.tools.StandardLocation;

import static org.lwjgl.util.generator.Utils.await;
import static org.lwjgl.util.generator.Utils.spawn;

/**
 * Generator visitor for the generator tool
 *
 * @author elias_naur <elias_naur@users.sourceforge.net>
 * @version $Revision$ $Id$
 */
public final class GeneratorVisitor extends ElementKindVisitor6<Void, Void> {
	private static final String FAST_PREFIX = "";
	private static final String SLOW_PREFIX = "\u001b[1;31m";
	private static final String TIMING_SUFFIX = "\u001b[0m";

	private static void printTiming(Messager messager, String message, long startTime) {
		long elapsed = System.currentTimeMillis() - startTime;
		if (elapsed > 5) {
			message += SLOW_PREFIX;
		} else {
			message += FAST_PREFIX;
		}
		messager.printMessage(Diagnostic.Kind.NOTE, message + " in " + elapsed + " ms" + TIMING_SUFFIX + '.');
	}

	private final ProcessingEnvironment		env;
	private final Path						gen_java_path;
	private final Path						gen_native_path;
	private final TypeMap					type_map;
	private final boolean					generate_error_checks;
	private final boolean					context_specific;
	private final boolean					validate;
	public	final ArrayDeque<Future<Void>>	futures = new ArrayDeque<>(128);

	public GeneratorVisitor(ProcessingEnvironment env, Path gen_java_path, Path gen_native_path, TypeMap type_map, boolean generate_error_checks, boolean context_specific, boolean validate) {
		this.env = env;
		this.gen_java_path = gen_java_path;
		this.gen_native_path = gen_native_path;
		this.type_map = type_map;
		this.generate_error_checks = generate_error_checks;
		this.context_specific = context_specific;
		this.validate = validate;
	}

	private void validateMethod(ExecutableElement method) {
		if ( method.isVarArgs() ) {
			throw new RuntimeException("Method " + method.getSimpleName() + " is variadic");
		}
		Collection<Modifier> modifiers = method.getModifiers();
		if ( !modifiers.contains(Modifier.PUBLIC) ) {
			throw new RuntimeException("Method " + method.getSimpleName() + " is not public");
		}
		if ( method.getThrownTypes().size() > 0 ) {
			throw new RuntimeException("Method " + method.getSimpleName() + " throws checked exceptions");
		}
		validateParameters(method);
		StripPostfix strip_annotation = method.getAnnotation(StripPostfix.class);
		if ( strip_annotation != null && method.getAnnotation(Alternate.class) == null ) {
			String postfix_param_name = strip_annotation.value();
			VariableElement postfix_param = Utils.findParameter(method, postfix_param_name);
			if ( Utils.isParameterMultiTyped(postfix_param) ) {
				throw new RuntimeException("Postfix parameter can't be the same as a multityped parameter in method " + method);
			}
			if ( Utils.getNIOBufferType(postfix_param.asType()) == null ) {
				throw new RuntimeException("Postfix parameter type must be a nio Buffer");
			}
		}
		if ( Utils.getResultParameter(method) != null && !method.getReturnType().equals(env.getTypeUtils().getNoType(TypeKind.VOID)) ) {
			throw new RuntimeException(method + " return type is not void but a parameter is annotated with Result");
		}
		if ( method.getAnnotation(CachedResult.class) != null ) {
			if ( Utils.getNIOBufferType(Utils.getMethodReturnType(method)) == null ) {
				throw new RuntimeException(method + " return type is not a Buffer, but is annotated with CachedResult");
			}
			if ( method.getAnnotation(AutoSize.class) == null ) {
				throw new RuntimeException(method + " is annotated with CachedResult but misses an AutoSize annotation");
			}
		}
		validateTypes(method, method.getAnnotationMirrors(), method.getReturnType());
	}

	private void validateType(ExecutableElement method, Class<? extends Annotation> annotation_type, Class type) {
		Class[] valid_types = type_map.getValidAnnotationTypes(type);
		for ( Class valid_type : valid_types ) {
			if ( valid_type.getName().equals(annotation_type) ) {
				return;
			}
		}
		throw new RuntimeException(type + " is annotated with invalid native type " + annotation_type
		                           + " in method " + method);
	}

	private void validateType(ExecutableElement method, Class<? extends Annotation> annotation_type, String type) {
		// TODO: reimplement this if we even care
	}

	private void validateTypes(ExecutableElement method, List<? extends AnnotationMirror> annotations, TypeMirror type_mirror) {
		if (true) {
			// TODO: remove if we reimplement validateType.
			return;
		}
		String desc = Utils.getDescriptor(type_mirror);
		if ( BUFFER_DESC.equals(desc) ) return;
		for ( AnnotationMirror annotation : annotations ) {
			NativeType native_type_annotation = NativeTypeTranslator.getAnnotation(annotation, NativeType.class);
			if ( native_type_annotation != null ) {
				Class<? extends Annotation> annotation_type = NativeTypeTranslator.getClassFromType(annotation.getAnnotationType());
				validateType(method, annotation_type, desc);
			}
		}
	}

	private static final String BUFFER_DESC = DescriptorTypeTranslator.forClass(Buffer.class);
	private static final String CHAR_SEQUENCE_DESC = DescriptorTypeTranslator.forClass(CharSequence.class);
	private static final String CHAR_SEQUENCE_ARRAY_DESC = DescriptorTypeTranslator.forArray(CHAR_SEQUENCE_DESC);

	private void validateParameters(ExecutableElement method) {
		method.getParameters().parallelStream().parallel().forEach(param -> {
			validateTypes(method, param.getAnnotationMirrors(), param.asType());
			String param_type = Utils.getDescriptor(param.asType());
			if (Utils.getNIOBufferType(param.asType()) != null && !CHAR_SEQUENCE_DESC.equals(param_type) && !CHAR_SEQUENCE_ARRAY_DESC.equals(param_type)) {
				Check parameter_check_annotation = param.getAnnotation(Check.class);
				NullTerminated null_terminated_annotation = param.getAnnotation(NullTerminated.class);
				if (parameter_check_annotation == null && null_terminated_annotation == null) {
					boolean found_auto_size_param = false;
					for (VariableElement inner_param : method.getParameters()) {
						AutoSize auto_size_annotation = inner_param.getAnnotation(AutoSize.class);
						if (auto_size_annotation != null && auto_size_annotation.value().equals(param.getSimpleName().toString())) {
							found_auto_size_param = true;
							break;
						}
					}
					if (!found_auto_size_param
					     && param.getAnnotation(Result.class) == null
					     && param.getAnnotation(Constant.class) == null
					     && !Utils.isReturnParameter(method, param)) {
						throw new RuntimeException(param + " has no Check, Result nor Constant annotation, is not the return parameter and no other parameter has"
						                           + " an @AutoSize annotation on it in method " + method);
					}
				}
				if ( param.getAnnotation(CachedReference.class) != null && param.getAnnotation(Result.class) != null ) {
					throw new RuntimeException(param + " can't be annotated with both CachedReference and Result");
				}
				if ( param.getAnnotation(BufferObject.class) != null && param.getAnnotation(Result.class) != null ) {
					throw new RuntimeException(param + " can't be annotated with both BufferObject and Result");
				}
				//if (param.getAnnotation(Constant.class) != null)
				//throw new RuntimeException("Buffer parameter " + param + " cannot be Constant");
			} else {
				if ( param.getAnnotation(BufferObject.class) != null ) {
					throw new RuntimeException(param + " type is not a buffer, but annotated as a BufferObject");
				}
				if ( param.getAnnotation(CachedReference.class) != null ) {
					throw new RuntimeException(param + " type is not a buffer, but annotated as a CachedReference");
				}
			}
		});
	}

	private static void generateMethodsNativePointers(PrintWriter writer, Collection<ExecutableElement> methods) {
		for ( ExecutableElement method : methods ) {
			if ( method.getAnnotation(Alternate.class) == null ) {
				generateMethodNativePointers(writer, method);
			}
		}
	}

	private static void generateMethodNativePointers(PrintWriter writer, ExecutableElement method) {
		if ( method.getAnnotation(Extern.class) == null ) {
			writer.append("static ");
		}
		writer.append(Utils.getTypedefName(method)).append(' ').append(method.getSimpleName()).append(";\n");
	}

	private void generateJavaSource(TypeElement d, PrintWriter java_writer, Collection<ExecutableElement> methods, Collection<VariableElement> fields) throws IOException {
		java_writer.append("/* MACHINE GENERATED FILE, DO NOT EDIT */\n\n");
		java_writer.append("package ").append(env.getElementUtils().getPackageOf(d).getQualifiedName().toString()).append(";\n\n");
		java_writer.append("import org.lwjgl.*;\n"
				+ "import java.nio.*;\n");
		Imports imports = d.getAnnotation(Imports.class);
		if ( imports != null ) {
			for ( String i : imports.value() ) {
				java_writer.append("import ").append(i).append(";\n");
			}
		}
		java_writer.append('\n');
		Utils.printDocComment(java_writer, d, env);
		if ( d.getAnnotation(Private.class) == null ) {
			java_writer.append("public ");
		}
		boolean is_final = Utils.isFinal(d);
		if ( is_final ) {
			java_writer.append("final ");
		}
		java_writer.append("class ").append(Utils.getSimpleClassName(d));
		List<? extends TypeMirror> super_interfaces = d.getInterfaces();
		if ( super_interfaces.size() > 1 ) {
			throw new RuntimeException(d + " extends more than one interface");
		}
		if ( super_interfaces.size() == 1 ) {
			TypeMirror super_interface = super_interfaces.iterator().next();
			java_writer.append(" extends ").append(Utils.getSimpleClassName(env.getElementUtils().getTypeElement(super_interface.toString())));
		}
		java_writer.append(" {\n");
		FieldsGenerator.generateFields(env, java_writer, fields);
		java_writer.append('\n');
		if ( is_final ) {
			// Write private constructor to avoid instantiation
			java_writer.append("\tprivate ").append(Utils.getSimpleClassName(d)).append("() {}\n");
		}
		if ( !methods.isEmpty() && !context_specific ) {
			java_writer.append('\n');
			java_writer.append("\tstatic native void ").append(Utils.STUB_INITIALIZER_NAME).append("() throws LWJGLException;\n");
		}
		JavaMethodsGenerator.generateMethodsJava(env, type_map, java_writer, d, generate_error_checks, context_specific);
		java_writer.append("}");
	}

	private void generateNativeSource(TypeElement d, Collection<ExecutableElement> methods, long startTime) throws IOException {
		if ( d.getKind().equals(ElementKind.ANNOTATION_TYPE) ) {
			return;
		}
		String qualified_interface_name = Utils.getQualifiedClassName(d);
		ByteArrayOutputStream native_writer1 = new ByteArrayOutputStream();
		PrintWriter native_writer = new PrintWriter(native_writer1);
    		
		native_writer.append("/* MACHINE GENERATED FILE, DO NOT EDIT */\n\n");
    	native_writer.append("#include <jni.h>\n");
    	type_map.printNativeIncludes(native_writer);
    	native_writer.append('\n');
		TypedefsGenerator.generateNativeTypedefs(type_map, native_writer, methods);
   		native_writer.append('\n');
   		if (!context_specific) {
   			generateMethodsNativePointers(native_writer, methods);
   			native_writer.append('\n');
   		}
   		NativeMethodStubsGenerator.generateNativeMethodStubs(env, type_map, native_writer, d, generate_error_checks, context_specific);
   		if (!context_specific) {
   			native_writer.append("JNIEXPORT void JNICALL ").append(Utils.getQualifiedNativeMethodName(qualified_interface_name, Utils.STUB_INITIALIZER_NAME));
   			native_writer.append("(JNIEnv *env, jclass clazz) {\n");
   			native_writer.append("\tJavaMethodAndExtFunction functions[] = {\n");
   			RegisterStubsGenerator.generateMethodsNativeStubBind(native_writer, d, generate_error_checks, context_specific);
   			native_writer.append("\t};\n");
   			native_writer.append("\tint num_functions = NUMFUNCTIONS(functions);\n");
   			native_writer.append('\t');
   			native_writer.append(type_map.getRegisterNativesFunctionName());
   			native_writer.append("(env, clazz, num_functions, functions);\n");
   			native_writer.append("}");
	   	}
		native_writer.flush();
		saveGeneratedCSource(env.getMessager(), this.gen_native_path, d, spawn((Callable<byte[]>) () -> native_writer1.toByteArray()), startTime);
	}

	private void doJavaGen(TypeElement e, Collection<ExecutableElement> methods, Collection<VariableElement> fields) {
		long startTime = System.currentTimeMillis();
		ByteArrayOutputStream java_writer = new ByteArrayOutputStream();
		PrintWriter pw = new PrintWriter(java_writer);
		try {
			generateJavaSource(e, pw, methods, fields);
		} catch (IOException ex) {
			throw new RuntimeException("Failed to generate the Java sources for " + e, ex);
		}
		pw.flush();
		saveGeneratedJavaSource(env.getMessager(), this.gen_java_path, e, spawn((Callable<byte[]>) () -> java_writer.toByteArray()), startTime);
	}

	private void doNativeGen(TypeElement e, Collection<ExecutableElement> methods) {
		long startTime = System.currentTimeMillis();
		boolean noNative = true;
		for (ExecutableElement method : methods) {
			Alternate alt_annotation = method.getAnnotation(Alternate.class);
			if ((alt_annotation == null || alt_annotation.nativeAlt()) && method.getAnnotation(Reuse.class) == null) {
				noNative = false;
				break;
			}
		}
		if (!noNative) {
			try {
				generateNativeSource(e, methods, startTime);
			} catch (IOException ex) {
				throw new RuntimeException(ex);
			}
		}
	}

	@Override
	public Void visitTypeAsInterface(TypeElement e, Void p) {
		final Collection<ExecutableElement> methods = List.copyOf(Utils.getMethods(e));
		final Collection<VariableElement> fields = List.copyOf(Utils.getFields(e));
		if (methods.isEmpty() && fields.isEmpty()) {
			return DEFAULT_VALUE;
		}

		//env.getMessager().printMessage(Kind.NOTE, "methods count : " + Utils.getMethods(e).size() + " fields count : " + Utils.getFields(e).size(), e);
		if (this.validate) {
			this.futures.push(spawn(() -> {
				long startTime = System.currentTimeMillis();
				methods.parallelStream().parallel().forEach(method -> validateMethod(method));
				printTiming(env.getMessager(), "Validated " + e, startTime);
				return null;
			}));
		}

		if (this.gen_java_path != null) {
			//this.futures.push(spawn(() -> {
				doJavaGen(e, methods, fields);
			//	return null;
			//}));
		}

		if (methods.size() > 0 && this.gen_native_path != null) {
			this.futures.push(spawn(() -> {
				doNativeGen(e, methods);
				return null;
			}));
		}

		return DEFAULT_VALUE;
	}

	public static void saveGeneratedJavaSource(Messager messager, Path gen_java_path, TypeElement e, Future<byte[]> content, long startTime) {
		String qualified_interface_name = Utils.getQualifiedClassName(e);
		saveGeneratedJavaSource(messager, gen_java_path, qualified_interface_name, content, startTime);
	}

	public static void saveGeneratedJavaSource(Messager messager, Path gen_java_path, String qualified_interface_name, Future<byte[]> content, long startTime) {
		final Path output = gen_java_path.resolve(qualified_interface_name.replace('.', '/') + ".java");
		saveGeneratedSource(messager, "java", qualified_interface_name, output, content, startTime);
	}

	public static void saveGeneratedCSource(Messager messager, Path gen_native_path, TypeElement e, Future<byte[]> content, long startTime) {
		String qualified_interface_name = Utils.getQualifiedClassName(e);
		final Path output = gen_native_path.resolve(Utils.getNativeQualifiedName(qualified_interface_name) + ".c");
		saveGeneratedSource(messager, "c", qualified_interface_name, output, content, startTime);
	}

	public static void saveGeneratedSource(Messager messager, String type, String name, Path output, Future<byte[]> newStr, long startTime) {
		try {
			if (Files.exists(output)) {
				Future<byte[]> actual = spawn(() -> Files.readAllBytes(output));
				byte[] expected = await(newStr);

				if (Arrays.equals(expected, await(actual))) {
					printTiming(messager, "Skipped identical '." + type + "' " + name + " at " + output, startTime);
					return;
				}
			} else {
				Files.createDirectories(output.getParent());
			}
		} catch (IOException e) {
			throw new RuntimeException("Failed to create the output file for " + name, e);
		}
		try {
			Files.write(output, await(newStr));
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
		printTiming(messager, "Generated '." + type + "' " + name + " at " + output, startTime);
	}
}