import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;

import org.jaxen.JaxenException;
import org.jaxen.jdom.JDOMXPath;
import org.jdom.Attribute;
import org.jdom.Element;
import org.jdom.JDOMException;
import org.jdom.Namespace;
import org.jdom.input.SAXBuilder;
import org.jdom.output.XMLOutputter;

/*
 * Created on Jul 1, 2004
 *
 */

/**
 * This is a simple class to fix up the XMI that is produced by Enterprise
 * Architect so that it works with XMI reader from NetBeans MDR. I'm using it to
 * read XMI 1.1 that represents an instance of a UML 1.3 model.
 * 
 * To use this class, you must check only the following options in the EA
 * "Export Package to XML" dialog: 
 * - Unisys/Rose Format 
 * - Format XML Output 
 * - Write Log File
 * 
 * Be sure to NOT check the XMI 1.0 option. The last two are probably optional,
 * but that's how I did it.
 * 
 * To compile and run, I used the following jars: 
 * - jaxen-1.0-FCS-full.jar: http://jaxen.org 
 * - jdom-1.0b8.jar: http://www.jdom.org 
 * - saxpath-1.0-FCS.jar: http://saxpath.sourceforge.net 
 * - xercesImpl-2.4.0.jar: http://xml.apache.org/xerces2-j/index.html 
 * - xml-apis-2.0.2.jar: http://xml.apache.org/xerces2-j/index.html
 * 
 * The latest of all of these should work fine.
 * 
 * If you have any trouble, just send me an email (or to the MDR users list
 * as I always monitor that).
 * 
 * But on the legal side of things... I take absolutely no responsibility for
 * the havoc this code may wreak on your XMI. Do whatever you like with this
 * code.
 * 
 * @author Joshua Phillips (joshua.a.phillips@comcast.net)
 *  
 */
public class FixEAXMI {

	/**
	 * Modifies a copy of the XMI file named by
	 * xmiIn and writes this copy to the file named by xmiOut.
	 * 
	 * @param xmiIn the EA XMI file
	 * @param xmiOut the fixed EA XMI file
	 * 
	 * @throws JDOMException
	 * @throws JaxenException
	 * @throws IOException
	 */
	public void fix(String xmiIn, String xmiOut) throws JDOMException,
			JaxenException, IOException {
		
		Element root = (new SAXBuilder(false)).build(new File(xmiIn))
				.getRootElement();

		//Add unique IDs to all elements
		Element model = (Element) (new JDOMXPath(
				"*[local-name()='XMI.content']/*[local-name()='Model']"))
				.selectSingleNode(root);
		fixIDs(model);

		//Move data from 'behaviour' TaggedValue elements into 
		// Operation.specification elements
		fixOperations(root);

		//Remove diagrams elements
		removeDiagramsElements(root);

		//Move TaggedValue elements to child of XMI.content
		fixTaggedValues(root);

		//Move Stereotype elements to child of XMI.content
		fixSterotypes(root);

		//Remove ModelElement.taggedValue elements
		removeMETaggedValues(root);

		//Remove ModelElement.stereotype elements
		removeMEStereotypes(root);

		//Remove XML.difference elements
		removeDifference(root);

		//Remove attributes that are redundant with child elements
		collapseAttributes(root);

		System.out.println("writing to: " + xmiOut);
		Writer writer = new FileWriter(new File(xmiOut));
		XMLOutputter xmlout = new XMLOutputter();
		xmlout.setIndent("   ");
		xmlout.setNewlines(true);
		writer.write(xmlout.outputString(root));
		writer.flush();
		writer.close();
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void removeDifference(Element root) throws JaxenException {
		String exp = "//*[local-name()='XMI.difference']";
		Collection nodes = (new JDOMXPath(exp)).selectNodes(root);
		List toDetach = new ArrayList();
		for (Iterator i = nodes.iterator(); i.hasNext();) {
			Element n = (Element) i.next();
			toDetach.add(n);
		}
		detachElements(toDetach);
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void removeMEStereotypes(Element root) throws JaxenException {
		String exp = "//*[local-name()='ModelElement.stereotype']";
		Collection nodes = (new JDOMXPath(exp)).selectNodes(root);
		List toDetach = new ArrayList();
		for (Iterator i = nodes.iterator(); i.hasNext();) {
			Element n = (Element) i.next();
			toDetach.add(n);
		}
		detachElements(toDetach);
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void removeMETaggedValues(Element root) throws JaxenException {
		String exp = "//*[local-name()='ModelElement.taggedValue']";
		Collection mtvs = (new JDOMXPath(exp)).selectNodes(root);
		List toDetach = new ArrayList();
		for (Iterator i = mtvs.iterator(); i.hasNext();) {
			Element mtv = (Element) i.next();
			//System.out.println(mtv.getParent().getAttributeValue("xmi.id"));
			toDetach.add(mtv);
		}
		detachElements(toDetach);
	}

	/**
	 * @param toDetach
	 */
	private void detachElements(List toDetach) {
		for (Iterator i = toDetach.iterator(); i.hasNext();) {
			Element e = (Element) i.next();
			e.detach();
		}
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void removeDiagramsElements(Element root) throws JaxenException {
		List toDetach = new ArrayList();
		String exp = "//*[local-name()='DiagramElement']";
		Collection nodes = (new JDOMXPath(exp)).selectNodes(root);
		for (Iterator i = nodes.iterator(); i.hasNext();) {
			Element n = (Element) i.next();
			toDetach.add(n);
		}
		for (Iterator i = toDetach.iterator(); i.hasNext();) {
			Element e = (Element) i.next();
			e.detach();
		}
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void fixSterotypes(Element root) throws JaxenException {
		String exp = "//*[local-name()='ModelElement.stereotype']";
		Collection nodes = (new JDOMXPath(exp)).selectNodes(root);
		Element contentEl = root.getChild("XMI.content");
		HashMap sHash = new HashMap();
		//First, fix existing stereotypes
		for (Iterator i = contentEl.getChildren().iterator(); i.hasNext();) {
			Element sterEl = (Element) i.next();
			if (sterEl.getName().endsWith("Stereotype")) {
				System.out.println("On stereotype: "
						+ sterEl.getAttributeValue("name"));
				Collection els = (new JDOMXPath(
						"*[local-name()='Stereotype.extendedElement']/*[local-name()='Foundation.Core.ModelElement']"))
						.selectNodes(sterEl);
				System.err.println("els.size() = " + els.size());
				for (Iterator j = els.iterator(); j.hasNext();) {
					Element e = (Element) j.next();
					String v = sterEl.getAttributeValue("extendedElement");
					if (v == null) {
						v = "";
					}
					sterEl.setAttribute("extendedElement", v + " "
							+ e.getAttributeValue("xmi.idref") + " ");
				}
				Element child = (Element) (new JDOMXPath(
						"//*[local-name()='Stereotype.extendedElement']"))
						.selectSingleNode(sterEl);
				child.detach();
				sHash.put(sterEl.getAttributeValue("name"), sterEl);
			}
		}

		for (Iterator i = nodes.iterator(); i.hasNext();) {
			Element meSEl = (Element) i.next();
			Element meParentEl = meSEl.getParent();
			String parentId = meParentEl.getAttributeValue("xmi.id");
			if (parentId == null) {
				throw new RuntimeException("no xmi.id for "
						+ meParentEl.getName());
			}
			int idSeq = 0;
			for (Iterator j = meSEl.getChildren().iterator(); j.hasNext(); idSeq++) {
				Element sEl = (Element) j.next();
				String sName = sEl.getAttributeValue("name");
				Element newSEl = (Element) sHash.get(sName);
				if (newSEl == null) {
					newSEl = new Element("Stereotype");
					contentEl.addContent(newSEl);
					newSEl.setNamespace(Namespace.getNamespace("UML",
							"href://org.omg/UML"));
					newSEl.setAttribute("xmi.id", parentId + "_fix_ster_"
							+ idSeq);
					newSEl.setAttribute("name", sName);
					newSEl.setAttribute("visibility", "public");
					newSEl.setAttribute("isSpecification", "false");
					newSEl.setAttribute("isRoot", "false");
					newSEl.setAttribute("isLeaf", "false");
					newSEl.setAttribute("isAbstract", "false");
					newSEl.setAttribute("baseClass", meParentEl.getName());
				}
				String ext = newSEl.getAttributeValue("extendedElement");
				if (ext == null) {
					ext = "";
				}
				if (ext.indexOf(" " + parentId + " ") == -1) {
					//System.out.println("adding me " + parentId + " to " +
					// sName);
					newSEl.setAttribute("extendedElement", ext + " " + parentId
							+ " ");
				}

			}
		}

	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void fixTaggedValues(Element root) throws JaxenException {
		String exp = "//*[local-name()='ModelElement.taggedValue']";
		Collection nodes = (new JDOMXPath(exp)).selectNodes(root);
		Element contentEl = root.getChild("XMI.content");
		for (Iterator i = nodes.iterator(); i.hasNext();) {
			Element meTVEl = (Element) i.next();
			Element meParentEl = meTVEl.getParent();
			String parentId = meParentEl.getAttributeValue("xmi.id");
			if (parentId == null) {
				throw new RuntimeException("no xmi.id for "
						+ meParentEl.getName());
			}
			int idSeq = 0;
			for (Iterator j = meTVEl.getChildren().iterator(); j.hasNext(); idSeq++) {
				Element tvEl = (Element) j.next();
				Element newTVEl = new Element("TaggedValue");
				contentEl.addContent(newTVEl);
				newTVEl.setNamespace(Namespace.getNamespace("UML",
						"href://org.omg/UML"));
				newTVEl.setAttribute("xmi.id", parentId + "_fix_" + idSeq);
				newTVEl.setAttribute("tag", tvEl.getAttributeValue("tag"));
				newTVEl.setAttribute("modelElement", parentId);
				newTVEl.setAttribute("value", tvEl.getAttributeValue("value"));
			}
		}
	}

	/**
	 * @param root
	 */
	private void fixIDs(Element parent) {
		int idSeq = 0;
		String parentId = parent.getAttributeValue("xmi.id");
		if (parentId == null) {
			throw new RuntimeException("xmi.id is null for " + parent.getName());
		}
		for (Iterator i = parent.getChildren().iterator(); i.hasNext(); idSeq++) {
			Element child = (Element) i.next();
			String childName = child.getName();
			if (!childName.endsWith(".taggedValue")
					&& !childName.endsWith(".stereotype")) {

				//Fix
				String childId = child.getAttributeValue("xmi.id");
				if (childId == null) {
					childId = parentId + "_fix_" + idSeq;
					child.setAttribute("xmi.id", childId);
				}

				fixIDs(child);
			}
		}
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void fixOperations(Element root) throws JaxenException {
		String exp = "//*[local-name()='Operation']";
		Collection ops = (new JDOMXPath(exp)).selectNodes(root);
		//System.out.println("Found " + ops.size() + " ops");
		for (Iterator i = ops.iterator(); i.hasNext();) {
			Element opEl = (Element) i.next();
			exp = "*[local-name()='ModelElement.taggedValue']/*[@tag='behaviour']";
			Collection specs = (new JDOMXPath(exp)).selectNodes(opEl);
			if (specs.size() > 1) {
				throw new RuntimeException("more than one spec for op "
						+ opEl.getAttributeValue("name"));
			}
			if (specs.size() == 1) {
				Element oldSpecEl = (Element) specs.iterator().next();
				Element newSpecEl = new Element("Operation.specification");
				newSpecEl.setNamespace(Namespace.getNamespace("UML",
						"href://org.omg.UML/1.3"));
				opEl.addContent(newSpecEl);
				newSpecEl.addContent(oldSpecEl.getAttributeValue("value"));
			}
			//System.out.println("Found " + specs.size() + " specs");
		}
	}

	/**
	 * @param root
	 * @throws JaxenException
	 */
	private void collapseAttributes(Element el) throws JaxenException {
		List toRemove = new ArrayList();
		for (Iterator i = el.getAttributes().iterator(); i.hasNext();) {
			Attribute att = (Attribute) i.next();
			String exp = "*[ends-with(local-name(), '" + att.getName() + "')]";
			Collection nodes = (new JDOMXPath(exp)).selectNodes(el);
			if (nodes.size() > 0) {
				toRemove.add(att);
			}
		}
		for (Iterator i = toRemove.iterator(); i.hasNext();) {
			el.removeAttribute((Attribute) i.next());
		}
		for (Iterator i = el.getChildren().iterator(); i.hasNext();) {
			collapseAttributes((Element) i.next());
		}
	}

	public static void main(String[] args) throws JaxenException,
			JDOMException, IOException {
		if (args.length != 2) {
			System.out
					.println("Usage: <xmiIn> <xmiOut>\n(e.g., my_eq_model.xmi my_fixed_eq_model.xml");
		} else {
			FixEAXMI f = new FixEAXMI();
			f.fix(args[0], args[1]);
		}
	}
}