package com.redfin.sitemapgenerator;

import org.xml.sax.SAXException;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.net.URL;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.GZIPOutputStream;

abstract class SitemapGenerator<U extends ISitemapUrl, THIS extends SitemapGenerator<U,THIS>> {
	
/** 50000 URLs per sitemap maximum */
	
public static final int MAX_URLS_PER_SITEMAP = 50000;
	

	
private final URL baseUrl;
	
private final File baseDir;
	
private final String fileNamePrefix;
	
private final String fileNameSuffix;
	
private final boolean allowEmptySitemap;
	
private final boolean allowMultipleSitemaps;
	
private final ArrayList<U> urls = new ArrayList<U>();
	
private final W3CDateFormat dateFormat;
	
private final int maxUrls;
	
private final boolean autoValidate;
	
private final boolean gzip;
	
private final ISitemapUrlRenderer<U> renderer;
	
private int mapCount = 0;
	
private boolean finished = false;
	

	
private final ArrayList<File> outFiles = new ArrayList<File>();
	

	
public SitemapGenerator(AbstractSitemapGeneratorOptions<?> options, ISitemapUrlRenderer<U> renderer) {
		
baseDir = options.baseDir;
		
baseUrl = options.baseUrl;
		
fileNamePrefix = options.fileNamePrefix;
		
W3CDateFormat dateFormat = options.dateFormat;
		
if (dateFormat == null) dateFormat = new W3CDateFormat();
		
this.dateFormat = dateFormat;
		
allowEmptySitemap = options.allowEmptySitemap;
		
allowMultipleSitemaps = options.allowMultipleSitemaps;
		
maxUrls = options.maxUrls;
		
autoValidate = options.autoValidate;
		
gzip = options.gzip;
		
this.renderer = renderer;

		
if(options.suffixStringPattern != null && !options.suffixStringPattern.isEmpty()) {
			
fileNameSuffix = gzip ? options.suffixStringPattern + ".xml.gz" : options.suffixStringPattern + ".xml";
		
}
		
else {
			
fileNameSuffix = gzip ? ".xml.gz" : ".xml";
		
}
	
}

	
/** Add one URL of the appropriate type to this sitemap.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or else write out one sitemap immediately.
	 
* @param url the URL to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrl(U url) {
		
if (finished) throw new RuntimeException("Sitemap already printed; you must create a new generator to make more sitemaps");
		
UrlUtils.checkUrl(url.getUrl(), baseUrl);
		
if (urls.size() == maxUrls) {
			
if (!allowMultipleSitemaps) throw new RuntimeException("More than " + maxUrls + " urls, but allowMultipleSitemaps is false. Enable allowMultipleSitemaps to split the sitemap into multiple files with a sitemap index.");
			
if (baseDir != null) {
				
if (mapCount == 0) mapCount++;
				
try {
					
writeSiteMap();
				
} catch(IOException ex) {
					
throw new RuntimeException("Closing of stream failed.", ex);
				
}
				
mapCount++;
				
urls.clear();
			
}
		
}
		
urls.add(url);
		
return getThis();
	
}
	

	
/** Add multiple URLs of the appropriate type to this sitemap, one at a time.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or write out one sitemap immediately.
	 
* @param urls the URLs to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrls(Iterable<? extends U> urls) {
		
for (U url : urls) addUrl(url);
		
return getThis();
	
}
	

	
/** Add multiple URLs of the appropriate type to this sitemap, one at a time.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or write out one sitemap immediately.
	 
* @param urls the URLs to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrls(U... urls) {
		
for (U url : urls) addUrl(url);
		
return getThis();
	
}
	

	
/** Add multiple URLs of the appropriate type to this sitemap, one at a time.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or write out one sitemap immediately.
	 
* @param urls the URLs to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrls(String... urls) {
		
for (String url : urls) addUrl(url);
		
return getThis();
	
}
	

	
/** Add one URL of the appropriate type to this sitemap.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or else write out one sitemap immediately.
	 
* @param url the URL to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrl(String url) {
		
U sitemapUrl;
		
try {
			
sitemapUrl = renderer.getUrlClass().getConstructor(String.class).newInstance(url);
			
return addUrl(sitemapUrl);
		
} catch (Exception e) {
			
throw new RuntimeException(e);
		
}
	
}
	

	
/** Add multiple URLs of the appropriate type to this sitemap, one at a time.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or write out one sitemap immediately.
	 
* @param urls the URLs to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrls(URL... urls) {
		
for (URL url : urls) addUrl(url);
		
return getThis();
	
}
	

	
/** Add one URL of the appropriate type to this sitemap.
	 
* If we have reached the maximum number of URLs, we'll throw an exception if
 
 
is false,
	 
* or write out one sitemap immediately.
	 
* @param url the URL to add to this sitemap
	 
* @return this
	 
*/

	
public THIS addUrl(URL url) {
		
U sitemapUrl;
		
try {
			
sitemapUrl = renderer.getUrlClass().getConstructor(URL.class).newInstance(url);
			
return addUrl(sitemapUrl);
		
} catch (Exception e) {
			
throw new RuntimeException(e);
		
}
	
}
	

	
@SuppressWarnings("unchecked")
	
THIS getThis() {
		
return (THIS)this;
	
}
	

	
/** Write out remaining URLs; this method can only be called once.
  
This is necessary so we can keep an accurate count for
 
.
	 
*
	 
* @return a list of files we wrote out to disk
	 
*/

	
public List<File> write() {
		
if (finished) throw new RuntimeException("Sitemap already printed; you must create a new generator to make more sitemaps");
		
if (!allowEmptySitemap && urls.isEmpty() && mapCount == 0) throw new RuntimeException("No URLs added, sitemap would be empty; you must add some URLs with addUrls");
		
try {
			
writeSiteMap();
		
} catch (IOException ex) {
			
throw new RuntimeException("Closing of streams has failed at some point.", ex);
		
}
		
finished = true;
		
return outFiles;
	
}
	

	
/**
	 
* Writes out the sitemaps as a list of strings.
	 
* Each string in the list is a formatted list of URLs.
	 
* We return a list because the URLs may not all fit --
	 
* google specifies a maximum of 50,000 URLs in one sitemap.
	 
* @return a list of XML-formatted strings
	 
*/

	
public List<String> writeAsStrings() {
		
List<String> listOfSiteMapStrings = new ArrayList<String>();
		
for (int start = 0; start < urls.size(); start += maxUrls) {
			
int end = start + maxUrls;
			
if (end > urls.size()) {
				
end = urls.size();
			
}
			
StringBuilder sb = new StringBuilder();
			
writeSiteMapAsString(sb, urls.subList(start, end));
			
listOfSiteMapStrings.add(sb.toString());
		
}
		
return listOfSiteMapStrings;
	
}
	

	
private void writeSiteMapAsString(StringBuilder sb, List<U> urls) {
		
sb.append("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
		
sb.append("<urlset xmlns=\"http://www.sitemaps.org/schemas/sitemap/0.9\" ");
		
if (renderer.getXmlNamespaces() != null) {
			
sb.append(renderer.getXmlNamespaces());
			
sb.append(' ');
		
}
		
sb.append(">\n");
		
for (U url : urls) {
			
renderer.render(url, sb, dateFormat);
		
}
		
sb.append("</urlset>");
	
}
	

	
/**
	 
* After you've called
 
, call this to generate a sitemap index of all sitemaps you generated.
	 
* The sitemap index is written to {baseDir}/sitemap_index.xml
	 
*/

	
public File writeSitemapsWithIndex() {
		
return writeSitemapsWithIndex(new File(baseDir, "sitemap_index.xml"));
	
}

	
/**
	 
* After you've called
 
, call this to generate a sitemap index of all sitemaps you generated.
	 
*/

	
public String writeSitemapsWithIndexAsString() {
		
return prepareSitemapIndexGenerator(null).writeAsString();
	
}

	
/**
	 
* After you've called
 
, call this to generate a sitemap index of all sitemaps you generated.
	 
*
	 
* @param outFile the destination file of the sitemap index.
	 
*/

	
public File writeSitemapsWithIndex(File outFile) {
		
prepareSitemapIndexGenerator(outFile).write();
		
return outFile;
	
}

	
private SitemapIndexGenerator prepareSitemapIndexGenerator(File outFile) {
		
if (!finished) throw new RuntimeException("Sitemaps not generated yet; call write() first");
		
SitemapIndexGenerator sig;
		
sig = new SitemapIndexGenerator.Options(baseUrl, outFile).dateFormat(dateFormat).autoValidate(autoValidate).build();
		
sig.addUrls(fileNamePrefix, fileNameSuffix, mapCount);
		
return sig;
	
}
	

	
private void writeSiteMap() throws IOException {
		
if (baseDir == null) {
			
throw new NullPointerException("To write to files, baseDir must not be null");
		
}
		
if (urls.isEmpty() && (mapCount > 0 || !allowEmptySitemap)) return;
		
String fileNamePrefix;
		
if (mapCount > 0) {
			
fileNamePrefix = this.fileNamePrefix + mapCount;
		
} else {
			
fileNamePrefix = this.fileNamePrefix;
		
}
		
File outFile = new File(baseDir, fileNamePrefix+fileNameSuffix);
		
outFiles.add(outFile);

		
OutputStreamWriter out = null;
		
try {
			
if (gzip) {
				
FileOutputStream fileStream = new FileOutputStream(outFile);
				
GZIPOutputStream gzipStream = new GZIPOutputStream(fileStream);
				
out = new OutputStreamWriter(gzipStream, Charset.forName("UTF-8").newEncoder());
			
} else {
				
out = new OutputStreamWriter(new FileOutputStream(outFile), Charset.forName("UTF-8").newEncoder());
			
}

			
writeSiteMap(out);
			
out.flush();

			
if (autoValidate) SitemapValidator.validateWebSitemap(outFile);
		
} catch (IOException e) {
			
throw new RuntimeException("Problem writing sitemap file " + outFile, e);
		
} catch (SAXException e) {
			
throw new RuntimeException("Sitemap file failed to validate (bug?)", e);
		
} finally {
			
if(out != null) {
				
out.close();
			
}
		
}
	
}
	

	
private void writeSiteMap(OutputStreamWriter out) throws IOException {
		
StringBuilder sb = new StringBuilder();
		
writeSiteMapAsString(sb, urls);
		
out.write(sb.toString());
	
}
	

}