/**
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 *
 * Copyright (C)
 * <a href="http://www.inab.org">Spanish National Institute of Bioinformatics (INB)</a>
 * <a href="http://www.bsc.es">Barcelona Supercomputing Center (BSC)</a>
 * <a href="http://inb.bsc.es">Computational Node 6</a>
 */

package org.inb.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;
import javax.tools.*;
import javax.tools.JavaFileObject.Kind;

/**
 * In-memory file manager to dynamically compile generated java classes
 * 
 * @author Dmitry Repchevsky
 */

public final class MemoryFileManager extends ForwardingJavaFileManager<JavaFileManager>
{
    private Map<String, ByteArrayOutputStream> classes; // compiled classes
    private List<JavaClassInArray> list;

    public MemoryFileManager(JavaFileManager fileManager)
    {
        super(fileManager);

        classes = new ConcurrentHashMap<String, ByteArrayOutputStream>();
        list = new ArrayList<JavaClassInArray>();
    }

    public void addJarFile(byte[] jar) throws IOException
    {
        JarInputStream in = new JarInputStream(new ByteArrayInputStream(jar));

        JarEntry entry;
        while((entry = in.getNextJarEntry()) != null)
        {
            String name = entry.getName();

            if (name.endsWith(Kind.CLASS.extension))
            {
                name = name.substring(0, name.length() - Kind.CLASS.extension.length());

                ByteArrayOutputStream out = new ByteArrayOutputStream();

                byte[] buf = new byte[1024];

                int i;
                while (( i = in.read(buf)) >= 0)
                {
                    out.write(buf, 0, i);
                }

                list.add(new JavaClassInArray(name, out.toByteArray()));

            }
        }
    }

    public Map<String, byte[]> getCompiledClasses()
    {
        Map<String, byte[]> compiled = new TreeMap<String, byte[]>();

        for(Map.Entry<String, ByteArrayOutputStream> entry : classes.entrySet())
        {
            ByteArrayOutputStream stream = entry.getValue();

            if (stream != null)
            {
                compiled.put(entry.getKey(), stream.toByteArray());
            }
        }

        return compiled;
    }

    @Override
    public JavaFileObject getJavaFileForInput(JavaFileManager.Location location,
                                          String className,
                                          JavaFileObject.Kind kind)
                                   throws IOException
    {
        return super.getJavaFileForInput(location, className, kind);
    }

    @Override
    public Iterable<JavaFileObject> list(JavaFileManager.Location location,
                                     String packageName,
                                     Set<JavaFileObject.Kind> kinds,
                                     boolean recurse)
                              throws IOException
    {
        List<JavaFileObject> result = new ArrayList<JavaFileObject>();

        for(JavaFileObject f : super.list(location, packageName, kinds, recurse))
        {
            result.add(f);
        }

        for (JavaClassInArray file : list)
        {
            String pkg = file.getPackage();

            if (pkg.startsWith(packageName))
            {
                result.add(file);
            }
        }

        return result;

    }

    @Override
    public String inferBinaryName(JavaFileManager.Location location, JavaFileObject file)
    {
        if (file instanceof JavaClassInArray)
        {
            return file.getName();
        }
        else
        {
            return super.inferBinaryName(location, file);
        }
    }

    @Override
    public JavaFileObject getJavaFileForOutput(Location location,
					       String name,
					       Kind kind,
					       FileObject originatingSource)
        throws UnsupportedOperationException
    {
        return new JavaClassInArray(name);
    }

    private class JavaClassInArray extends SimpleJavaFileObject
    {
        private String pkg;
        private String name;

        private byte[] clazz;

        JavaClassInArray(String name)
        {
            super(URI.create("file:///" + name.replace('.','/') + Kind.CLASS.extension), Kind.CLASS);
            this.name = name.replace('/', '.');
        }

        JavaClassInArray(String name, byte[] clazz)
        {
            super(URI.create("file:///" + name.replace('.','/') + Kind.CLASS.extension), Kind.CLASS);

            this.name = name.replace('/', '.');
            this.clazz = clazz;
        }

        @Override
        public String getName()
        {
            return name;
        }

        public String getPackage()
        {
            if (pkg == null)
            {
                int idx = name.lastIndexOf('.');

                if (idx < 0)
                {
                    pkg = "";

                }
                else
                {
                    pkg = name.substring(0, idx);
                }
            }

            return pkg;
        }

        @Override
        public OutputStream openOutputStream()
        {
            ByteArrayOutputStream out = new ByteArrayOutputStream();

            classes.put(name, out);

            return out;
        }

        @Override
        public InputStream openInputStream()
        {
            return new ByteArrayInputStream(clazz);
        }
    }

    public static class JavaSourceFromString extends SimpleJavaFileObject
    {
        final String code;

        public JavaSourceFromString(String name, String code)
        {
            super(URI.create("file:///" + name.replace('.', '/') + Kind.SOURCE.extension), Kind.SOURCE);
            this.code = code;
        }

        @Override
        public CharSequence getCharContent(boolean ignoreEncodingErrors)
        {
            return code;
        }
    }
}