package org.biomoby.shared;

import org.xbill.DNS.*;  //dnsjava

import org.w3c.dom.*;

import javax.xml.namespace.QName;
import javax.xml.parsers.*;
import javax.xml.transform.*;
import javax.xml.transform.stream.*;
import javax.xml.ws.*;
import javax.xml.xpath.*;

import java.io.*;
import java.net.URL;
import java.util.*;
import java.util.logging.*;
import java.util.regex.*;

/**
 * Implementation of LSID resolution for data and metadata. LSIDs are a OMG standard
 * for referring to pieces of biomedical data using URNs. See http://lsids.sourceforge.net/ 
 * Currently, this implementation only supports HTTP GET/POST resolution (resolution
 * for SOAP-based LSID servers is in the works).  Essentially, given a URN such as
 * <code>urn:lsid:biomoby.org:servicetype:Retrieval:2001-09-21T16-00-00Z</code>, the
 * resolvcer takes these steps:
 * <ol>
 *  <li>Parses out the authority "biomoby.org" from the URN</li>
 *  <li>Uses DNS SRV records to determine if the authority has a LSID 
 *      resolution server (e.g. moby.ucalgary.ca, port 80 for biomoby.org)</li>
 *  <li>Does an HTTP GET for /authority/ on the server to get the WSDL for the getAvailableServices operation</li>
 *  <li>Calls getAvailableServices on the LSID, which returns a WSDL of things you can do with the LSID</li>
 *  <li>Determines from the WSDL how to retrieve the data or metadata for the LSID</li>
 *  <li>Resolves any lsid:latest predicate in the response (if the requested LSID was versionless)</li>
 *  <li>Returns the URL where the data or metadata can be retrieved</li>
 * </ol>
 *
 * For the moment, the First Well-Known Rule of LSID resolution is ignored because 
 * lsid.urn.arpa is not registered in DNS, and all known hosts have DNS SRV records.  
 * Steps 1 and 2 described here implement the Second Well-Known Rule in the LSID DDDS process (RFC 3405).
 */
public class LSIDResolver{
    public static final String SRV_PREFIX = "_lsid._tcp.";
    public static final String AUTHORITY_WSDL_LOCATION = "/authority/";
    public static final String AUTHORITY_SERVICE_HTTP_BINDING = "LSIDAuthorityHTTPBinding";
    public static final String AUTHORITY_SERVICE_SOAP_BINDING = "LSIDAuthoritySOAPBinding";
    public static final String LSID_DATA_NAMESPACE = "http://www.omg.org/LSID/2003/DataServiceSOAPBindings";
    public static final String LSID_HTTP_NAMESPACE = "http://www.omg.org/LSID/2003/AuthorityServiceHTTPBindings";
    public static final String LSID_SOAP_NAMESPACE = "http://www.omg.org/LSID/2003/AuthorityServiceSOAPBindings";
    public static final String LSID_WSDL_NAMESPACE = "http://www.omg.org/LSID/2003/Standard/WSDL";
    private boolean warnedOfDNSAccess = false;

    private static Logger logger = Logger.getLogger(LSIDResolver.class.getName());    

    private XPath xPath;
    private DocumentBuilder docBuilder;

    public LSIDResolver(){
        //PG Temporarily use xalan while Google App Engine XPath bug exists
        //XPathFactory xPathFactory = XPathFactory.newInstance();
        XPathFactory xPathFactory = new org.apache.xpath.jaxp.XPathFactoryImpl();
	try{
	    xPath = xPathFactory.newXPath();
	    xPath.setNamespaceContext(new NamespaceContextImpl());
	} catch(Exception e){
            e.printStackTrace();
        } 

        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        try{
            docBuilder = dbf.newDocumentBuilder();
        } catch(Exception e){
            e.printStackTrace();
        }
    }

    public static boolean isLSID(String id){
	return id.matches("urn:lsid:(\\S+?):.*:.*");
    }

    public URL findAuthorityWSDL(String lsid) throws Exception{
	String lsidAuthority = lsid.replaceFirst("^urn:lsid:(\\S+?):.*:.*$", "$1");
	if(lsidAuthority == null || lsidAuthority.equals(lsid)){
	    throw new Exception("The LSID to resolve (" + lsid + 
				") does not seem to be a proper LSID syntax, " +
				"cannot extract the lookup authority.");
	}

	Record[] records = null;
	try{
	    records = new Lookup(SRV_PREFIX + lsidAuthority, Type.SRV).run();
	} catch(Throwable t){
	    // If we can't access DNA resources (e.g. in a restricted environment 
	    // such as an applet or Google App Engine), do a last ditch guess
	    // as to the location of the server rather than using the second well-known rule
	    if(!warnedOfDNSAccess){
		warnedOfDNSAccess = true;
		logger.log(Level.WARNING, "Cannot access DNS SRV records for LSID resolution, will use hostname directly", t);
	    }
	    return new URL("http", 
			   lsidAuthority, 
			   80, 
			   AUTHORITY_WSDL_LOCATION);
	}
	if(records == null || records.length == 0){
	    throw new Exception("While trying to use LSID resolution's Second Well-Known Rule: " +
				" The authority for the lsid (" + lsidAuthority + 
				") did not have a SRV record in the DNS for " + SRV_PREFIX +  
				lsidAuthority + " as expected (see RFC 2782).  Either the lsid authority" +
				" is incorrect, or the authority no longer advertises LSID resolution."); 
	}
	Map resolverRecords = new TreeMap();
	// TODO: handle priority and weight values from the SRV records for load balancing/redundancy
	for(Record record: records){
	    if(!(record instanceof SRVRecord)){
		// They give you ARecords (Inet Addresses) too, even though we asked for SRVRecord...
		continue;
	    }
	    SRVRecord srvRecord = (SRVRecord) record;
	    String hostName = srvRecord.getTarget().toString();
	    // Sometimes there's an extraneous . at end of the hostname
	    if(hostName.charAt(hostName.length()-1) == '.'){ 
		hostName = hostName.substring(0, hostName.length()-1);
	    }

	    return new URL("http", 
			   hostName, 
			   srvRecord.getPort(), 
			   AUTHORITY_WSDL_LOCATION);
	}
	throw new Exception("While trying to use LSID resolution's Second Well-Known Rule: " +
			    " The authority for the lsid (" + lsidAuthority + 
			    ") did not have a SRV record in the DNS for " + SRV_PREFIX +  
			    lsidAuthority + " as expected (see RFC 2782).  Either the lsid authority" +
			    " is incorrect, or the authority no longer advertises LSID resolution."); 
    }

    // HTTP binding
    public URL resolveDataURL(String lsid) throws Exception{
	return resolveURL(lsid, "LSIDDataHTTPBinding");
    }

    public URL resolveMetadataURL(String lsid) throws Exception{
	return resolveURL(lsid, "LSIDMetadataHTTPBinding");
    }

    protected URL resolveURL(String lsid, String bindingName) throws Exception{
	URL authorityWsdlUrl = findAuthorityWSDL(lsid);

	URL resultURL = null;
	URL getAvailableServicesURLHTTP = null;
	String getAvailableServicesURLHTTPVerb = "GET";
	
	Document wsdlDoc = docBuilder.parse(authorityWsdlUrl.openStream());
	Element[] es = findGetAvailableServices(wsdlDoc, authorityWsdlUrl);
	if(es == null || es.length == 0){
	    throw new Exception("Could not find any definition of the " +
				"getAvailableServices operation in the WSDL ("+
				authorityWsdlUrl+")");
	}
	for(Element e: es){
	    NodeList ops = e.getElementsByTagNameNS(MobyPrefixResolver.HTTP_NAMESPACE, "operation");
	    if(ops.getLength() == 1){
		getAvailableServicesURLHTTP = new URL(authorityWsdlUrl, 
						      ((Element) ops.item(0)).getAttribute("location"));
		break;
	    }
	}
	if(getAvailableServicesURLHTTP != null){
	    //System.err.println("HTTP for getAvailableServices: " + getAvailableServicesURLHTTP);
	    URL servicesURL = null;
	    if("GET".equals(getAvailableServicesURLHTTPVerb)){
		servicesURL = new URL(getAvailableServicesURLHTTP, "?lsid="+lsid);
	    }
	    else if("POST".equals(getAvailableServicesURLHTTPVerb)){
		// TODO
	    }
	    else{
		throw new Exception("Method (" + getAvailableServicesURLHTTPVerb + 
				    ") was neither POST nor GET in the WSDL for " +
				    "HTTP version of getAvailableServices");
	    }
	    Document servicesDoc = docBuilder.parse(servicesURL.openStream());
	    resultURL = findLSIDHTTPBindingAddr(lsid, bindingName, servicesDoc, servicesURL);
	}
	else{
	    throw new Exception("Could not find getAvailableServices operation using HTTP in the WSDL");
	}	    

	String[] lsidFields = lsid.split(":");
	if(lsidFields.length == 5){ // missing the version, see if it points us to another LSID...
	    resultURL = getLatestURL(resultURL, bindingName);
	}

	return resultURL;
    }

    private URL getLatestURL(URL currentURL, String bindingName) throws Exception{

	Document rdfDoc = null;
	try{
	    rdfDoc = docBuilder.parse(currentURL.openStream());
	} catch(Exception e){
	    // Can't parse the data, so don't try resolving the latest version info (it may not even be XML)
	    return currentURL;
	}

	NodeList latestElems = rdfDoc.getDocumentElement().
	    getElementsByTagNameNS(MobyPrefixResolver.LSID_NAMESPACE, "latest");
	if(latestElems.getLength() == 1){
	    String newLSID = ((Element) latestElems.item(0)).getTextContent();
	    if(newLSID == null){
		return currentURL;
	    }
	    String[] lsidFields = newLSID.split(":");
	    if(lsidFields.length == 6){
		// update the URL for the LSID
		currentURL = resolveURL(newLSID, bindingName);
	    }
	    else{
		throw new Exception("The latest LSID predicate (" + newLSID + 
				    ") did not resolve to an LSID with a version number");
	    }
	}
	else if(latestElems.getLength() > 1){
	    throw new Exception("More than one latest LSID predicate was found when resolving a versionless LSID");
	}
	// else keep the current URL, as we have no latest-version data
	return currentURL;
    }

    // WSDL
    public Dispatch resolveMetadataService(String lsid) throws Exception{
	URL authorityWsdlUrl = findAuthorityWSDL(lsid);

	QName getAvailableServicesSOAPPortName = null;
	QName getAvailableServicesSOAPServiceName = null;
	Service authorityService = Service.create(authorityWsdlUrl, getAvailableServicesSOAPServiceName);
	Dispatch<Source> dispatch = authorityService.createDispatch(getAvailableServicesSOAPPortName,
								    Source.class,
								    Service.Mode.PAYLOAD);
	
	Source response = dispatch.invoke(createSource(dispatch,
						       new QName(LSID_DATA_NAMESPACE,
								 "getAvailableServices"),
						       lsid));
	// 	    System.err.println(getResponseData(response));
	return null;
    }


    // Because the JAX-WS interface needs the QName of the service in order to parse the WSDL, 
    // we need to extract the names from the WSDL doc first (we'll use XPath rather than 
    // building a whole WSDL model ourselves)
    private QName[] getServiceQNamesFromWSDL(Document wsdlDoc) throws Exception{
	NodeList targetNamespaceAttrs = (NodeList) xPath.evaluate("//wsdl:definitions/@targetNamespace",
								  wsdlDoc,
								  XPathConstants.NODESET);
	if(targetNamespaceAttrs.getLength() == 0){
	    throw new Exception("Could not find the target namespace in the WSDL file");
	}
	if(targetNamespaceAttrs.getLength() > 1){
	    throw new Exception("Found more than one (hence ambiguous) target namespace in the WSDL file");
	}
	if(!(targetNamespaceAttrs.item(0) instanceof Attr)){
	    throw new Exception("Result from target namespace attr XPath was of the wrong class (" +
				targetNamespaceAttrs.item(0).getClass().getName()+"), which is not a W3C Attr");
	}
	String targetNamespace = ((Attr) targetNamespaceAttrs.item(0)).getValue();

	NodeList serviceNameAttrs = (NodeList) xPath.evaluate("//wsdl:service/@name",
							      wsdlDoc,
							      XPathConstants.NODESET);
	List<QName> serviceQNames = new Vector<QName>();
	if(serviceNameAttrs.getLength() == 0){
	    throw new Exception("Could not find any service attributes in the WSDL file");
	}
	for(int i = 0; i < serviceNameAttrs.getLength(); i++){
	    Node resultNode = serviceNameAttrs.item(i);
	    if(!(resultNode instanceof Attr)){
		System.err.println("Ignoring result from service name attr XPath results, wrong class (" +
				   resultNode.getClass().getName()+"), which is not a W3C Attr");
	    }
	    serviceQNames.add(new QName(targetNamespace, ((Attr) resultNode).getValue()));
	}
	return serviceQNames.toArray(new QName[serviceQNames.size()]);
    }

    // Recursively find the getAvaialbleServices operation
    private Element[] findGetAvailableServices(Document wsdlDoc, URL baseURL) throws Exception{
	return findXPath("//wsdl:operation[@name=\"getAvailableServices\"]", wsdlDoc, baseURL);
    }

    private URL findLSIDHTTPBindingAddr(String lsid, String bindingName, Document wsdlDoc, URL baseURL) throws Exception{
	// The following XPath is not 100% correct, but will work in 99% of cases, i.e. those without local name collisions
	Element[] es = findXPath("//wsdl:port[contains(@binding,\":"+bindingName+"\")]", wsdlDoc, baseURL);
	if(es == null || es.length == 0){
	    throw new Exception("Could not find the definition of the  in the WSDL"+
				"(maybe no metadata is available for this LSID?)");
	}
	NodeList addrs = es[0].getElementsByTagNameNS(MobyPrefixResolver.HTTP_NAMESPACE, "address");
	if(addrs == null || addrs.getLength() == 0){
	    throw new Exception("Could not find the definition of the HTTP address within " +
				"the WSDL's LSIDMetadataHTTPBinding spec, maybe a non-existant LSID? (" + lsid + ")");
	}
	else{
	    // baseURL is a bad assumption, as the node may be from an import in another location
	    URL metaURL = new URL(baseURL, ((Element) addrs.item(0)).getAttribute("location"));
	    //System.err.println("Meta URL is " + metaURL);
	    return new URL(metaURL.getProtocol()+"://"+metaURL.getHost()+
			   metaURL.getPath()+"?lsid="+lsid); //assumes GET for now
	}
    }
 
    private Element[] findXPath(String xPathStr, Document wsdlDoc, URL baseURL) throws Exception{
	// Is the service defined in the main doc?
	NodeList availServiceElems = (NodeList) xPath.evaluate(xPathStr,
							       wsdlDoc,
							       XPathConstants.NODESET);
	if(availServiceElems.getLength() == 0){
	    NodeList importElems = (NodeList) xPath.evaluate("(//wsdl:import | //import)/@location",
							     wsdlDoc,
							     XPathConstants.NODESET);
	    // do the import and see if it's defined there
	    for(int i = 0; i < importElems.getLength(); i++){
		URL importURL = new URL(baseURL, ((Attr) importElems.item(i)).getValue());
		Document importWsdlDoc = docBuilder.parse(importURL.openStream());
		Element[] op = findGetAvailableServices(importWsdlDoc, importURL);
		if(op != null){
		    return op; // found the operation element in an import
		}
	    }
	    return null; // fail recursion terminal: not in this doc or its imports
	}
	else if(availServiceElems.getLength() > 1){
	    throw new Exception("The operation getAvailableServices is defined more than once in the WSDL");
	}
	else if(!(availServiceElems.item(0) instanceof Element)){
	    throw new Exception("The XPath to retrieve the operation getAvailableServices returned " +
				"something other than a W3C Element (bad XPath?): " + 
				availServiceElems.item(0).getClass().getName());
	}
	else{  // success recursion terminal: found it properly if we got to here
	    return new Element[]{(Element) availServiceElems.item(0)};
	}
    }

    private Source createSource(Dispatch<Source> dispatch, QName operationName, String lsid){
	Map<String,Object> context = dispatch.getRequestContext();
	context.put(Dispatch.SOAPACTION_USE_PROPERTY, Boolean.TRUE);
	context.put(Dispatch.SOAPACTION_URI_PROPERTY, 
		    operationName.getNamespaceURI()+"#"+operationName.getLocalPart());

	return new StreamSource(new StringReader("<"+operationName.getLocalPart()+
						 " xmlns=\""+operationName.getNamespaceURI()+
						 "\"><lsid>"+lsid+
						 "</lsid></"+operationName.getLocalPart()+">"));
    }

    private String getResponseData(Source source) throws Exception{
	StringWriter stringWriter = new StringWriter();
	TransformerFactory transformerFactory = TransformerFactory.newInstance();
	Transformer nullTransformer = transformerFactory.newTransformer();
	nullTransformer.transform(source, new StreamResult(stringWriter));
	return stringWriter.toString();
    }
}