package com.solvebio.xmlsplit;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.util.List;

import net.sourceforge.argparse4j.ArgumentParsers;
import net.sourceforge.argparse4j.inf.ArgumentParser;
import net.sourceforge.argparse4j.inf.ArgumentParserException;
import net.sourceforge.argparse4j.inf.Namespace;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.LineIterator;

import com.google.common.collect.Lists;

public class XmlSplit {
	public static void main(String[] args) throws IOException {
		ArgumentParser parser = ArgumentParsers.newArgumentParser("XmlSplit").defaultHelp(true)
				.description("Split a single, large XML document into 'N' smaller documents.");
		parser.addArgument("-n", "--number").type(Integer.class).setDefault(10).help("number of output files");
		parser.addArgument("input_filepath").metavar("[ input filepath ]").help("input XML filepath");
		parser.addArgument("output_directory").metavar("[ output directory ]").help("root XML output directory");
		parser.addArgument("tag").metavar("[ element tag ]").help("the element to split on");
		Namespace ns = null;
		try {
			ns = parser.parseArgs(args);
		} catch (ArgumentParserException e) {
			parser.handleError(e);
			System.exit(1);
		}

		long t0 = System.currentTimeMillis();

		int writerIdx = 0;
		Writer[] fileWriters = new BufferedWriter[ns.get("number")];
		for (int i = 0; i < fileWriters.length; i++) {
			String fname = String.format("./sample_output/%s.xml", i);
			FileWriter fw = new FileWriter(fname);
			System.out.println("Creating output file: " + fname);
			fileWriters[i] = new BufferedWriter(fw);
		}

		String tag = ns.getString("tag").replace("<", "").replace(">", "");
		List<String> header = Lists.newArrayList();
		List<String> lines = Lists.newArrayList();
		List<String> footer = Lists.newArrayList();

		int readCnt = 0;
		int elemCnt = 0;
		LineIterator iter = FileUtils.lineIterator(new File(ns.getString("input_filepath")));
		boolean foundOne = false;
		while (iter.hasNext()) {
			if (++readCnt % 1000000 == 0) {
				System.out.println(String.format("read %s lines", readCnt));
			}

			// sanity check. if we've gone 10,000 lines without hitting an opening tag, puke
			if (readCnt == 10000 && !foundOne) {
				System.out.println(String.format("Couldn't find opening tag '%s' after 10K lines. Exiting.", tag));
				System.exit(1);
			}

			String line = iter.nextLine();
			String lineTrim = line.trim();

			// skip empty lines
			if (lineTrim.isEmpty()) {
				continue;
			}

			// populate header
			if (!foundOne && !lineTrim.startsWith("<" + tag)) {
				header.add(line);
				continue;
			} else {
				if (!foundOne) {
					// write headers for all files
					System.out.println("writing headers...");
					for (Writer w : fileWriters) {
						IOUtils.writeLines(header, IOUtils.LINE_SEPARATOR, w);
					}
				}
				foundOne = true;
			}

			lines.add(line);
			if (lineTrim.startsWith("</" + tag)) {
				if (++elemCnt % 10000 == 0) {
					System.out.println(String.format("processed %s elements", elemCnt));
				}

				// write
				IOUtils.writeLines(lines, IOUtils.LINE_SEPARATOR, fileWriters[(++writerIdx + 1) % fileWriters.length]);

				// round-robin
				// writerIdx = (writerIdx + 1) % fileWriters.length;

				lines.clear();
				footer.clear();
			} else {
				footer.add(line);
			}
		}

		// sanity check. if we've gone 10,000 lines without hitting an opening tag, puke
		if (!foundOne) {
			System.out.println(String.format("Finished reading the file but couldn't find opening tag '%s'. Exiting.",
					tag));
			// TODO: cleanup outputs
			System.exit(1);
		}

		// write footer!
		System.out.println("writing footers...");
		for (Writer w : fileWriters) {
			IOUtils.writeLines(footer, IOUtils.LINE_SEPARATOR, w);
		}

		// close all writers
		for (Writer w : fileWriters) {
			w.close();
		}

		long t1 = System.currentTimeMillis();
		System.out.println(String.format("done! read %s lines. processed %s elements. took: %s ms.", readCnt, elemCnt,
				t1 - t0));
	}
}

@@ -433,24 +433,68 @@
 .FileUtils;%0A
+import org.apache.commons.io.FilenameUtils;%0A
 import org.a
@@ -933,18 +933,17 @@
 Default(
-10
+5
 ).help(%22
@@ -991,26 +991,25 @@
 ument(%22input
-_f
+F
 ilepath%22).me
@@ -1096,10 +1096,9 @@
 tput
-_d
+D
 irec
@@ -1449,24 +1449,477 @@
 eMillis();%0A%0A
+%09%09// setup paths%0A%09%09String inputFilepath = ns.getString(%22inputFilepath%22);%0A%09%09String outputDirectory = ns.getString(%22outputDirectory%22);%0A%09%09String filename = FilenameUtils.getBaseName(inputFilepath);%0A%09%09String ext = %22.%22 + FilenameUtils.getExtension(inputFilepath);%0A%09%09String outputFilepathFmt = FilenameUtils.concat(outputDirectory, filename + %22_%2502d%22 + ext);%0A%0A%09%09// ensure output director%0A%09%09FileUtils.forceMkdir(new File(outputDirectory));%0A%0A%09%09// setup writers%0A
 %09%09int writer
@@ -2075,77 +2075,28 @@
 mat(
-%22./sample_output/%25s.xml%22, i);%0A%09%09%09FileWriter fw = new FileWriter(fname
+outputFilepathFmt, i
 );%0A%09
@@ -2194,10 +2194,29 @@
 ter(
-fw
+new FileWriter(fname)
 );%0A%09
@@ -2523,38 +2523,21 @@
 ile(
-ns.getString(%22
 input
-_f
+F
 ilepath
-%22)
 ));%0A

write to specified outputdir