/********************************************************************************
 * Copyright (c) 2022, 2024 Contributors to the Eclipse Foundation
 *
 * See the NOTICE file(s) distributed with this work for additional
 * information regarding copyright ownership.
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ********************************************************************************/
package org.eclipse.emt4j.test.common;

import javax.tools.*;
import java.io.*;
import java.net.URI;
import java.util.*;

public class CompileUtil {
    private static final boolean PRINT_SOURCE = true;

    public static void compile(JavaSource[] javaSources, File classesDir, List<String> options) throws IOException {
        Map<String, JavaByteObject> byteObjectMap = compileInMemory(javaSources, options);
        writeToFile(javaSources, byteObjectMap, classesDir);
    }

    private static void writeToFile(JavaSource[] javaSources, Map<String, JavaByteObject> byteObjectMap, File classesDir) throws IOException {
        Set<String> writtenSet = new HashSet<>();
        for (JavaSource javaSource : javaSources) {
            JavaByteObject jbo = byteObjectMap.get(javaSource.getClassName());
            byte[] classBytes = jbo.getBytes();
            if (classBytes.length == 0) {
                throw new RuntimeException("Class : " + javaSource.getClassName() + " compile failed!");
            }
            doWrite(classesDir, javaSource, jbo);
            writtenSet.add(javaSource.getClassName());
        }

        //inner classes
        for (Map.Entry<String, JavaByteObject> entry : byteObjectMap.entrySet()) {
            if (!writtenSet.contains(entry.getKey())) {
                doWrite(classesDir, entry.getKey(), entry.getValue());
            }
        }
    }

    private static void doWrite(File classesDir, JavaSource javaSource, JavaByteObject jbo) throws IOException {
        String destFile;
        if (javaSource.getPackageName() != null) {
            preparePackageDir(classesDir, javaSource.getPackageName());
            destFile = javaSource.getPackageName().replace('.', File.separatorChar) + File.separator + javaSource.getName() + ".class";
        } else {
            destFile = javaSource.getName() + ".class";
        }
        try (FileOutputStream fos = new FileOutputStream(new File(classesDir, destFile))) {
            fos.write(jbo.getBytes());
        }
    }

    private static void doWrite(File classesDir, String className, JavaByteObject jbo) throws IOException {
        String destFile;
        int idx = className.lastIndexOf('.');
        if (idx != -1) {
            String packageName = className.substring(0, idx);
            preparePackageDir(classesDir, className.substring(0, idx));
            destFile = packageName.replace('.', File.separatorChar) + File.separator + className.substring(idx + 1) + ".class";
        } else {
            destFile = className + ".class";
        }
        try (FileOutputStream fos = new FileOutputStream(new File(classesDir, destFile))) {
            fos.write(jbo.getBytes());
        }
    }

    private static void preparePackageDir(File classesDir, String packageName) {
        File f = new File(classesDir, packageName.replace('.', File.separatorChar));
        if (!f.exists() && !f.mkdirs()) {
            throw new RuntimeException("Prepare package dir :" + f + " failed!Classes Dir: " + classesDir + ",package name: " + packageName);
        }
    }

    private static Map<String, JavaByteObject> compileInMemory(JavaSource[] javaSources, List<String> options) throws IOException {
        JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
        DiagnosticCollector<JavaFileObject> diagnostics =
                new DiagnosticCollector<>();

        final Map<String, JavaByteObject> byteObjectMap = new HashMap<>();
        for (JavaSource javaSource : javaSources) {
            byteObjectMap.put(javaSource.getClassName(), new JavaByteObject(javaSource.getClassName()));
        }

        StandardJavaFileManager standardFileManager =
                compiler.getStandardFileManager(diagnostics, null, null);

        JavaFileManager fileManager = createFileManager(standardFileManager,
                byteObjectMap);

        if (options == null) {
            options = new ArrayList<>();
        }
        // disable annotation processors
        options.add("-proc:none");

        JavaCompiler.CompilationTask task = compiler.getTask(null,
                fileManager, diagnostics, options, null, getCompilationUnits(javaSources));

        if (!task.call()) {
            diagnostics.getDiagnostics().forEach(System.out::println);
        }
        fileManager.close();
        return byteObjectMap;
    }

    public static Iterable<? extends JavaFileObject> getCompilationUnits(JavaSource[] javaSources) {
        List<JavaFileObject> fileObjectList = new ArrayList<>();
        for (JavaSource src : javaSources) {
            fileObjectList.add(new JavaStringObject(src.getClassName(), src.buildJavaSourceCode()));
            if (PRINT_SOURCE) {
                System.out.println("================= " + src.getClassName() + "=================");
                System.out.println(src.buildJavaSourceCode());
            }
        }
        return fileObjectList;
    }


    private static JavaFileManager createFileManager(StandardJavaFileManager fileManager,
                                                     Map<String, JavaByteObject> byteObjectMap) {
        return new ForwardingJavaFileManager<StandardJavaFileManager>(fileManager) {
            @Override
            public JavaFileObject getJavaFileForOutput(Location location,
                                                       String className, JavaFileObject.Kind kind,
                                                       FileObject sibling) throws IOException {
                JavaFileObject javaFileObject = byteObjectMap.get(className);
                if (null == javaFileObject) {
                    javaFileObject = new JavaByteObject(className);
                    byteObjectMap.put(className, (JavaByteObject) javaFileObject);
                }
                return javaFileObject;
            }
        };
    }


    private static class JavaByteObject extends SimpleJavaFileObject {
        private ByteArrayOutputStream outputStream;

        protected JavaByteObject(String name) {
            super(URI.create("bytes:///" + name + name.replaceAll("\\.", "/")), Kind.CLASS);
            outputStream = new ByteArrayOutputStream();
        }

        @Override
        public OutputStream openOutputStream() throws IOException {
            return outputStream;
        }

        public byte[] getBytes() {
            return outputStream.toByteArray();
        }
    }

    public static class JavaStringObject extends SimpleJavaFileObject {
        private final String source;

        protected JavaStringObject(String name, String source) {
            super(URI.create("string:///" + name.replaceAll("\\.", "/") +
                    Kind.SOURCE.extension), Kind.SOURCE);
            this.source = source;
        }

        @Override
        public CharSequence getCharContent(boolean ignoreEncodingErrors)
                throws IOException {
            return source;
        }
    }
}
