from multiprocessing import Pipe
from multiprocessing.connection import Connection

from typing import TypeVar, Generic, Type, Optional, Tuple, overload

SendT = TypeVar("SendT")
RecvT = TypeVar("RecvT")
T = TypeVar("T")

class LimitedPipe(Generic[SendT, RecvT]):
    def __init__(self, send_type: Optional[Type[SendT]], recv_type: Optional[Type[RecvT]], connection: Connection):            
        self._connection = connection

    def send(self, data: SendT):
        self._connection.send(data)

    def recv(self) -> RecvT:
        return self._connection.recv()
    
    def close(self) -> None:
        self._connection.close()

    def poll(self, timeout: Optional[float] = None):
        return self._connection.poll(timeout=timeout)

    @property
    def closed(self) -> bool:
        return self._connection.closed

@overload
def make_limited_pipes(
    send_type: Type[SendT], 
) -> Tuple[LimitedPipe[SendT, None], LimitedPipe[None, SendT]]: ...

@overload
def make_limited_pipes(
    send_type: Type[SendT], recv_type: Type[RecvT]
) -> Tuple[LimitedPipe[SendT, RecvT], LimitedPipe[RecvT, SendT]]: ...


def make_limited_pipes(
    send_type: Type[SendT], recv_type: Optional[Type[RecvT]] = None,
) -> Tuple[LimitedPipe[SendT, RecvT], LimitedPipe[RecvT, SendT]]:
    if recv_type is None:
        c1, c2 = Pipe(duplex=False)

        # NOTE (swapnil) - based on multiprocessing.Pipe documentation, its the second tuple elem
        # that can send, not the first
        return LimitedPipe(send_type, recv_type, c2), LimitedPipe(recv_type, send_type, c1)

    c1, c2 = Pipe(duplex=True)        
    return LimitedPipe(send_type, recv_type, c2), LimitedPipe(recv_type, send_type, c1)





if __name__=="__main__":
    c1, c2 = make_limited_pipes(int, float)