ImgScroll/src/main/java/com/rarchives/ripme/ripper/AlbumRipper.java

package com.rarchives.ripme.ripper;

import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import org.jsoup.Connection;
import org.jsoup.Connection.Method;
import org.jsoup.Connection.Response;
import org.jsoup.Jsoup;

import com.rarchives.ripme.ui.RipStatusMessage;
import com.rarchives.ripme.ui.RipStatusMessage.STATUS;
import com.rarchives.ripme.utils.Utils;

public abstract class AlbumRipper extends AbstractRipper {

    protected Map<URL, File> itemsPending = Collections.synchronizedMap(new HashMap<URL, File>());
    protected Map<URL, File> itemsCompleted = Collections.synchronizedMap(new HashMap<URL, File>());
    protected Map<URL, String> itemsErrored = Collections.synchronizedMap(new HashMap<URL, String>());

    public AlbumRipper(URL url) throws IOException {
        super(url);
    }

    public abstract boolean canRip(URL url);
    public abstract URL sanitizeURL(URL url) throws MalformedURLException;
    public abstract void rip() throws IOException;
    public abstract String getHost();
    public abstract String getGID(URL url) throws MalformedURLException;

    public boolean allowDuplicates() {
        return false;
    }

    public void addURLToDownload(URL url, File saveAs, String referrer, Map<String,String> cookies) {
        if (!allowDuplicates()
                && ( itemsPending.containsKey(url)
                  || itemsCompleted.containsKey(url)
                  || itemsErrored.containsKey(url) )) {
            // Item is already downloaded/downloading, skip it.
            logger.info("[!] Skipping " + url + " -- already attempted: " + Utils.removeCWD(saveAs));
            return;
        }
        itemsPending.put(url, saveAs);
        DownloadFileThread dft = new DownloadFileThread(url,  saveAs,  this);
        if (referrer != null) {
            dft.setReferrer(referrer);
        }
        if (cookies != null) {
            dft.setCookies(cookies);
        }
        threadPool.addThread(dft);
    }

    @Override
    public void addURLToDownload(URL url, File saveAs) {
        addURLToDownload(url, saveAs, null, null);
    }

    /**
     * Queues image to be downloaded and saved.
     * Uses filename from URL to decide filename.
     * @param url
     *      URL to download
     */
    public void addURLToDownload(URL url) {
        // Use empty prefix and empty subdirectory
        addURLToDownload(url, "", "");
    }

    @Override
    public void downloadCompleted(URL url, File saveAs) {
        if (observer == null) {
            return;
        }
        try {
            String path = Utils.removeCWD(saveAs);
            RipStatusMessage msg = new RipStatusMessage(STATUS.DOWNLOAD_COMPLETE, path);
            itemsPending.remove(url);
            itemsCompleted.put(url, saveAs);
            observer.update(this, msg);

            checkIfComplete();
        } catch (Exception e) {
            logger.error("Exception while updating observer: ", e);
        }
    }

    @Override
    public void downloadErrored(URL url, String reason) {
        if (observer == null) {
            return;
        }
        itemsPending.remove(url);
        itemsErrored.put(url, reason);
        observer.update(this, new RipStatusMessage(STATUS.DOWNLOAD_ERRORED, url + " : " + reason));

        checkIfComplete();
    }

    @Override
    public void downloadProblem(URL url, String message) {
        if (observer == null) {
            return;
        }
        
        itemsPending.remove(url);
        itemsErrored.put(url, message);
        observer.update(this, new RipStatusMessage(STATUS.DOWNLOAD_WARN, url + " : " + message));
            
        checkIfComplete();
    }

    /**
     * Notifies observers and updates state if all files have been ripped.
     */
    @Override
    protected void checkIfComplete() {
        if (observer == null) {
            return;
        }
        if (itemsPending.isEmpty()) {
            super.checkIfComplete();
        }
    }

    /**
     * Sets directory to save all ripped files to.
     * @param url
     *      URL to define how the workin directory should be saved.
     */
    @Override
    public void setWorkingDir(URL url) throws IOException {
        String path = Utils.getWorkingDirectory().getCanonicalPath();
        if (!path.endsWith(File.separator)) {
            path += File.separator;
        }
        String title = getAlbumTitle(this.url);
        title = Utils.filesystemSafe(title);
        path += title + File.separator;
        this.workingDir = new File(path);
        if (!this.workingDir.exists()) {
            logger.info("[+] Creating directory: " + Utils.removeCWD(this.workingDir));
            this.workingDir.mkdirs();
        }
        logger.debug("Set working directory to: " + this.workingDir);
    }

    /**
     * @return
     *      Integer between 0 and 100 defining the progress of the album rip.
     */
    @Override
    public int getCompletionPercentage() {
        double total = itemsPending.size()  + itemsErrored.size() + itemsCompleted.size();
        return (int) (100 * ( (total - itemsPending.size()) / total));
    }

    /**
     * @return
     *      Human-readable information on the status of the current rip.
     */
    @Override
    public String getStatusText() {
        StringBuilder sb = new StringBuilder();
        sb.append(getCompletionPercentage())
          .append("% ")
          .append("- Pending: "  ).append(itemsPending.size())
          .append(", Completed: ").append(itemsCompleted.size())
          .append(", Errored: "  ).append(itemsErrored.size());
        return sb.toString();
    }
    
    public Response getResponse(String url, 
                                Method method,
                                String userAgent, 
                                String referrer, 
                                Map<String,String> cookies, 
                                boolean ignoreContentType)
                    throws IOException {
        Connection connection = Jsoup.connect(url);

        if (method == null) {
            method = Method.GET;
        }
        connection.method(method);

        if (userAgent == null) {
            userAgent = USER_AGENT;
        }
        connection.userAgent(userAgent);

        if (cookies != null) {
            connection.cookies(cookies);
        }

        if (referrer != null) {
            connection.referrer(referrer);
        }
        connection.ignoreContentType(ignoreContentType);

        connection.maxBodySize(0);
        Response response = null;
        int retries = Utils.getConfigInteger("download.retries", 1);;
        while (retries >= 0) {
            retries--;
            try {
                response = connection.execute();
            } catch (IOException e) {
                logger.warn("Error while loading " + url, e);
                continue;
            }
        }
        return response;
    }
}