"""Bridge only the host's loopback bonsai port across a networkless container."""

from __future__ import annotations

import argparse
import asyncio
from pathlib import Path


async def relay(reader, writer):
    try:
        while data := await reader.read(65536):
            writer.write(data)
            await writer.drain()
        if writer.can_write_eof():
            writer.write_eof()
    except (ConnectionError, OSError):
        pass


async def bridge(reader, writer, connect):
    upstream = None
    try:
        other, upstream = await connect()
        await asyncio.gather(relay(reader, upstream), relay(other, writer))
    except (ConnectionError, OSError):
        pass
    finally:
        writer.close()
        if upstream:
            upstream.close()


async def serve(path, command):
    if command:
        server = await asyncio.start_server(
            lambda r, w: bridge(r, w, lambda: asyncio.open_unix_connection(path)),
            "127.0.0.1",
            8899,
        )
        async with server:
            process = await asyncio.create_subprocess_exec(*command)
            return await process.wait()
    server = await asyncio.start_unix_server(
        lambda r, w: bridge(r, w, lambda: asyncio.open_connection("127.0.0.1", 8899)),
        path,
    )
    Path(path).chmod(0o666)
    async with server:
        await server.serve_forever()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("socket")
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    command = args.command
    if command[:1] == ["--"]:
        command = command[1:]
    raise SystemExit(asyncio.run(serve(args.socket, command)))


if __name__ == "__main__":
    main()