D-Net

package eu.dnetlib.data.collector.plugins.oai;

import java.io.StringReader;

import java.io.UnsupportedEncodingException;

import java.net.URLEncoder;

import java.util.Iterator;

import java.util.Map;

import java.util.Queue;

import java.util.concurrent.PriorityBlockingQueue;

import com.google.common.collect.Maps;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.dom4j.Document;

import org.dom4j.DocumentException;

import org.dom4j.Node;

import org.dom4j.io.SAXReader;

import eu.dnetlib.data.collector.plugins.HttpConnector;

import eu.dnetlib.data.collector.rmi.CollectorServiceException;

import eu.dnetlib.data.collector.plugins.oai.engine.XmlCleaner;

public class OaiIterator implements Iterator<String> {

	private static final Log log = LogFactory.getLog(OaiIterator.class); // NOPMD by marko on 11/24/08 5:02 PM

	private Queue<String> queue = new PriorityBlockingQueue<String>();

	private SAXReader reader = new SAXReader();

	private String baseUrl;

	private String set;

	private String mdFormat;

	private String fromDate;

	private String untilDate;

	private String token;

	private boolean started;

	private HttpConnector httpConnector;

	private Map<String, String> requestHeader;

	public OaiIterator(final String baseUrl, final String mdFormat, final String set, final String fromDate, final String untilDate, final HttpConnector httpConnector) {

		this.baseUrl = baseUrl;

		this.mdFormat = mdFormat;

		this.set = set;

		this.fromDate = fromDate;

		this.untilDate = untilDate;

		this.started = false;

		this.httpConnector = httpConnector;

		this.requestHeader = requestHeader != null ? requestHeader : Maps.newHashMap();

	}

	private void verifyStarted() {

		if (!this.started) {

			this.started = true;

			try {

				this.token = firstPage();

			} catch (CollectorServiceException e) {

				throw new RuntimeException(e);

			}

		}

	}

	@Override

	public boolean hasNext() {

		synchronized (queue) {

			verifyStarted();

			return !queue.isEmpty();

		}

	}

	@Override

	public String next() {

		synchronized (queue) {

			verifyStarted();

			final String res = queue.poll();

			while (queue.isEmpty() && (token != null) && !token.isEmpty()) {

				try {

					token = otherPages(token);

				} catch (CollectorServiceException e) {

					throw new RuntimeException(e);

				}

			}

			return res;

		}

	}

	@Override

	public void remove() {}

	private String firstPage() throws CollectorServiceException {

		try {

			String url = baseUrl + "?verb=ListRecords&metadataPrefix=" + URLEncoder.encode(mdFormat,"UTF-8");

			if ((set != null) && !set.isEmpty()) {

				url += "&set=" + URLEncoder.encode(set,"UTF-8");

			} 

			// ISO8601 , http://www.openarchives.org/OAI/openarchivesprotocol.html#Dates , eg. 1957-03-20T20:30:00Z

			if ((fromDate != null) && (fromDate.matches("\\d{4}-\\d{2}-\\d{2}") || fromDate.matches("\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z"))  ) {

				url += "&from=" + URLEncoder.encode(fromDate,"UTF-8");

			}

			if ((untilDate != null) && (untilDate.matches("\\d{4}-\\d{2}-\\d{2}") || untilDate.matches("\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z")) ) {

				url += "&until=" + URLEncoder.encode(untilDate,"UTF-8");

			}

			log.info("Start harvesting using url: " + url);

			return downloadPage(url);

		} catch(UnsupportedEncodingException e) {

			throw new CollectorServiceException(e);

		}

	}

	private String extractResumptionToken(final String xml) {

		final String s = StringUtils.substringAfter(xml, "<resumptionToken");

		if (s == null){

			return null;

		}

		final String result = StringUtils.substringBetween(s, ">", "</");

		if (result == null)

			return null;

		return  result.trim();

	}

	private String otherPages(final String resumptionToken) throws CollectorServiceException {

		try {

			return downloadPage(baseUrl + "?verb=ListRecords&resumptionToken=" + URLEncoder.encode(resumptionToken,"UTF-8"));

		} catch (UnsupportedEncodingException e) {

			throw new CollectorServiceException(e);

		}

	}

    private String downloadPage(final String url) throws CollectorServiceException {

        final String xml = httpConnector.getInputSource(url);

		Document doc;

        try {

            doc = reader.read(new StringReader(xml));

        } catch (DocumentException e) {

            log.warn("Error parsing xml, I try to clean it", e);

            final String cleaned = XmlCleaner.cleanAllEntities(xml);

            try {

                doc = reader.read(new StringReader(cleaned));

            } catch (DocumentException e1) {

                final String resumptionToken = extractResumptionToken(xml);

                if (resumptionToken == null)

                    throw new CollectorServiceException("Error parsing cleaned document", e1);

                return resumptionToken;

            }

        }

        final Node errorNode = doc.selectSingleNode("/*[local-name()='OAI-PMH']/*[local-name()='error']");

        if (errorNode != null) {

            final String code = errorNode.valueOf("@code");

            if ("noRecordsMatch".equalsIgnoreCase(code.trim())) {

                log.warn("noRecordsMatch for oai call: " + url);

                return null;

            } else {

                throw new CollectorServiceException(code + " - " + errorNode.getText());

            }

        }

		int added_records = 0;

        for (Object o : doc.selectNodes("//*[local-name()='ListRecords']/*[local-name()='record']")) {

            queue.add(((Node) o).asXML());

            added_records++;

        }

        if(added_records == 0){

        	throw new RuntimeException("OAI Iterator Exception: Empty records list for " + url);

		}

        return doc.valueOf("//*[local-name()='resumptionToken']");

    }

}