"""Tests for tcp_proxy.py. Plain asserts, no pytest: `python3 test_tcp_proxy.py`.

Every test that needs a listener takes an ephemeral port, so nothing here can
ever collide with the fixed ports `demo.py` uses. Where a test needs two
events ordered it waits on a counter, exactly as the demo does -- there is no
`sleep()` in this file.
"""
import os
import socket
import threading

import tcp_proxy as T

KEYS = ["alpha", "beta", "gamma"]


def free_ephemeral():
    """A port the kernel just handed out, and immediately gave back."""
    probe = socket.socket()
    probe.bind((T.HOST, 0))
    port = probe.getsockname()[1]
    probe.close()
    return port


def load(px_port, n_clients, n_requests, timeout=2.0):
    """n_clients concurrent clients, n_requests each, released by a barrier."""
    got = [None] * n_clients
    barrier = threading.Barrier(n_clients)

    def one(i):
        sock = T.connect(px_port)
        lines = []
        barrier.wait()
        for r in range(n_requests):
            key = KEYS[(i + r) % len(KEYS)]
            sock.sendall(b"GET " + key.encode() + b"\n")
            lines.append((key, T.read_line(sock, timeout=timeout)))
        got[i] = lines
        sock.close()

    threads = [threading.Thread(target=one, args=(i,), daemon=True)
               for i in range(n_clients)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=20)
    ok = bad = 0
    for lines in got:
        for key, reply in lines or []:
            want = ("%s=%s\n" % (key, T.VALUES[key])).encode()
            ok, bad = (ok + 1, bad) if reply == want else (ok, bad + 1)
    return ok, bad


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

def test_upstream_answers_a_known_key():
    up = T.Upstream(free_ephemeral())
    sock = T.connect(up.port)
    sock.sendall(b"GET alpha\n")
    assert T.read_line(sock) == b"alpha=AAAAAAAAAAAA\n"
    sock.close(); up.close()


def test_upstream_errors_on_an_unknown_key():
    up = T.Upstream(free_ephemeral())
    sock = T.connect(up.port)
    sock.sendall(b"GET nope\n")
    assert T.read_line(sock) == b"ERR unknown key 'GET nope'\n"
    sock.close(); up.close()


def test_upstream_frames_on_the_delimiter_not_on_recv():
    """Two writes, one request: the frame is the "\\n", not the send."""
    up = T.Upstream(free_ephemeral())
    sock = T.connect(up.port)
    sock.sendall(b"GET al")
    up.wait_bytes(6)
    assert up.requests == []            # half a request is not a request
    sock.sendall(b"pha\n")
    assert T.read_line(sock) == b"alpha=AAAAAAAAAAAA\n"
    assert [line for _, line in up.requests] == [b"GET alpha"]
    sock.close(); up.close()


# --- the naive relay: correct, and wrong ------------------------------------

def test_naive_relay_forwards_every_byte_exactly_once():
    """The relay is not buggy. Byte conservation holds through it."""
    up = T.Upstream(free_ephemeral())
    px = T.NaiveProxy(free_ephemeral(), up.addr())
    a, b = T.connect(px.port), T.connect(px.port)
    a.sendall(b"GET alpha\n")
    px.wait_forwarded(10)
    b.sendall(b"GET beta\n")
    px.wait_forwarded(19)
    up.wait_requests(2)
    assert up.bytes_in == px.forwarded == 19
    assert [line for _, line in up.requests] == [b"GET alpha", b"GET beta"]
    a.close(); b.close(); px.close(); up.close()


def test_naive_hands_one_clients_reply_to_another():
    """The headline: no interleaving, no corruption, wrong recipient."""
    gate = threading.Event()
    up = T.Upstream(free_ephemeral(), gate=gate.wait)
    px = T.NaiveProxy(free_ephemeral(), up.addr(), routing="last")
    a, b = T.connect(px.port), T.connect(px.port)
    a.sendall(b"GET alpha\n")
    px.wait_forwarded(10)
    up.wait_requests(1)
    b.sendall(b"GET beta\n")
    px.wait_forwarded(19)               # B is now the last writer
    gate.set()
    up.wait_requests(2)
    assert T.drain(a) == b""
    assert T.drain(b) == b"alpha=AAAAAAAAAAAA\nbeta=BBBBBBBBBBBB\n"
    a.close(); b.close(); px.close(); up.close()


def test_naive_broadcast_routing_is_wrong_too():
    """The other defensible policy. There is no third one."""
    gate = threading.Event()
    up = T.Upstream(free_ephemeral(), gate=gate.wait)
    px = T.NaiveProxy(free_ephemeral(), up.addr(), routing="broadcast")
    a, b = T.connect(px.port), T.connect(px.port)
    a.sendall(b"GET alpha\n")
    px.wait_forwarded(10)
    up.wait_requests(1)
    b.sendall(b"GET beta\n")
    px.wait_forwarded(19)
    gate.set()
    up.wait_requests(2)
    both = b"alpha=AAAAAAAAAAAA\nbeta=BBBBBBBBBBBB\n"
    assert T.drain(a) == both           # A gets B's answer as well as its own
    assert T.drain(b) == both           # and B gets A's
    a.close(); b.close(); px.close(); up.close()


def test_naive_interleave_manufactures_requests_nobody_sent():
    up = T.Upstream(free_ephemeral())
    px = T.NaiveProxy(free_ephemeral(), up.addr())
    a, b = T.connect(px.port), T.connect(px.port)
    a.sendall(b"GET alp")
    up.wait_bytes(7)
    b.sendall(b"GET beta\n")
    up.wait_bytes(16)
    up.wait_requests(1)
    px.wait_routed(1)
    a.sendall(b"ha\n")
    up.wait_requests(2)
    assert [line for _, line in up.requests] == [b"GET alpGET beta", b"ha"]
    a.close(); b.close(); px.close(); up.close()


def test_naive_is_perfect_at_concurrency_one():
    """The boundary. Same relay, 8 clients, one request in flight."""
    up = T.Upstream(free_ephemeral())
    px = T.NaiveProxy(free_ephemeral(), up.addr())
    socks = [T.connect(px.port) for _ in range(8)]
    ok = bad = 0
    for r in range(3):
        for i, sock in enumerate(socks):
            key = KEYS[(i + r) % len(KEYS)]
            sock.sendall(b"GET " + key.encode() + b"\n")
            want = ("%s=%s\n" % (key, T.VALUES[key])).encode()
            ok, bad = ((ok + 1, bad) if T.read_line(sock) == want
                       else (ok, bad + 1))
    assert (ok, bad) == (24, 0)
    assert up.conns == 1                # 8 client conns : 1 upstream conn
    for sock in socks:
        sock.close()
    px.close(); up.close()


# --- the framed pool -------------------------------------------------------

def test_read_reply_stops_at_the_delimiter_and_keeps_the_rest():
    px = T.FramedProxy(free_ephemeral(), (T.HOST, 1))   # never dialled
    near, far = socket.socketpair()
    far.sendall(b"alpha=AAAAAAAAAAAA\nbeta=BBBBBBBBBBBB\n")
    reply, rest = px._read_reply(near, b"")
    assert reply == b"alpha=AAAAAAAAAAAA\n"
    assert rest == b"beta=BBBBBBBBBBBB\n"               # stays with the conn
    near.close(); far.close(); px.close()


def test_framed_pool_is_right_at_every_cap():
    for cap in (1, 2, 4, 8):
        up = T.Upstream(free_ephemeral())
        px = T.FramedProxy(free_ephemeral(), up.addr(), pool_cap=cap)
        assert load(px.port, 8, 3) == (24, 0), cap
        px.close(); up.close()


def test_pool_cap_above_the_concurrency_level_opens_nothing_extra():
    up = T.Upstream(free_ephemeral())
    up.gate = lambda: up.wait_requests(8)   # pin the count: hold 8 replies
    px = T.FramedProxy(free_ephemeral(), up.addr(), pool_cap=16)
    assert load(px.port, 8, 3) == (24, 0)
    assert up.conns == 8                    # not 16
    px.close(); up.close()


def test_releasing_on_a_timeout_breaks_when_the_upstream_pauses():
    """The load-bearing line, both ways round."""
    for release, want in (("framed", (2, 0)), ("timeout", (0, 2))):
        up = T.Upstream(free_ephemeral(), reply_delay=0.12, reply_split=6)
        px = T.FramedProxy(free_ephemeral(), up.addr(), pool_cap=1,
                           release=release, release_timeout=0.05)
        assert load(px.port, 2, 1) == want, release
        px.close(); up.close()


# --- the appendix: the 1:1 relay and the half-close ------------------------

def test_passthrough_opens_one_upstream_conn_per_client_conn():
    up = T.Upstream(free_ephemeral())
    px = T.PassthroughProxy(free_ephemeral(), up.addr())
    assert load(px.port, 8, 3) == (24, 0)
    assert up.conns == 8                    # no multiplexing whatsoever
    px.close(); up.close()


def test_half_close_passes_the_fin_on_and_the_reply_arrives():
    gate = threading.Event()
    up = T.Upstream(free_ephemeral(), gate=gate.wait)
    px = T.PassthroughProxy(free_ephemeral(), up.addr(), half_close=True)
    sock = T.connect(px.port)
    sock.sendall(b"GET alpha\n")
    up.wait_requests(1)
    sock.shutdown(socket.SHUT_WR)
    px.wait_eofs(1)
    gate.set()
    assert T.read_line(sock, timeout=1.0) == b"alpha=AAAAAAAAAAAA\n"
    sock.close(); px.close(); up.close()


def test_closing_on_the_fin_loses_the_whole_reply():
    gate = threading.Event()
    up = T.Upstream(free_ephemeral(), gate=gate.wait)
    px = T.PassthroughProxy(free_ephemeral(), up.addr(), half_close=False)
    sock = T.connect(px.port)
    sock.sendall(b"GET alpha\n")
    up.wait_requests(1)
    sock.shutdown(socket.SHUT_WR)
    px.wait_eofs(1)
    gate.set()
    assert T.read_line(sock, timeout=1.0) == b""
    sock.close(); px.close(); up.close()


# --- the claim the commentary makes about its own size ---------------------

def test_the_toy_is_over_the_repos_line_ceiling():
    """Pinned so the commentary's admission cannot quietly go stale."""
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                        "tcp_proxy.py")
    lines = open(path).read().splitlines()
    code = [l for l in lines if l.strip() and not l.strip().startswith("#")]
    assert (len(lines), len(code)) == (433, 357), (len(lines), len(code))
    assert len(lines) > 300


def main():
    tests = [value for name, value in sorted(globals().items())
             if name.startswith("test_")]
    for test in tests:
        test()
        print("ok   %s" % test.__name__)
    print("\n%d tests passed" % len(tests))


if __name__ == "__main__":
    main()