import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.io.UncheckedIOException;
import java.lang.module.ModuleDescriptor;
import java.lang.module.ModuleFinder;
import java.lang.module.ModuleReference;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HexFormat;
import java.util.List;
import java.util.Map;
import java.util.Scanner;
import java.util.Set;
import java.util.function.Function;
import java.util.spi.ToolProvider;
import java.util.stream.Collectors;

public class BuildShaSumFile {

    // Mirrored from ResourcePoolEntry.Type:
    private enum Type {
        CLASS_OR_RESOURCE(0),
        CONFIG(1),
        HEADER_FILE(2),
        LEGAL_NOTICE(3),
        MAN_PAGE(4),
        NATIVE_CMD(5),
        NATIVE_LIB(6),
        TOP(7);
        
        int intVal;
        
        Type(int val) {
            this.intVal = val;
        }
        
        Integer getOrdinal() {
            return intVal;
        }
    }
    
    private static final String RESOURCE_FILE_FORMAT = "jdk/tools/jlink/internal/runtimelink/fs_%s_files";
    private static final Map<Integer, Type> TYPE_MAP = Arrays.stream(Type.values())
                .collect(Collectors.toMap(Type::getOrdinal, Function.identity()));
    private static final Module JLINK_MODULE = ToolProvider.findFirst("jlink").orElseThrow().getClass().getModule();
    private static final Set<Type> BINARIES_LIBS = Set.of(Type.NATIVE_CMD, Type.NATIVE_LIB);
    
    private final String outPutFile;
    
    public BuildShaSumFile(String output) {
        this.outPutFile = output;
    }
    
    /*
     * Use the fs_<module-name>_files resource part of jdk.jlink to
     * figure out the files part of a module from the JDK image.
     * 
     * Note that the file is of the following format:
     * <type-id>|<symlink>|<sha512sum>|<file-path>
     * 
     * Example:
     * 5|0|428e61a697dc05f585333888f0d9baaef7dad088c279f052e189e401471d0c94e99711f567c5d1e5d2e08259c3b320b9d239cbb610f7c347b2eff4eec8ce712c|bin/jmod
     */
    private List<String> getModuleFiles(String module, Set<Type> includeTypes) {
        List<String> moduleFiles = new ArrayList<>();
        String resourceFile = String.format(RESOURCE_FILE_FORMAT, module);
        try (InputStream in = JLINK_MODULE.getResourceAsStream(resourceFile);
             Scanner scanner = new Scanner(in)) {
            while (scanner.hasNextLine()) {
                String[] tokens = scanner.nextLine().split("\\|", 4);
                Integer typeInt = Integer.valueOf(tokens[0]);
                Type type = TYPE_MAP.get(typeInt);
                boolean isSymlink = Integer.valueOf(tokens[1]) == 1;
                if (includeTypes.contains(type) && !isSymlink) {
                    moduleFiles.add(tokens[3]);
                }
            }
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
        return moduleFiles;
    }
    
    private String sha512Sum(String file) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-512");
            try (InputStream is = Files.newInputStream(Path.of(System.getProperty("java.home")).resolve(Path.of(file)))) {
                byte[] buf = new byte[1024];
                int readBytes = -1;
                while ((readBytes = is.read(buf)) != -1) {
                    digest.update(buf, 0, readBytes);
                }
            }
            byte[] hashSum = digest.digest();
            return HexFormat.of().formatHex(hashSum);
        } catch (Exception e) {
            throw new RuntimeException("Failed to generate hash sum");
        }
    }
    
    public void produceHashSumFile() {
        // Sanity check JEP 493 enabled builds:
        String resourceFile = String.format(RESOURCE_FILE_FORMAT, "java.base");
        try (InputStream in = JLINK_MODULE.getResourceAsStream(resourceFile)) {
            if (in == null) {
                System.out.println("Not a JEP 493 enabled build. Aborting!");
                return;
            }
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
        try (PrintWriter pw = new PrintWriter(new FileWriter(Path.of(outPutFile).toFile()))) {
            ModuleFinder.ofSystem().findAll().stream()
                    .map(ModuleReference::descriptor)
                    .map(ModuleDescriptor::name).forEach(m -> {
                List<String> moduleFiles = getModuleFiles(m, BINARIES_LIBS);
                for (String file: moduleFiles) {
                    String shaSum = sha512Sum(file);
                    pw.println(String.format("%s|%s|%s", m, file, shaSum));
                }
            });
            
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
        System.out.println("File " + outPutFile + " written. " + 
            "You can feed that now to 'jlink's --sha-overrides option.");
    }

    public static void main(String[] args) {
        if (args.length != 1) {
            System.err.println("Usage: " + BuildShaSumFile.class.getSimpleName() + " <sha-sum-file>");
            System.exit(1);
        }
        BuildShaSumFile b = new BuildShaSumFile(args[0]);
        b.produceHashSumFile();
    }

}