
import os
import threading
import sys
import logging


logger = logging.getLogger("wandb")


class Unbuffered(object):
   def __init__(self, stream):
       self.stream = stream
   def write(self, data):
       self.stream.write(data)
       self.stream.flush()
   def writelines(self, datas):
       self.stream.writelines(datas)
       self.stream.flush()
   def __getattr__(self, attr):
       return getattr(self.stream, attr)


def _pipe_relay(stopped, fd, name, cb, tee):
    while not stopped.isSet():
        logger.info("relay read")
        try:
            data = os.read(fd, 1000)
        except OSError as e:
            # TODO(jhr): handle this
            logger.info("got read error")
            return
        if len(data) == 0:
            logger.info("got no data")
            break
        if tee:
            logger.info("tee write")
            os.write(tee, data)
        #print("[%s]" % data, file=sys.stderr)
        if cb:
            logger.info("callback")
            cb(name, data)
            logger.info("callback done")
    logger.info("relay done")
    os.close(fd)
    logger.info("relay done done")


class Redirect(object):
    def __init__(self, src, dest, unbuffered=False, tee=False):
        self._installed = False
        self._stream = src
        self._dest = dest
        self._unbuffered = unbuffered
        self._tee = tee

        self._old_fd = None
        self._old_fp = None

    def _redirect(self, to_fd, unbuffered=False):
        fp = getattr(sys, self._stream)
        fp.close()
        os.dup2(to_fd, self._old_fd)
        setattr(sys, self._stream, os.fdopen(self._old_fd, 'w'))
        if unbuffered:
            setattr(sys, self._stream, Unbuffered(getattr(sys, self._stream)))

    def install(self):
        logger.info("install start")

        fp = getattr(sys, self._stream)
        fd = fp.fileno()
        old_fp = os.fdopen(os.dup(fd), 'w')

        if self._tee:
            self._dest._set_tee(old_fp.fileno())
        self._dest._start()
        self._installed = True

        self._old_fd = fd
        self._old_fp = old_fp

        self._redirect(to_fd=self._dest._get_writer(), unbuffered=self._unbuffered)
        logger.info("install stop")

    def uninstall(self):
        logger.info("uninstall start")
        self._dest._stop()
        self._redirect(to_fd=self._old_fp.fileno())
        logger.info("uninstall done")


class Capture(object):
    def __init__(self, name, cb):
        self._started = False
        self._name = name
        self._cb = cb
        self._stopped = None
        self._thread = None
        self._tee = None

        self._pipe_rd = None
        self._pipe_wr = None

    def _get_writer(self):
        assert self._started
        return self._pipe_wr

    def _set_tee(self, tee):
        self._tee = tee

    def _start(self):
        rd, wr = os.pipe()
        self._pipe_rd = rd
        self._pipe_wr = wr
        self._started = True

        self._stopped = threading.Event()
        # NB: daemon thread is used because we use atexit to determine when a user process is finished.  the atexit handler is responsible for flushing, joining, and closing
        read_thread = threading.Thread(name=self._name, target=_pipe_relay, daemon=True, args=(self._stopped, self._pipe_rd, self._name, self._cb, self._tee))
        read_thread.start()
        self._thread = read_thread

    def _stop(self):
        logger.info("_stop")
        self._stopped.set()
        os.close(self._pipe_wr)
        logger.info("_stop closed")
        # FIXME: need to shut this down cleanly since it is a daemon thread
        # self._thread.join()

@@ -82,24 +82,63 @@
 r(%22wandb%22)%0A%0A
+_LAST_WRITE_TOKEN = %22L@stWr!t3T0k3n%5Cn%22%0A
 %0Aclass Unbuf
@@ -525,62 +525,13 @@
 ile 
-not stopped.isSet():%0A        logger.info(%22relay read%22)
+True:
 %0A   
@@ -647,50 +647,8 @@
 his%0A
-            logger.info(%22got read error%22)%0A
     
@@ -705,384 +705,315 @@
     
-logger.info(%22got no data%22)%0A            break%0A        if tee:%0A            logger.info(%22tee write%22)%0A            os.write(tee, data)%0A        #print(%22%5B%25s%5D%22 %25 data, file=sys.stderr)%0A        if cb:%0A            logger.info(%22callback%22)%0A            cb(name, data)%0A            logger.info(%22callback done%22)%0A    logger.info(%22relay done%22)%0A    os.close(fd)%0A    logger.info(%22relay done done%22
+break%0A        if stopped.isSet():%0A            if data == _LAST_WRITE_TOKEN.encode():%0A                logger.info(%22relay done saw last write: %25s%22, name)%0A                break%0A        if tee:%0A            os.write(tee, data)%0A        if cb:%0A            cb(name, data)%0A    logger.info(%22relay done done: %25s%22, name
 )%0A%0A%0A
@@ -2167,35 +2167,8 @@
 t%22)%0A
-        self._dest._stop()%0A
     
@@ -2207,32 +2207,59 @@
 ld_fp.fileno())%0A
+        self._dest._stop()%0A
         logger.i
@@ -3284,32 +3284,58 @@
 ef _stop(self):%0A
+        name = self._name%0A
         logger.i
@@ -3344,18 +3344,29 @@
 o(%22_stop
-%22)
+: %25s%22, name)%0A
 %0A       
@@ -3382,24 +3382,84 @@
 opped.set()%0A
+        os.write(self._pipe_wr, _LAST_WRITE_TOKEN.encode())%0A
         os.c
@@ -3478,16 +3478,17 @@
 ipe_wr)%0A
+%0A
         
@@ -3512,17 +3512,27 @@
 p closed
-%22
+: %25s%22, name
 )%0A      
@@ -3609,18 +3609,16 @@
         
-# 
 self._th
@@ -3627,10 +3627,291 @@
 ad.join(
+timeout=30)%0A        if self._thread.isAlive():%0A            logger.error(%22Thread did not join: %25s%22, self._name)%0A            # TODO(jhr): do something better%0A        logger.info(%22_stop joined: %25s%22, name)%0A        os.close(self._pipe_rd)%0A        logger.info(%22_stop rd closed: %25s%22, name
 )%0A

make redirect flush all buffers