package com.qmth.boot.tools.io;

import org.apache.commons.lang3.StringUtils;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.zip.ZipEntry;
import java.util.zip.ZipFile;

/**
 * ZIP格式文件读取工具
 */
public class ZipReader {

    private ZipFile zipFile;

    private Node root;

    /**
     * 从本地文件初始化，默认使用UTF8字符集
     *
     * @param file
     * @throws IOException
     */
    public ZipReader(File file) throws IOException {
        this(file, StandardCharsets.UTF_8);
    }

    /**
     * 从本地文件初始化，指定字符集
     *
     * @param file
     * @param charset
     * @throws IOException
     */
    public ZipReader(File file, Charset charset) throws IOException {
        if (file == null) {
            throw new IllegalArgumentException("file should not be null");
        }
        if (charset == null) {
            throw new IllegalArgumentException("charset should not be null");
        }
        this.root = new Node();
        this.zipFile = new ZipFile(file, charset);
        zipFile.stream().forEach(entry -> {
            String[] paths = entry.getName().split(IOUtils.ZIP_SEPERATOR);
            int size = paths.length;
            Node current = root;
            for (int i = 0; i < (size - 1); i++) {
                current = current.addChildren(paths[i], false);
            }
            current.addChildren(paths[size - 1], !entry.isDirectory());
        });
    }

    /**
     * 关闭IO流，清空目录索引，不再使用
     */
    public void close() {
        try {
            this.root = null;
            this.zipFile.close();
        } catch (Exception e) {
        }
    }

    /**
     * 读取指定路径下的文件/目录列表，若读取根目录则path可以不传
     *
     * @param isFile
     * @param path
     * @return
     */
    public List<String> list(boolean isFile, String... path) {
        Node parent = root;
        if (path != null) {
            for (String dir : path) {
                parent = parent.children.get(dir);
                if (parent == null) {
                    return Collections.EMPTY_LIST;
                }
            }
        }
        return parent.children.values().stream().filter(node -> isFile == node.isFile).map(node -> node.name)
                .collect(Collectors.toList());
    }

    /**
     * 按照指定路径读取ZIP内文件，获取输入流
     *
     * @param path
     * @return
     * @throws IOException
     */
    public InputStream read(String... path) throws IOException {
        ZipEntry entry = zipFile.getEntry(StringUtils.join(path, IOUtils.ZIP_SEPERATOR));
        if (entry != null) {
            return zipFile.getInputStream(entry);
        } else {
            return null;
        }
    }

    private class Node {

        boolean isFile;

        String name;

        Map<String, Node> children = new HashMap<>();

        private Node addChildren(String name, boolean isFile) {
            return children.computeIfAbsent(name, key -> {
                Node child = new Node();
                child.name = key;
                child.isFile = isFile;
                return child;
            });
        }
    }

}
