class WebSocketClient:
    def __init__(self, hostname, port):
        self.buffer = b""
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.hostname = hostname
        self.port = port
        self.connect()

    def __getattr__(self, method, *args, **kwargs):
        if method in self.__dict__.keys():
            return self.__dict__[method]

        def call(*args, **kwargs):
            self.write(json.dumps({"method": method, "args": args, "kwargs": kwargs}))
            return json.loads(self.read())

        return call

    def connect(self):
        self.socket.connect((self.hostname, self.port))
        key = base64.b64encode(b"1234123412341234").decode("utf-8")
        handshake = (
            f"GET /db HTTP/1.1\r\n"
            f"Host: localhost:3131\r\n"
            f"Upgrade: websocket\r\n"
            f"Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            f"Sec-WebSocket-Version: 13\r\n\r\n"
        )
        self.socket.sendall(handshake.encode("utf-8"))
        response = self.read_until(b"\r\n\r\n")
        if b"101 Switching Protocols" not in response:
            raise Exception("Failed to connect to WebSocket")

    def write(self, message):
        message_bytes = message.encode("utf-8")
        length = len(message_bytes)
        if length <= 125:
            self.socket.sendall(b"\x81" + bytes([length]) + message_bytes)
        elif length >= 126 and length <= 65535:
            self.socket.sendall(
                b"\x81" + bytes([126]) + length.to_bytes(2, "big") + message_bytes
            )
        else:
            self.socket.sendall(
                b"\x81" + bytes([127]) + length.to_bytes(8, "big") + message_bytes
            )

    def read_until(self, delimiter):
        while True:
            find_pos = self.buffer.find(delimiter)
            if find_pos != -1:
                data = self.buffer[: find_pos + 4]
                self.buffer = self.buffer[find_pos + 4 :]
                return data

            chunk = self.socket.recv(1024)
            if not chunk:
                return None
            self.buffer += chunk

    def read_exactly(self, length):
        while len(self.buffer) < length:
            chunk = self.socket.recv(length - len(self.buffer))
            if not chunk:
                return None
            self.buffer += chunk
        response = self.buffer[:length]
        self.buffer = self.buffer[length:]
        return response

    def read(self):
        frame = None
        frame = self.read_exactly(2)
        length = frame[1] & 127
        if length == 126:
            length = int.from_bytes(self.read_exactly(2), "big")
        elif length == 127:
            length = int.from_bytes(self.read_exactly(8), "big")
        message = self.read_exactly(length)
        return message

    def close(self):
        self.socket.close()
