/*
 * Copyright (c) 2012, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */
package com.sun.tools.jdeps;

import com.sun.tools.classfile.ClassFile;
import com.sun.tools.classfile.ConstantPoolException;
import com.sun.tools.classfile.Dependencies.ClassFileError;
import java.io.*;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * ClassFileReader reads ClassFile(s) of a given path that can be
 * a .class file, a directory, or a JAR file.
 */
public class ClassFileReader implements Iterator<ClassFile> {
    /**
     * Returns a ClassFileReader instance of a given path.
     */
    public static ClassFileReader newInstance(File path) throws IOException {
        if (!path.exists()) {
            throw new FileNotFoundException(path.getAbsolutePath());
        }

        if (path.isDirectory()) {
            return new DirectoryArchive(path.toPath());
        } else if (path.getName().endsWith(".jar")) {
            return new JarFileArchive(path.toPath());
        } else {
            return new ClassFileReader(path.toPath());
        }
    }

    protected final Path path;
    protected int classCount;
    private ClassFileReader(Path path) throws IOException {
        this.path = path;
        if (path.toFile().getName().endsWith(".class")) {
            this.classCount = 0;
        } else {
            this.classCount = -1;
        }
    }

    public String getFileName() {
        return path.toFile().getName();
    }

    public ClassFile getClassFile(String classname) {
        String fn = classname.replace('.', File.separatorChar) + ".class";
        if (!fn.equals(getFileName())) {
            return null;
        }

        try (InputStream is = Files.newInputStream(path)) {
            return ClassFile.read(is);
        } catch (IOException | ConstantPoolException e) {
            throw new ClassFileError(e);
        }
    }

    public Iterable<ClassFile> getClassFiles() {
        final Iterator<ClassFile> iter = this;
        return new Iterable<ClassFile>() {
            public Iterator<ClassFile> iterator() {
                return iter;
            }
        };
    }

    @Override
    public boolean hasNext() {
        return classCount == 0;
    }

    @Override
    public ClassFile next() {
        if (!hasNext()) {
            throw new NoSuchElementException();
        }

        try (InputStream is = Files.newInputStream(path)) {
            classCount++;
            return ClassFile.read(is);
        } catch (IOException | ConstantPoolException e) {
            throw new ClassFileError(e);
        }
    }

    @Override
    public void remove() {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    @Override
    public String toString() {
        return path.toString();
    }

    private static class DirectoryArchive extends ClassFileReader {
        final List<Path> classes;
        private int index;

        DirectoryArchive(Path path) throws IOException {
            super(path);
            classes = walkTree(path);
            index = 0;
        }

        @Override
        public ClassFile getClassFile(String classname) {
            String fn = classname.replace('.', '/') + ".class";
            Path p = path.resolve(fn);
            if (!p.toFile().exists()) {
                return null;
            }
            try (InputStream is = Files.newInputStream(p)) {
                return ClassFile.read(is);
            } catch (IOException | ConstantPoolException e) {
                throw new ClassFileError(e);
            }
        }

        List<Path> walkTree(Path dir) throws IOException {
            final List<Path> files = new ArrayList<>();

            Files.walkFileTree(dir, new SimpleFileVisitor<Path>() {
                public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
                        throws IOException {
                    if (file.toFile().getName().endsWith(".class")) {
                        files.add(file);
                    }
                    return FileVisitResult.CONTINUE;
                }

                public FileVisitResult postVisitDirectory(Path dir, IOException e)
                        throws IOException {
                    if (e == null) {
                        return FileVisitResult.CONTINUE;
                    } else {
                        // directory iteration failed
                        throw e;
                    }
                }
            });
            return files;
        }

        @Override
        public boolean hasNext() {
            return index != classes.size();
        }

        @Override
        public ClassFile next() {
            if (!hasNext()) {
                throw new NoSuchElementException();
            }
            Path path = classes.get(index++);
            try (InputStream is = Files.newInputStream(path)) {
                classCount++;
                return ClassFile.read(is);
            } catch (IOException | ConstantPoolException e) {
                throw new ClassFileError(e);
            }
        }
    }

    private static class JarFileArchive extends ClassFileReader {
        final JarFile jarfile;
        final Enumeration<JarEntry> entries;
        JarEntry nextEntry;

        JarFileArchive(Path path) throws IOException {
            super(path);
            this.jarfile = new JarFile(path.toFile());
            this.entries = jarfile.entries();
            while (entries.hasMoreElements()) {
                JarEntry e = entries.nextElement();
                String name = e.getName();
                if (name.endsWith(".class")) {
                    nextEntry = e;
                    break;
                }
            }
        }

        @Override
        public ClassFile getClassFile(String classname) {
            String fn = classname.replace('.', '/') + ".class";
            JarEntry e = jarfile.getJarEntry(fn);
            if (e == null) {
                return null;
            }

            try (InputStream is = jarfile.getInputStream(e)) {
                classCount++;
                return ClassFile.read(is);
            } catch (IOException | ConstantPoolException ex) {
                throw new ClassFileError(ex);
            }
        }

        @Override
        public boolean hasNext() {
            return nextEntry != null;
        }

        @Override
        public ClassFile next() {
            if (!hasNext()) {
                throw new NoSuchElementException();
            }

            ClassFile cf;
            try {
                cf = ClassFile.read(jarfile.getInputStream(nextEntry));
            } catch (IOException | ConstantPoolException ex) {
                throw new ClassFileError(ex);
            }
            JarEntry entry = nextEntry;
            nextEntry = null;
            while (entries.hasMoreElements()) {
                JarEntry e = entries.nextElement();
                String name = e.getName();
                if (name.endsWith(".class")) {
                    nextEntry = e;
                    break;
                }
            }
            return cf;
        }
    }
}