"""tcp-proxy: a byte-relay proxy and a request-aware pool, side by side.

    python3 demo.py            # the aha
    python3 test_tcp_proxy.py  # the tests

One line-protocol upstream, three proxies in front of it:

    NaiveProxy        many clients onto ONE shared upstream connection, as a
                      pure byte relay -- "multiplex many onto few" taken
                      literally. It misdelivers.
    FramedProxy       parse a request, borrow a pooled upstream connection,
                      read one response, hand the connection back. The parse
                      is what makes the pool possible at all.
    PassthroughProxy  one upstream connection per client connection: the
                      honest L4 relay, which multiplexes nothing.

Everything runs in one process over real loopback sockets, with threads
rather than an event loop, so the hooks the demo sequences on are ordinary
condition variables rather than a scheduler.
"""
import queue
import socket
import threading
import time

HOST = "127.0.0.1"

# The whole application protocol: "GET <key>\n" -> "<key>=<value>\n".
VALUES = {
    "alpha": "AAAAAAAAAAAA",
    "beta": "BBBBBBBBBBBB",
    "gamma": "GGGGGGGGGGGG",
}


def free_port(port):
    """True if nothing is bound to `port` on the loopback interface."""
    sock = socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        sock.bind((HOST, port))
        return True
    except OSError:
        return False
    finally:
        sock.close()


# ---------------------------------------------------------------------------
# the upstream


class Upstream:
    """The origin server: 'GET <key>\\n' in, '<key>=<value>\\n' out.

    It keeps the accounting the demo sequences on (bytes in, requests parsed,
    connections accepted) plus two counterfactual knobs: `gate` holds a reply
    until a condition is met, and `reply_split`/`reply_delay` cut a reply into
    two writes with a pause between them.
    """

    def __init__(self, port, gate=None, reply_delay=0.0, reply_split=0):
        self.port = port
        self.gate = gate or (lambda: None)
        self.reply_delay = reply_delay
        self.reply_split = reply_split
        self.cv = threading.Condition()
        self.bytes_in = 0
        self.requests = []              # [(conn_id, request_line_bytes)]
        self.conns = 0
        self.running = True
        self.lsock = socket.socket()
        self.lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.lsock.bind((HOST, port))
        self.lsock.listen(64)
        # listen() has returned before __init__ does, so a caller that
        # connects next cannot lose the race. No readiness probe needed.
        threading.Thread(target=self._accept, daemon=True).start()

    def addr(self):
        return (HOST, self.port)

    def _accept(self):
        while self.running:
            try:
                conn, _ = self.lsock.accept()
            except OSError:
                return
            with self.cv:
                self.conns += 1
                cid = self.conns
                self.cv.notify_all()
            threading.Thread(target=self._handle, args=(conn, cid),
                             daemon=True).start()

    def _handle(self, conn, cid):
        buf = b""
        while True:
            try:
                data = conn.recv(4096)
            except OSError:
                break
            if not data:
                break
            with self.cv:
                self.bytes_in += len(data)
                self.cv.notify_all()
            buf += data
            # The upstream frames on "\n" like everyone else has to. One
            # recv can carry half a request, or three of them.
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                self._respond(conn, cid, line)
        conn.close()

    def _respond(self, conn, cid, line):
        with self.cv:
            self.requests.append((cid, line))
            self.cv.notify_all()
        self.gate()
        text = line.decode("latin1")
        if text.startswith("GET ") and text[4:] in VALUES:
            reply = "%s=%s\n" % (text[4:], VALUES[text[4:]])
        else:
            reply = "ERR unknown key %r\n" % text
        raw = reply.encode()
        try:
            if self.reply_split:
                conn.sendall(raw[:self.reply_split])
                time.sleep(self.reply_delay)
                conn.sendall(raw[self.reply_split:])
            else:
                conn.sendall(raw)
        except OSError:
            pass

    # --- sequencing hooks: wait on a real event, never on a sleep ----------

    def wait_bytes(self, n):
        """Block until the upstream has received at least n bytes."""
        with self.cv:
            self.cv.wait_for(lambda: self.bytes_in >= n, timeout=5)

    def wait_requests(self, n):
        """Block until the upstream has parsed at least n whole requests."""
        with self.cv:
            self.cv.wait_for(lambda: len(self.requests) >= n, timeout=5)

    def close(self):
        self.running = False
        self.lsock.close()


# ---------------------------------------------------------------------------
# the proxies


class _Listener:
    """Accept loop shared by the three proxies. One thread per client."""

    def __init__(self, port):
        self.port = port
        self.running = True
        self.lsock = socket.socket()
        self.lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.lsock.bind((HOST, port))
        self.lsock.listen(64)
        threading.Thread(target=self._accept, daemon=True).start()

    def _accept(self):
        while self.running:
            try:
                conn, _ = self.lsock.accept()
            except OSError:
                return
            threading.Thread(target=self.handle, args=(conn,),
                             daemon=True).start()

    def close(self):
        self.running = False
        self.lsock.close()


class NaiveProxy(_Listener):
    """ONE upstream connection, shared by every client. A pure byte relay.

    "Multiplex many client connections onto few upstream connections" taken
    literally, with no idea what a protocol is. It forwards every byte, in
    order, exactly once, and it is wrong.
    """

    def __init__(self, port, upstream_addr, routing="last"):
        self.routing = routing          # "last" | "broadcast"
        self.clients = []
        self.last = None
        self.forwarded = 0              # bytes relayed upstream
        self.routed = 0                 # replies relayed downstream
        self.cv = threading.Condition()
        self.up = socket.create_connection(upstream_addr)
        threading.Thread(target=self._pump_upstream, daemon=True).start()
        _Listener.__init__(self, port)

    def handle(self, conn):
        self.clients.append(conn)
        while True:
            try:
                data = conn.recv(4096)
            except OSError:
                break
            if not data:
                break
            self.last = conn            # <-- the routing decision, such as it is
            self.up.sendall(data)
            with self.cv:
                self.forwarded += len(data)
                self.cv.notify_all()

    def _pump_upstream(self):
        while True:
            try:
                data = self.up.recv(4096)
            except OSError:
                return
            if not data:
                return
            if self.routing == "last":
                targets = [self.last] if self.last else []
            else:
                targets = list(self.clients)
            for target in targets:
                try:
                    target.sendall(data)
                except OSError:
                    pass
            with self.cv:
                self.routed += 1
                self.cv.notify_all()

    # --- sequencing hooks --------------------------------------------------

    def wait_forwarded(self, n):
        """Block until n bytes have been relayed upstream."""
        with self.cv:
            self.cv.wait_for(lambda: self.forwarded >= n, timeout=5)

    def wait_routed(self, n):
        """Block until n upstream replies have been relayed downstream."""
        with self.cv:
            self.cv.wait_for(lambda: self.routed >= n, timeout=5)


class FramedProxy(_Listener):
    """L7 pool: parse a request, borrow a connection, read one response,
    give the connection back. The parse IS the pooling.

    `release` picks how the borrower decides the response ended: "framed"
    reads to the delimiter, "timeout" reads until the upstream has been quiet
    for `release_timeout` seconds -- the only framing-free way to guess.
    """

    def __init__(self, port, upstream_addr, pool_cap=2, release="framed",
                 release_timeout=0.05):
        self.upstream_addr = upstream_addr
        self.pool_cap = pool_cap
        self.release = release          # "framed" | "timeout"
        self.release_timeout = release_timeout
        self.free = queue.Queue()       # of (socket, leftover_bytes)
        self.made = 0
        self.made_lock = threading.Lock()
        self.blocked_gets = 0           # times a borrower had to wait
        _Listener.__init__(self, port)

    def _checkout(self):
        try:
            return self.free.get_nowait()
        except queue.Empty:
            pass
        with self.made_lock:
            if self.made < self.pool_cap:
                self.made += 1
                return (socket.create_connection(self.upstream_addr), b"")
        # At the cap and nothing free: this is head-of-line blocking, and it
        # is the price the pool charges for not opening a connection here.
        self.blocked_gets += 1
        return self.free.get()

    def handle(self, conn):
        buf = b""
        while True:
            while b"\n" not in buf:
                try:
                    data = conn.recv(4096)
                except OSError:
                    return
                if not data:
                    return
                buf += data
            line, buf = buf.split(b"\n", 1)
            sock, rest = self._checkout()
            try:
                sock.sendall(line + b"\n")
                reply, rest = self._read_reply(sock, rest)
            finally:
                self.free.put((sock, rest))
            try:
                conn.sendall(reply)
            except OSError:
                return

    def _read_reply(self, sock, rest):
        """Read exactly one reply off a pooled connection.

        The `timeout` branch is pooling without parsing: it returns when the
        upstream goes quiet, a guess about the upstream's scheduling. The
        `framed` branch returns at the delimiter, a fact about the message,
        and keeps any bytes past it with the connection.
        """
        if self.release == "timeout":
            sock.settimeout(self.release_timeout)
            out = rest
            try:
                while True:
                    data = sock.recv(4096)
                    if not data:
                        break
                    out += data
            except (socket.timeout, OSError):
                pass
            sock.settimeout(None)
            return out, b""
        while b"\n" not in rest:
            data = sock.recv(4096)
            if not data:
                return rest, b""
            rest += data
        line, rest = rest.split(b"\n", 1)
        return line + b"\n", rest


class PassthroughProxy(_Listener):
    """The honest L4 relay: one upstream connection per client connection.

    Appendix material: it teaches a second mechanism (the half-close) and is
    why this file is over the repo's 300-line ceiling. `half_close` picks
    what the relay does with the client's FIN -- pass it on, or tear the
    whole connection down.
    """

    def __init__(self, port, upstream_addr, half_close=True):
        self.upstream_addr = upstream_addr
        self.half_close = half_close
        self.eofs = 0                   # relay directions that have seen EOF
        self.cv = threading.Condition()
        _Listener.__init__(self, port)

    def handle(self, conn):
        up = socket.create_connection(self.upstream_addr)
        back = threading.Thread(target=self._relay, args=(up, conn),
                                daemon=True)
        back.start()
        self._relay(conn, up)
        back.join(timeout=2)
        conn.close()
        up.close()

    def _relay(self, src, dst):
        while True:
            try:
                data = src.recv(4096)
            except OSError:
                break
            if not data:
                break
            try:
                dst.sendall(data)
            except OSError:
                break
        if self.half_close:
            try:
                dst.shutdown(socket.SHUT_WR)   # "no more from this side"
            except OSError:
                pass
        else:
            try:
                dst.close()                    # ...and the reply dies with it
            except OSError:
                pass
        with self.cv:
            self.eofs += 1
            self.cv.notify_all()

    def wait_eofs(self, n):
        """Block until n relay directions have reacted to an EOF."""
        with self.cv:
            self.cv.wait_for(lambda: self.eofs >= n, timeout=5)


# ---------------------------------------------------------------------------
# client helpers


def connect(port):
    return socket.create_connection((HOST, port))


def drain(sock, timeout=0.4):
    """Everything the socket has to give, up to `timeout` of quiet."""
    sock.settimeout(timeout)
    out = b""
    try:
        while True:
            data = sock.recv(4096)
            if not data:
                break
            out += data
    except (socket.timeout, OSError):
        pass
    return out


def read_line(sock, timeout=2.0):
    """One framed reply, or whatever arrived before `timeout`."""
    sock.settimeout(timeout)
    out = b""
    try:
        while b"\n" not in out:
            data = sock.recv(4096)
            if not data:
                break
            out += data
    except (socket.timeout, OSError):
        pass
    return out