"""
Control functions.
"""
import os
import socket
import select
import json
import binascii
import sys

from . import utils

class Control(object):

    def __init__(
            self,
            path,
            bind=False,
            **kwargs):

        super(Control, self).__init__(**kwargs)

        dirname = os.path.dirname(path)
        if not os.path.exists(dirname) or not os.path.isdir(dirname):
            try:
                os.makedirs(dirname)
            except OSError:
                # Did we catch a race condition?
                if not os.path.exists(dirname) or not os.path.isdir(dirname):
                    raise

        self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        if bind:
            if os.path.exists(path):
                os.remove(path)
            utils.cleanup(os.remove, path)
            self._sock.bind(path)
            self._sock.listen(1024)
        else:
            self._sock.connect(path)

    def fd(self):
        return self._sock.fileno()

    def run(self, command, environment=None, cwd="/"):
        if environment is None:
            environment = os.environ

        fobj = self._sock.makefile()
        fobj.write("NOVM RUN\n")

        # Write the initial run command.
        json.dump({
            "command": command,
            "environment": [
                "%s=%s" % (k,v) for (k,v) in environment.items()
            ],
            "cwd": cwd
        }, fobj)
        fobj.flush()

        # Read the initial result.
        res = json.loads(fobj.readline())
        if not "pid" in res:
            raise Exception(res)

        try:
            # Poll and transform the event stream.
            # This will basically turn this process
            # into a proxy for the remote process.
            while True:
                to_read, _, _ = select.select([fobj, sys.stdin], [], [])

                if fobj in to_read:
                    obj = json.loads(fobj.readline())

                    if "stderr" in obj:
                        if obj["data"] is None:
                            if obj.get("stderr"):
                                sys.stderr.close()
                            else:
                                sys.stdout.close()
                        else:
                            data = binascii.a2b_base64(obj["data"])
                            if obj.get("stderr"):
                                sys.stderr.write(data)
                            else:
                                sys.stdout.write(data)

                    elif "exitcode" in obj:
                        sys.exit(obj["exitcode"])

                if sys.stdin in to_read:
                    data = os.read(sys.stdin.fileno(), 4096)
                    fobj.write(data)
                    fobj.flush()

        except IOError:
            # Socket eventually will be closed.
            # This is a clean exit scenario.
            sys.exit(0)

@@ -1612,24 +1612,78 @@
 ption(res)%0A%0A
+        # Our process exitcode.%0A        exitcode = 0%0A%0A
         try:
@@ -1837,16 +1837,57 @@
 rocess.%0A
+            read_set = %5Bfobj, sys.stdin%5D%0A
         
@@ -1948,33 +1948,24 @@
 .select(
-%5Bfobj, sys.stdin%5D
+read_set
 , %5B%5D, %5B%5D
@@ -2027,40 +2027,230 @@
     
-obj = json.loads(fobj.readline()
+data = fobj.readline()%0A%0A                    # Server has closed the socket?%0A                    if not data:%0A                        break%0A%0A                    # Decode the object.%0A                    obj = json.loads(data
 )%0A%0A 
@@ -2883,25 +2883,27 @@
         
-sys.exit(
+exitcode = 
 obj%5B%22exi
@@ -2909,17 +2909,16 @@
 itcode%22%5D
-)
 %0A%0A      
@@ -2944,32 +2944,33 @@
 din in to_read:%0A
+%0A
                 
@@ -3014,16 +3014,49 @@
 , 4096)%0A
+                    if data:%0A    
         
@@ -3096,32 +3096,36 @@
                 
+    
 fobj.flush()%0A%0A  
@@ -3112,32 +3112,162 @@
     fobj.flush()
+%0A                    else:%0A                        # We don't close the socket.%0A                        read_set.remove(sys.stdin)
 %0A%0A        except
@@ -3369,24 +3369,38 @@
 enario.%0A    
+        pass%0A%0A
         sys.
@@ -3404,11 +3404,18 @@
 ys.exit(
-0
+exitcode
 )%0A

Enhance control client.