from dotenv import load_dotenv, find_dotenv
import os
import zerorpc
import gevent
import time
from io import BytesIO
import base64
from PIL import Image

load_dotenv(find_dotenv())

class CortexRPC(object):
  def __init__(self, get_fn, set_fn, ready_fn):
    super().__init__() 
    self._get_fn = get_fn
    self._set_fn = set_fn
    self.working = False
    self.connect(ready_fn)

  def connect(self, ready_fn):
    print('Connecting to RPC on port {}'.format(os.getenv('RPC_PORT')))
    self.rpc = zerorpc.Server(self)
    self.rpc.bind('tcp://127.0.0.1:' + os.getenv('RPC_PORT'))
    gevent.spawn(self.connected)
    gevent.spawn(lambda: ready_fn(self))
    self.rpc.run()

  def connected(self):
    print('RPC running on port {}'.format(os.getenv('RPC_PORT')))
    print('Connecting to Relay on port {}'.format(os.getenv('RELAY_PORT')))
    self.relay = zerorpc.Client()
    self.relay.connect('tcp://127.0.0.1:' + os.getenv('RELAY_PORT'))
    val = None
    try:
      val = self.relay.connected('rpc')
      print('RPC connected...'.format(val))
    except zerorpc.TimeoutExpired:
      print('RPC timeout :(')
    return val
    # print('defer sending a file...')
    # gevent.sleep(1)
    # return self.get_last_frame()

  """
  def get_last_frame(self):
    if self.working:
      self.working = False
      return 'working'
    path = os.getenv('TEST_FRAMES_PATH')
    onlyfiles = sorted([f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))])
    self.working = True

    for f in onlyfiles:
      gevent.sleep(1/10)
      if not self.working:
        return
      output = BytesIO()
      im = Image.open(os.path.join(path, f))
      # im = im.convert('RGB')
      im.save(output, format='JPEG')
      frame = output.getvalue()
      output.close()
      self.send_frame(f.replace('png', 'jpg'), frame)

    self.working = False
    return 'ok'
  """

  def set_param(self, key, value):
    print('aloha. {}: {}'.format(key, value))
    if self._set_fn is not None:
      self._set_fn(key, value)
    else:
      print('No listener provided to CortexRPC')
    return 'ok'

  def get_params(self):
    print('Get params call received!')
    return self._get_fn()

  def send_pil_image(self, fn, im):
    output = BytesIO()
    # im = im.convert('RGB')
    im.save(output, format='JPEG')
    frame = output.getvalue()
    output.close()
    self.send_frame(fn.replace('png', 'jpg'), frame)

  def send_frame(self, fn, frame):
    if self.relay is not None:
      # print('Send frame ' + fn)
      self.relay.send_frame(fn, frame)
    else:
      print('Can\'t send frame, relay not ready...')

if __name__ == '__main__':
  def log_param(key, value):
    print('Log param. {}: {}'.format(key, value))

  # this will block
  rpc = CortexRPC(log_param)