"""Glide HTTP Server
This module implements the functions necessary to turn Glide into a persistent
HTTP server that accepts ligands via POST requests and sends the poses back.
To use, just add the following lines to a Glide input file:
    CLIENT_MODULE schrodinger.application.glide.http_server
    CLIENT_OPTIONS "host=localhost; port=8000"
The server may then be tested using a web browser by opening
http://localhost:8000/. For programmatic access, see
schrodinger.application.glide.http_client.py.
The server responds to the following paths:
    /               a form that can be used for testing from a browser
    /shutdown       break out of the ligand loop and terminate
    /dock_ligand    POST a ligand and get the poses back
NOTE: the server is single-threaded, single-process, hence it's not designed to
accept concurrent connections. While Glide is busy docking a ligand, the
server won't be accepting connections. This server is meant for uses where
there is a single client that only needs to do one ligand at a time!
"""
import cgi
import http.server
import json
import re
import socket
import sys
import traceback
import urllib
from schrodinger import structure
from schrodinger.application.glide import http_client
from schrodinger.infra import mm
from schrodinger.job import jobcontrol
from schrodinger.utils import log
ENCODING = 'utf-8'
DEFAULT_HOST = 'localhost'
DEFAULT_PORT = 8000
# SUBJOBNAME is set by the C API, but we declare it here for documentation
# purposes and to keep flake8 happy.
SUBJOBNAME = None
logger = log.get_output_logger("http_server")
# GlideHTTPServer object
httpd = None
# Current reference ligand.
reflig_handle = mm.MMCT_INVALID_CT
DEFAULT_FORM = """
<!DOCTYPE HTML PUBLIC "-//W3C//DTD HTML 4.01 Transitional//EN"
    "http://www.w3.org/TR/html4/loose.dtd">
<html>
<head>
    <title>Glide HTTP server test form</title>
</head>
<body>
 <h1>Glide HTTP server test form</h1>
 <form method="post" action="/dock_ligand" enctype="multipart/form-data">
  <p>One-ligand file (uncompressed .mae): <input type="file" name="lig" /><br/>
  <input type="submit"></p>
 </form>
 <form method="get" action="/dock_smiles">
  <p>SMILES: <input name="smiles" />
  <input type="submit"></p>
 </form>
 <p><a href="/shutdown">[Shut down the server]</a></p>
</body>
</html>
"""
[docs]class GlideHTTPHandler(http.server.BaseHTTPRequestHandler):
    """This class, derived from BaseHTTPRequestHandler, implements
    the do_GET and do_POST methods. Unlike the parent class, this
    handler does not "finish" immediately after calling do_GET/do_POST,
    but waits until glide_finish() is called.
    Properties:
    * glide_data: a dictionary containing the posted form data.
    * glide_stop: a boolean, set to True if the client asked us to stop.
    * glide_form: the form to send out when getting "/".
    """
[docs]    def setup(self):
        self.glide_data = None
        self.glide_stop = False
        self.glide_form = DEFAULT_FORM
        http.server.BaseHTTPRequestHandler.setup(self) 
[docs]    def do_POST(self):
        if self.path in ('/dock_ligand', '/set_reflig'):
            fs = cgi.FieldStorage(self.rfile,
                                  headers=self.headers,
                                  environ={'REQUEST_METHOD': 'POST'})
            self.glide_data = {k: fs.getlist(k) for k in fs}
        else:
            self.send_response(404) 
[docs]    def do_GET(self):
        if self.path == '/':
            self.glide_send_response("text/html", self.glide_form)
        elif self.path == '/shutdown':
            self.glide_stop = True
            self.glide_send_response("text/plain", "Server is shutting down.\n")
        elif self.path.startswith('/dock_smiles?'):
            url_components = urllib.parse.urlparse(self.path)
            self.glide_data = urllib.parse.parse_qs(url_components.query)
        else:
            self.send_response(404) 
[docs]    def glide_send_response(self, ctype, body):
        """Convenience method to send the response line, content-type header,
        and body in just one call."""
        self.send_response(200)
        self.send_header("Content-Type", '%s; charset=%s' % (ctype, ENCODING))
        self.end_headers()
        self.wfile.write(body.encode(ENCODING)) 
[docs]    def finish(self):
        pass  # so we don't really finish handling the request immediately, but 
        # instead wait until glide_finish is called
[docs]    def glide_finish(self):
        """Finish the handler by calling the finish() method from the parent
        class. Among other things, this closes the connection."""
        http.server.BaseHTTPRequestHandler.finish(self)  
[docs]class GlideHTTPServer(http.server.HTTPServer):
    """This is a variant on HTTPServer that doesn't shut down requests
    immediately, but keeps them around until glide_shutdown_request is
    called. This allows us to split the processing of the request into
    two steps: one to get the request, and the other to respond to it.
    In the meantime, the handler object is kept around in the
    glide_http_handler property."""
[docs]    def finish_request(self, request, client_address):
        # This overrides the finish_request in the parent class. We don't
        # really "finish" the request here; we let the caller of
        # .handle_request() to take care of that by explicitly calling
        # glide_shutdown_request.
        # We'll keep the handler around so we can communicate with it
        self.glide_http_handler = self.RequestHandlerClass(
            request, client_address, self) 
[docs]    def shutdown_request(self, request):
        # save the request object; we won't actually shut down the request
        # here, but instead wait until glide_shutdown_request is called
        self.glide_request = request 
[docs]    def glide_shutdown_request(self):
        """Shut down the current request by calling the shutdown_request
        method from the parent class."""
        if self.glide_http_handler:
            self.glide_http_handler.glide_finish()
        if self.glide_request:
            http.server.HTTPServer.shutdown_request(self, self.glide_request) 
[docs]    def handle_request(self):
        self.glide_timeout = False
        self.glide_http_handler = None
        self.glide_request = None
        http.server.HTTPServer.handle_request(self) 
[docs]    def handle_timeout(self):
        self.glide_timeout = True  
[docs]def start(options):
    """Start the HTTP server. Takes a string as an argument that may specify
    the host and port as, for example, "host=localhost; port=8000; timeout=0".
    These are in fact the default values. To accept connections from remote
    hosts, set host to an empty string (i.e., "host="). If the timeout value
    is greater than zero, pull_ligand will return -1, indicating no more
    ligands, after waiting for that time in seconds."""
    global httpd
    host = DEFAULT_HOST
    port = DEFAULT_PORT
    m = re.search(r'\bhost *= *([\w.-]*)', options, re.I)
    if m:
        host = m.group(1)
    m = re.search(r'\bport *= *(\d+)', options, re.I)
    if m:
        port = int(m.group(1))
    httpd = GlideHTTPServer((host, port), GlideHTTPHandler)
    timeout = None
    m = re.search(r'\btimeout *= *(-?\d+)', options, re.I)
    if m:
        timeout = int(m.group(1))
        if timeout > 0:
            httpd.timeout = timeout 
[docs]def write_config():
    """
    Write a JSON file with host and port information so the client knows
    that the server is ready and where to connect. This is particularly
    needed when using automated port selection.
    When running under job control, the file is copied back to the launch
    directory immediately.
    """
    ip_addr, port = httpd.server_address
    logger.info("Listening on %s:%d" % (ip_addr, port))
    config_filename = '%s_http.json' % SUBJOBNAME
    backend = jobcontrol.get_backend()
    with open(config_filename, 'w') as fh:
        config = {'port': port, 'host': ip_addr}
        if backend:
            config['jobid'] = backend.getJob().JobId
        json.dump(config, fh, indent=2)
    if backend:
        backend.copyOutputFile(config_filename) 
[docs]def pull_ligand():
    """Wait until someone POSTs a ligand and return its mmct handle.
    If we were asked to shut down by the client, return -1."""
    try:
        return _pull_ligand()
    except:
        traceback.print_exc()
        raise
    finally:
        sys.stdout.flush()
        sys.stderr.flush() 
config_written = False
def _pull_ligand():
    global config_written
    if not config_written:
        # We write the config file the first time pull_ligand() is called and
        # not from the more obvious start() because the pull_ligand() call
        # indicates that Glide is actually ready to listen, and the presence of
        # the config file indicates the same to the client. (Due to historical
        # accident, start() is called by Glide before reading the grid files and
        # other initializations which may take several seconds).
        write_config()
        config_written = True
    # Wait until we get a ligand or a shutdown request or time out.
    while True:
        try:
            got_lig_to_dock = False
            # This only actually "takes" the request; the
            # response will be sent below or from _push_ligand()
            httpd.handle_request()
            if httpd.glide_timeout:
                logger.error("Glide server timed out (%d s)" % httpd.timeout)
                return -1
            handler = httpd.glide_http_handler
            if handler.glide_stop:
                return -1
            data = handler.glide_data
            if data is not None:
                try:
                    if 'lig' in data:
                        lig_str = data['lig'][0].decode('utf-8')
                        ct = next(
                            structure.MaestroReader("", input_string=lig_str))
                    elif 'smiles' in data:
                        smiles = data['smiles'][0]
                        ct = structure.SmilesStructure(smiles).get2dStructure()
                    else:
                        raise ValueError('Query had neither lig nor smiles')
                except (ValueError, KeyError) as e:
                    # problem parsing ligand; return HTTP error and try again
                    logger.error(e)
                    result = http_client.GlideResult([], str(e))
                    httpd.glide_http_handler.glide_send_response(
                        "application/json", result.toJson())
                    httpd.glide_shutdown_request()
                    continue
                # Call private method to release c++ refcounting and automated
                # deletion since glide will delete it.
                ct._cpp_structure.releaseOwnership()
                if handler.path == '/set_reflig':
                    global reflig_handle
                    reflig_handle = ct.handle
                    result = http_client.GlideResult([], "Updated reflig")
                    handler.glide_send_response("application/json",
                                                result.toJson())
                else:
                    got_lig_to_dock = True
                    return ct.handle
        finally:
            # The request is complete unless we got a ligand ct to dock.
            # If we did, the request will be completed in _push_ligand().
            if not got_lig_to_dock:
                httpd.glide_shutdown_request()
[docs]def push_ligand(pose_handles, msg):
    """
    Send the HTTP response as an m2io file of docked poses.
    Takes an array of mmct handles and an error message (the latter
    is currently unused.)
    :param pose_handles: mmct handles for docked poses
    :type: iterable of int
    :param msg: status message from Glide
    :type msg: str
    """
    try:
        return _push_ligand(pose_handles, msg)
    except:
        traceback.print_exc()
        raise
    finally:
        sys.stdout.flush()
        sys.stderr.flush() 
def _push_ligand(pose_handles, msg):
    poses = []
    for handle in pose_handles:
        st = structure.Structure(handle)
        #Call private method to release c++ refcounting and automated deletion
        #since glide will delete it
        st._cpp_structure.releaseOwnership()
        poses.append(st)
    try:
        result = http_client.GlideResult(poses, msg)
        body = result.toJson()
        httpd.glide_http_handler.glide_send_response("application/json", body)
        httpd.glide_shutdown_request()
    except socket.error as e:
        # we can get socket exceptions if the client died, the connection
        # timed out, etc. But that's not reason enough to have the server
        # die, so we'll just ignore it and put the burden on the client to
        # figure out what to do.
        logger.warning("push_ligand caught socket exception: %s", e)
[docs]def stop():
    """Delete the HTTP server object and stop listening."""
    global httpd
    httpd.server_close()
    httpd = None 
if __name__ == '__main__':
    # for testing as a standalone script, we'll just echo back the ligands
    # we get from the client (after converting from string into mmct
    # handles and back!
    start("")
    while True:
        cth = pull_ligand()
        logger.info('ct handle %s', cth)
        if cth < 0:
            break
        push_ligand([cth], "")
    stop()
[docs]def get_reflig():
    """
    Return the mmct handle to the new reference ligand (MMCT_INVALID_CT if
    there's no new reference ligand). The handle is then owned by the caller and
    the function call has the side effect of making this module forget the
    current handle.
    :returns: mmct handle
    :rtype: int
    """
    global reflig_handle
    retval = reflig_handle
    reflig_handle = mm.MMCT_INVALID_CT
    return retval