import socket
import threading
import time

class P2PTransport:
    def __init__(self, listen_port: int, peer_host: str, peer_port: int):
        """
        Args:
            listen_port: port to bind() and accept()
            peer_host:   host to dial()
            peer_port:   port to dial()
        """
        self.listen_port = listen_port
        self.peer_host = peer_host
        self.peer_port = peer_port

    def connect(self) -> (socket.socket, bool):
        """
        Race bind+listen vs. dial:
          - If dial succeeds first, return (sock, True)  # we are initiator
          - If accept succeeds first, return (sock, False) # we are responder
        """
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind(('0.0.0.0', self.listen_port))
        server.listen(1)

        result = {}
        event = threading.Event()
        lock = threading.Lock()

        def accept_thread():
            try:
                conn, _ = server.accept()
                with lock:
                    if not event.is_set():
                        result['sock'] = conn
                        result['initiator'] = False
                        event.set()
            except Exception:
                pass

        def dial_thread():
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1.0)
            while not event.is_set():
                try:
                    sock.connect((self.peer_host, self.peer_port))
                    with lock:
                        if not event.is_set():
                            result['sock'] = sock
                            result['initiator'] = True
                            event.set()
                    return
                except (ConnectionRefusedError, socket.timeout):
                    time.sleep(0.1)
                except Exception:
                    break

        t1 = threading.Thread(target=accept_thread, daemon=True)
        t2 = threading.Thread(target=dial_thread, daemon=True)
        t1.start()
        t2.start()

        event.wait()
        sock, initiator = result['sock'], result['initiator']
        # close the listening socket—we’ve got our P2P link
        server.close()
        # ensure this socket is in blocking mode (no lingering timeouts)
        sock.settimeout(None)
        return sock, initiator

    def send_packet(self, sock: socket.socket, data: bytes) -> None:
        """
        Send a 2-byte big-endian length prefix followed by data.
        """
        length = len(data).to_bytes(2, 'big')
        sock.sendall(length + data)

    def recv_packet(self, sock: socket.socket) -> bytes:
        """
        Receive a 2-byte length prefix, then that many payload bytes.
        """
        hdr = self._read_exact(sock, 2)
        length = int.from_bytes(hdr, 'big')
        return self._read_exact(sock, length)

    @staticmethod
    def _read_exact(sock: socket.socket, n: int) -> bytes:
        buf = bytearray()
        while len(buf) < n:
            chunk = sock.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("Socket closed during read")
            buf.extend(chunk)
        return bytes(buf)

    def close(self, sock: socket.socket) -> None:
        sock.close()
