import time
import weakref
from functools import partial as bind
from collections import deque

import numpy as np

from . import sockets


class Client:

  def __init__(
      self, address, identity=None, name='Client', ipv6=False, pings=10,
      maxage=120, maxinflight=16, errors=True, connect=False):
    if identity is None:
      identity = int(np.random.randint(2 ** 32))
    self.address = address
    self.identity = identity
    self.name = name
    self.maxinflight = maxinflight
    self.errors = errors
    self.resolved = None
    self.socket = sockets.ClientSocket(identity, ipv6, pings, maxage)
    self.futures = weakref.WeakValueDictionary()
    self.queue = deque()
    connect and self.connect()

  def __getattr__(self, name):
    if name.startswith('__'):
      raise AttributeError(name)
    try:
      return bind(self.call, name)
    except AttributeError:
      raise ValueError(name)

  def stats(self):
    return {
        'futures': len(self.futures),
        'inflight': len(self.queue),
    }

  def connect(self, retry=True, timeout=10):
    while True:
      self.resolved = self._resolve(self.address)
      self._print(f'Connecting to {self.resolved}')
      try:
        self.socket.connect(self.resolved, timeout)
        self._print('Connection established')
        return
      except sockets.ProtocolError as e:
        self._print(f'Ignoring unexpected message: {e}')
      except sockets.ConnectError:
        pass
      if retry:
        continue
      else:
        raise sockets.ConnectError

  def call(self, method, data):
    assert len(self.futures) < 1000, (
        f'Too many unresolved requests in client {self.name}.\n' +
        f'Futures: {len(self.futures)}\n' +
        f'Resolved: {sum([x.done() for x in self.futures.values()])}')
    if self.maxinflight:
      while sum(not x.done() for x in self.queue) >= self.maxinflight:
        self.queue[0].check()
        time.sleep(0.001)
    if self.errors:
      try:
        while self.queue[0].done():
          self.queue.popleft().result()
      except IndexError:
        pass
    assert isinstance(data, dict)
    data = {k: np.asarray(v) for k, v in data.items()}
    data = sockets.pack(data)
    rid = self.socket.send_call(method, data)
    future = Future(self._receive, rid)
    self.futures[rid] = future
    if self.errors or self.maxinflight:
      self.queue.append(future)
    return future

  def close(self):
    return self.socket.close()

  def _receive(self, rid, retry):
    while rid in self.futures and not self.futures[rid].done():
      result = self._listen()
      if result is None and not retry:
        return
      time.sleep(0.0001)

  def _listen(self):
    try:
      result = self.socket.receive()
      if result is not None:
        other, payload = result
        if other in self.futures:
          self.futures[other].set_result(sockets.unpack(payload))
      return result
    except sockets.NotAliveError:
      self._print('Server is not responding')
      raise
    except sockets.RemoteError as e:
      self._print(f'Received error response: {e.args[1]}')
      other = e.args[0]
      if other in self.futures:
        self.futures[other].set_error(sockets.RemoteError(e.args[1]))
    except sockets.ProtocolError as e:
      self._print(f'Ignoring unexpected message: {e}')

  def _resolve(self, address):
    protocol, address = address.split('://', 1)
    if address.startswith('/bns/'):
      assert self.ipv6, (address, self.ipv6)
      self._print(f'BNS address detected: {address}')
      from google3.chubby.python.public import pychubbyutil
      while True:
        try:
          address, port = pychubbyutil.ResolveBNSName(address)
          break
        except pychubbyutil.NoResultsError:
          self._print('BNS address not found, retrying')
          time.sleep(10)
      address = f'[{address}]:{port}'
      self._print(f'BNS address resolved to: {address}')
    return f'{protocol}://{address}'

  def _print(self, text):
    print(f'[{self.name}] {text}')


class Future:

  def __init__(self, waitfn, *args):
    self._waitfn = waitfn
    self._args = args
    self._status = 0
    self._result = None
    self._error = None

  def check(self):
    if self._status == 0:
      self._waitfn(*self._args, retry=False)

  def done(self):
    return self._status > 0

  def result(self):
    if self._status == 0:
      self._waitfn(*self._args, retry=True)
    if self._status == 1:
      return self._result
    elif self._status == 2:
      raise self._error
    else:
      assert False

  def set_result(self, result):
    self._status = 1
    self._result = result

  def set_error(self, error):
    self._status = 2
    self._error = error
