import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;

import Glacier2.CannotCreateSessionException;
import Glacier2.PermissionDeniedException;

import omero.ServerError;
import omero.client;
import omero.api.ExporterPrx;
import omero.api.IContainerPrx;
import omero.api.ServiceFactoryPrx;
import omero.model.Dataset;
import omero.model.IObject;
import omero.model.Image;
import omero.sys.ParametersI;


public class OmeroDatasetExport {

    private final omero.client client;
    private final ServiceFactoryPrx sf;
    private final IContainerPrx containerService;
    private static final int BUFFER_SIZE = 1024 * 1024;  // 1MB

    public OmeroDatasetExport(String host, int port, String username, String password)
            throws CannotCreateSessionException, PermissionDeniedException, ServerError {
        client = new client(host, port);
        sf = client.createSession(username, password);
        containerService = sf.getContainerService();
    }

    public void exportImages(long datasetId) throws ServerError, IOException {
        ParametersI param = new ParametersI();
        param.leaves();
        List<omero.model.IObject> results =
                containerService.loadContainerHierarchy(
                        Dataset.class.getName(), Arrays.asList(datasetId), param);
        if (results.size() == 0){
            throw new RuntimeException("No dataset with ID: " + datasetId);
        }

        for (IObject result : results) {
            Dataset dataset = (Dataset) result;
            for (Image image : dataset.linkedImageList()) {
                long imageId = image.getId().getValue();
                ExporterPrx exporter = sf.createExporter();
                try {
                    exporter.addImage(imageId);
                    exporter.generateTiff();

                    File file = new File(System.getProperty("user.dir"),
                                         imageId + ".ome.tif");
                    FileOutputStream fos = new FileOutputStream(file);
                    long offset = 0;
                    int readLength = BUFFER_SIZE;
                    while (readLength == BUFFER_SIZE) {
                        byte[] buf = exporter.read(offset,  BUFFER_SIZE);
                        fos.write(buf);
                        readLength = buf.length;
                        offset += readLength;
                    }
                    System.out.println("Exported Image:" + imageId +
                                       " to " + file.getAbsolutePath());
                } finally {
                    exporter.close();
                }
            }
        }
    }

    public void cleanup() throws ServerError {
        client.closeSession();
    }

    public static void main(String[] args) throws Exception {
        if (args.length != 5) {
            System.err.println("Usage: OmeroDatasetExport <host> <port> " +
                               "<username> <password> <dataset_id>");
        }
        String host = args[0];
        int port = Integer.parseInt(args[1]);
        String username = args[2];
        String password = args[3];
        long datasetId = Long.parseLong(args[4]);
        OmeroDatasetExport test = null;
        try {
            test = new OmeroDatasetExport(host, port, username, password);
            test.exportImages(datasetId);
        } finally {
            if (test != null) {
                test.cleanup();
            }
        }
    }
}