D-Net

package eu.dnetlib.data.collector.plugins.schemaorg.sitemapindex;

import eu.dnetlib.data.collector.plugins.schemaorg.Utils;

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

import org.apache.commons.io.FileUtils;

import org.apache.commons.io.IOUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import java.io.*;

import java.net.URL;

import java.nio.charset.Charset;

import java.util.*;

public class SitemapFileIterator implements Iterator<String> {

	private static final Log log = LogFactory.getLog(SitemapFileIterator.class);

	public static class Options {

		public enum SitemapFileType{

			Text,

			GZ

		}

		public enum SitemapSchemaType{

			Text,

			Xml

		}

		public Options(){}

		public Options(URL fileUrl, Charset charset, SitemapSchemaType schemaType, SitemapFileType fileType) {

			this.fileUrl = fileUrl;

			this.charset = charset;

			this.schemaType = schemaType;

			this.fileType = fileType;

		}

		private SitemapFileType fileType;

		private SitemapSchemaType schemaType;

		private URL fileUrl;

		private Charset charset;

		public Charset getCharset() {

			return charset;

		}

		public void setCharset(Charset charset) {

			this.charset = charset;

		}

		public URL getFileUrl() {

			return fileUrl;

		}

		public void setFileUrl(URL fileUrl) {

			this.fileUrl = fileUrl;

		}

		public SitemapFileType getFileType() {

			return fileType;

		}

		public void setFileType(SitemapFileType fileType) {

			this.fileType = fileType;

		}

		public SitemapSchemaType getSchemaType() {

			return schemaType;

		}

		public void setSchemaType(SitemapSchemaType schemaType) {

			this.schemaType = schemaType;

		}

		@Override

		public Object clone(){

			Options clone = new Options();

			clone.setCharset(this.getCharset());

			clone.setFileType(this.getFileType());

			clone.setFileUrl(this.getFileUrl());

			clone.setSchemaType(this.getSchemaType());

			return clone;

		}

	}

	private Options options;

	private File downloadedFile;

	private File contentFile;

	private Queue<String> locations;

	public SitemapFileIterator(Options options){

		this.options = options;

	}

	public void bootstrap() {

		LinkedList<String> endpoints = null;

		try {

			String path = new java.io.File( "." ).getCanonicalPath();

			log.debug(String.format("bootstrapping sitemapindex file access for sitemapindex %s", this.options.getFileUrl()));

			this.downloadedFile = File.createTempFile(UUID.randomUUID().toString(), ".tmp");

			this.downloadedFile.deleteOnExit();

			FileUtils.copyURLToFile(this.options.getFileUrl(), this.downloadedFile);

			log.debug(String.format("downloaded file: %s has size %d", this.downloadedFile.toString(), this.downloadedFile.length()));

			switch (this.options.getFileType()) {

				case Text: {

					this.contentFile = this.downloadedFile;

					break;

				}

				case GZ: {

					this.contentFile = File.createTempFile(UUID.randomUUID().toString(), ".tmp");

					this.contentFile.deleteOnExit();

					Utils.decompressGZipTo(this.downloadedFile, this.contentFile);

					log.debug(String.format("extracted gz file: %s has size %d", this.contentFile.toString(), this.contentFile.length()));

					break;

				}

				default:

					throw new CollectorServiceException("unrecognized file type " + this.options.getFileType());

			}

			List<String> content = this.collectContentLocations();

			log.debug(String.format("extracted %d sitemapindex endpoints", content.size()));

			endpoints = new LinkedList<>(content);

		}catch(Exception ex){

			log.error(String.format("error processing sitemapindex %s. returning 0 endpoints",this.options.getFileUrl()), ex);

			endpoints = new LinkedList<>();

		}finally {

			if (this.contentFile != null) {

				this.contentFile.delete();

			}

			if (this.downloadedFile != null) {

				this.downloadedFile.delete();

			}

		}

		this.locations = endpoints;

	}

	private List<String> collectContentLocations() throws Exception{

		switch(this.options.getSchemaType()) {

			case Text:{

				return this.collectTextContentLocations();

			}

			case Xml:{

				return this.collectXmlContentLocations();

			}

			default: throw new CollectorServiceException("unrecognized file type "+this.options.getFileType());

		}

	}

	private List<String> collectTextContentLocations() throws Exception {

		log.debug(String.format("reading endpoint locations from text sitemapindex"));

		try (FileInputStream in = new FileInputStream(this.contentFile)) {

			return IOUtils.readLines(in, this.options.getCharset());

		}

	}

	private List<String> collectXmlContentLocations() throws Exception {

		log.debug(String.format("reading endpoint locations from xml sitemapindex"));

		return Utils.collectAsStrings(this.contentFile,"/urlset/url/loc/text()");

	}

	@Override

	public boolean hasNext() {

		return !this.locations.isEmpty();

	}

	@Override

	public String next() {

		return this.locations.poll();

	}

}