# this handler responds with a file

import sys
import os
import tnetstring
import zmq


class Session(object):
    def __init__(self):
        self.file = None
        self.out_seq = 0
        self.credits = 0


filename = sys.argv[1]

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

        if "type" in req:
            print("wrong packet type")
            continue

        s = Session()
        sessions[id] = s

        s.file = open(filename, "r")

        if "credits" in req:
            s.credits += req["credits"]

        body = ""
        eof = False
        while s.credits > 0:
            buf = s.file.read(s.credits)
            s.credits -= len(buf)
            body += buf
            if len(buf) == 0:
                eof = True
                break

        resp = dict()
        resp["from"] = client_id
        resp["id"] = req["id"]
        resp["seq"] = s.out_seq
        resp["code"] = 200
        resp["reason"] = "OK"
        headers = list()
        headers.append(["Content-Type", "text/plain"])
        headers.append(["Content-Length", str(os.stat(filename).st_size)])
        resp["headers"] = headers
        if len(body) > 0:
            resp["body"] = body
        if not eof:
            resp["more"] = True

        s.out_seq += 1

        if eof:
            del sessions[id]

        print("OUT %s" % resp)
        out_sock.send(req["from"] + " T" + tnetstring.dumps(resp))
    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)

        # we are only streaming output, so subsequent input messages must be credits
        if "type" not in req or req["type"] != "credit":
            print("wrong packet type")
            continue

        id = req["id"]

        s = sessions.get(id)
        if s is None:
            continue

        if "credits" in req:
            s.credits += req["credits"]

        body = ""
        eof = False
        while s.credits > 0:
            buf = s.file.read(s.credits)
            s.credits -= len(buf)
            body += buf
            if len(buf) == 0:
                eof = True
                break

        resp = dict()
        resp["from"] = client_id
        resp["id"] = id
        resp["seq"] = s.out_seq

        s.out_seq += 1

        if len(body) > 0:
            resp["body"] = body
        if not eof:
            resp["more"] = True

        if eof:
            del sessions[id]

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