/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://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.
 */

package org.apache.ignite.ml.inference.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.HashMap;
import java.util.Map;

/**
 * Utils class that helps to serialize directory content as a has map and then deserialize it.
 */
public class DirectorySerializer {
    /**
     * Serializes directory content.
     *
     * @param path Path to the directory.
     * @return Serialized directory content.
     * @throws IOException If directory cannot be serialized.
     */
    public static byte[] serialize(Path path) throws IOException {
        File file = path.toFile();

        if (!file.isDirectory())
            throw new IllegalStateException("Path is not directory [path=\"" + path + "\"]");

        Map<String, byte[]> data = new HashMap<>();
        serialize(data, path, file);

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ObjectOutputStream oos = new ObjectOutputStream(baos);
        oos.writeObject(data);
        oos.flush();

        return baos.toByteArray();
    }

    /**
     * Deserializes directory content.
     *
     * @param path Path to the directory.
     * @param data Serialized content.
     * @throws IOException If the directory cannot be deserialized.
     * @throws ClassNotFoundException If the directory cannot be deserialized.
     */
    @SuppressWarnings("unchecked")
    public static void deserialize(Path path, byte[] data) throws IOException, ClassNotFoundException {
        ByteArrayInputStream bais = new ByteArrayInputStream(data);
        ObjectInputStream ois = new ObjectInputStream(bais);
        Map<String, byte[]> files = (Map<String, byte[]>)ois.readObject();

        for (Map.Entry<String, byte[]> file : files.entrySet()) {
            Path dst = path.resolve(file.getKey());
            File dstFile = dst.toFile();
            Files.createDirectories(dstFile.getParentFile().toPath());
            Files.createFile(dstFile.toPath());
            FileOutputStream fos = new FileOutputStream(dstFile);
            fos.write(file.getValue());
            fos.flush();
        }
    }

    /**
     * Removes the specified directory.
     *
     * @param path Path to the directory.
     */
    public static void deleteDirectory(Path path) {
        File file = path.toFile();
        if (file.isDirectory()) {
            File[] children = file.listFiles();
            if (children != null) {
                for (File child : children)
                    deleteDirectory(child.toPath());
            }
        }

        try {
            Files.delete(path);
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Serializes directory content or file.
     *
     * @param data Storage to keep pairs of file name and file content.
     * @param basePath Base path to the serialized directory.
     * @param file File to be serialized.
     * @throws IOException If the file cannot be serialized.
     */
    private static void serialize(Map<String, byte[]> data, Path basePath, File file) throws IOException {
        if (file.isFile()) {
            String relative = basePath.relativize(file.toPath()).toString();
            byte[] bytes = Files.readAllBytes(file.toPath());
            data.put(relative, bytes);
        }
        else {
            File[] children = file.listFiles();
            if (children != null) {
                for (File child : children)
                    serialize(data, basePath, child);
            }
        }
    }
}

@@ -1841,32 +1841,37 @@
 file);%0A%0A        
+try (
 ByteArrayOutputS
@@ -1916,24 +1916,29 @@
 ();%0A        
+    
+ 
 ObjectOutput
@@ -1978,18 +1978,24 @@
 am(baos)
-;%0A
+) %7B%0A    
         
@@ -2021,24 +2021,28 @@
 a);%0A        
+    
 oos.flush();
@@ -2039,24 +2039,28 @@
 s.flush();%0A%0A
+    
         retu
@@ -2078,24 +2078,34 @@
 yteArray();%0A
+        %7D%0A
     %7D%0A%0A    /
@@ -2536,16 +2536,21 @@
         
+try (
 ByteArra
@@ -2605,24 +2605,29 @@
 a);%0A        
+     
 ObjectInputS
@@ -2665,18 +2665,24 @@
 am(bais)
-;%0A
+) %7B%0A    
         
@@ -2753,24 +2753,28 @@
 );%0A%0A        
+    
 for (Map.Ent
@@ -2831,16 +2831,20 @@
         
+   
+ 
 Path dst
@@ -2883,24 +2883,28 @@
             
+    
 File dstFile
@@ -2924,32 +2924,36 @@
 ();%0A            
+  
+  
 Files.createDire
@@ -2987,32 +2987,36 @@
 le().toPath());%0A
+    
             File
@@ -3055,24 +3055,33 @@
             
+    try (
 FileOutputSt
@@ -3120,18 +3120,28 @@
 dstFile)
-;%0A
+) %7B%0A        
         
@@ -3180,24 +3180,32 @@
             
+        
 fos.flush();
@@ -3201,24 +3201,56 @@
 os.flush();%0A
+                %7D%0A            %7D%0A
         %7D%0A  

Add try-with-resources to close streams in DirectorySerializer.