import aiosched
import socket


# Pls notice me njsmith senpai :>


async def proxy_one_way(source: aiosched.socket.AsyncSocket, sink: aiosched.socket.AsyncSocket):
    while True:
        data = await source.receive(1024)
        if not data:
            await sink.shutdown(socket.SHUT_WR)
            break
        await sink.send_all(data)


async def proxy_two_way(a: aiosched.socket.AsyncSocket, b: aiosched.socket.AsyncSocket):
    async with aiosched.create_pool() as pool:
        await pool.spawn(proxy_one_way, a, b)
        await pool.spawn(proxy_one_way, b, a)


async def main():
    print("Starting two-way proxy server")
    async with aiosched.skip_after(10):
        a = aiosched.socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        b = aiosched.socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        await a.connect(("localhost", 12345))
        await b.connect(("localhost", 54321))
        async with a, b:
            await proxy_two_way(a, b)
    print("Two-way proxy shutting down")


try:
    aiosched.run(main)
except KeyboardInterrupt:
    print(f"Ctrl+C caught")