# this handler accepts a POST and responds with it

import tnetstring
import zmq


class Session(object):
    def __init__(self):
        self.id = None
        # 0 = receiving, 1 = responding
        self.state = 0
        self.sent_credits = False
        self.content_type = None
        self.data = ""
        self.offset = 0
        self.out_seq = 0
        self.credits = 0

    def take_data(self):
        left = len(self.data) - self.offset
        if self.credits >= left:
            take = left
        else:
            take = self.credits
        body = self.data[self.offset : self.offset + take]
        self.offset += take
        self.credits -= take
        return body


sessions = dict()

client_id = "zhttp-test"

ctx = zmq.Context()
in_sock = ctx.socket(zmq.PULL)
in_sock.connect("ipc:///tmp/zhttp-test-out")
in_stream_sock = ctx.socket(zmq.DEALER)
in_stream_sock.setsockopt(zmq.IDENTITY, client_id)
in_stream_sock.connect("ipc:///tmp/zhttp-test-out-stream")
out_sock = ctx.socket(zmq.PUB)
out_sock.connect("ipc:///tmp/zhttp-test-in")

poller = zmq.Poller()
poller.register(in_sock, zmq.POLLIN)
poller.register(in_stream_sock, zmq.POLLIN)

while True:
    socks = dict(poller.poll())
    if socks.get(in_sock) == zmq.POLLIN:
        m_raw = in_sock.recv()
        req = tnetstring.loads(m_raw[1:])
        print("IN %s" % req)

        id = req["id"]
        if id in sessions:
            print("session already active")
            continue

        s = Session()
        s.id = id
        headers = req.get("headers")
        if headers:
            for h in headers:
                if h[0].lower() == "content-type":
                    s.content_type = h[1]
                    break
        sessions[id] = s

    elif socks.get(in_stream_sock) == zmq.POLLIN:
        parts = in_stream_sock.recv_multipart()
        req = tnetstring.loads(parts[1][1:])
        print("IN stream %s" % req)

        s = sessions.get(id)
        if s is None:
            print("no such session")
            continue

    ptype = req.get("type")
    if ptype is None or (ptype is not None and ptype == "credit"):
        if "credits" in req:
            s.credits += req["credits"]

    if ptype is None:
        assert s.state == 0
        body = ""
        if "body" in req:
            body = req["body"]

        s.data += body

        if not req.get("more"):
            s.state = 1  # responding

            resp = dict()
            resp["from"] = client_id
            resp["id"] = s.id
            resp["seq"] = s.out_seq
            s.out_seq += 1
            resp["code"] = 200
            resp["reason"] = "OK"
            headers = list()
            if s.content_type:
                headers.append(["Content-Type", s.content_type])
            headers.append(["Content-Length", str(len(s.data))])
            resp["headers"] = headers
            resp["body"] = s.take_data()
            if s.offset < len(s.data):
                resp["more"] = True
            else:
                del sessions[s.id]

            print("OUT %s" % resp)
            out_sock.send(req["from"] + " T" + tnetstring.dumps(resp))
        else:
            # send credits
            resp = dict()
            resp["from"] = client_id
            resp["id"] = s.id
            resp["seq"] = s.out_seq
            s.out_seq += 1
            resp["type"] = "credit"
            if not s.sent_credits:
                resp["credits"] = 200000
                s.sent_credits = True
            else:
                resp["credits"] = len(body)

            print("OUT %s" % resp)
            out_sock.send(req["from"] + " T" + tnetstring.dumps(resp))
    elif ptype is not None and ptype == "credit":
        if s.state == 1 and s.credits > 0:
            resp = dict()
            resp["from"] = client_id
            resp["id"] = id
            resp["seq"] = s.out_seq
            s.out_seq += 1
            resp["body"] = s.take_data()
            if s.offset < len(s.data):
                resp["more"] = True
            else:
                del sessions[s.id]

            print("OUT %s" % resp)
            out_sock.send(req["from"] + " T" + tnetstring.dumps(resp))
