import json, itertools
from twisted.internet import defer
from twisted.internet.defer import inlineCallbacks
from autobahn.twisted import websocket

class WSClient(websocket.WebSocketClientProtocol):
    def __init__(self):
        websocket.WebSocketClientProtocol.__init__(self)
        self.events = []
        self.errors = []
        self.d = None
        self.ping_counter = itertools.count(0)
    def onOpen(self):
        self.factory.d.callback(self)
    def onMessage(self, payload, isBinary):
        assert not isBinary
        event = json.loads(payload.decode("utf-8"))
        if event["type"] == "error":
            self.errors.append(event)
        if self.d:
            assert not self.events
            d,self.d = self.d,None
            d.callback(event)
            return
        self.events.append(event)

    def close(self):
        self.d = defer.Deferred()
        self.transport.loseConnection()
        return self.d
    def onClose(self, wasClean, code, reason):
        if self.d:
            self.d.callback((wasClean, code, reason))

    def next_event(self):
        assert not self.d
        if self.events:
            event = self.events.pop(0)
            return defer.succeed(event)
        self.d = defer.Deferred()
        return self.d

    @inlineCallbacks
    def next_non_ack(self):
        while True:
            m = yield self.next_event()
            if isinstance(m, tuple):
                print("unexpected onClose", m)
                raise AssertionError("unexpected onClose")
            if m["type"] != "ack":
                return m

    def strip_acks(self):
        self.events = [e for e in self.events if e["type"] != "ack"]

    def send(self, mtype, **kwargs):
        kwargs["type"] = mtype
        payload = json.dumps(kwargs).encode("utf-8")
        self.sendMessage(payload, False)

    def send_notype(self, **kwargs):
        payload = json.dumps(kwargs).encode("utf-8")
        self.sendMessage(payload, False)

    @inlineCallbacks
    def sync(self):
        ping = next(self.ping_counter)
        self.send("ping", ping=ping)
        # queue all messages until the pong, then put them back
        old_events = []
        while True:
            ev = yield self.next_event()
            if ev["type"] == "pong" and ev["pong"] == ping:
                self.events = old_events + self.events
                return None
            old_events.append(ev)

class WSFactory(websocket.WebSocketClientFactory):
    protocol = WSClient
