import usocket as socket
import ubinascii
import uos

class WebSocketError(Exception):
    pass

class WebSocketClient:
    def __init__(self, uri, timeout=5):
        self.sock = None
        self.uri = uri
        self.timeout = timeout
        self.unread_messages = [] # Queue for buffered messages
        self.connect()

    def connect(self):
        uri = self.uri
        assert uri.startswith("ws://")

        uri = uri[5:]
        if "/" in uri:
            host, path = uri.split("/", 1)
        else:
            host, path = uri, ""
        path = "/" + path

        if ":" in host:
            host, port = host.split(":")
            port = int(port)
        else:
            port = 80

        print(f"Connecting to {host}:{port}{path}...")
        self.sock = socket.socket()
        
        # Add timeout
        self.sock.settimeout(self.timeout) 

        addr_info = socket.getaddrinfo(host, port)
        addr = addr_info[0][-1]
        print(f"Resolved address: {addr}")
        
        try:
            self.sock.connect(addr)
        except OSError as e:
            print(f"Socket connect failed: {e}")
            if e.args[0] == 113:
                print("Hint: Check firewall settings on server or if server is running.")
            raise

        # Random key
        key = ubinascii.b2a_base64(uos.urandom(16)).strip()

        
        req = "GET {} HTTP/1.1\r\n".format(path)
        req += "Host: {}:{}\r\n".format(host, port)
        req += "Connection: Upgrade\r\n"
        req += "Upgrade: websocket\r\n"
        req += "Sec-WebSocket-Key: {}\r\n".format(key.decode())
        req += "Sec-WebSocket-Version: 13\r\n"
        req += "\r\n"
        
        self.sock.write(req.encode())

        # Read handshake response
        header = b""
        while b"\r\n\r\n" not in header:
            chunk = self.sock.read(1)
            if not chunk:
                raise WebSocketError("Connection closed during handshake")
            header += chunk

        if b" 101 " not in header:
            raise WebSocketError("Handshake failed: " + header.decode())
        
        print("WebSocket connected!")

    def is_connected(self):
        return self.sock is not None

    def send(self, data, opcode=1): # 1=Text, 2=Binary
        if not self.sock:
            print("WebSocket is not connected (send called on closed socket)")
            raise WebSocketError("Connection closed")
            
        if isinstance(data, str):
            data = data.encode('utf-8')

        header = bytearray()
        header.append(0x80 | opcode) # FIN + Opcode

        length = len(data)
        if length < 126:
            header.append(0x80 | length) # Masked + length
        elif length < 65536:
            header.append(0x80 | 126)
            header.extend(length.to_bytes(2, 'big'))
        else:
            header.append(0x80 | 127)
            header.extend(length.to_bytes(8, 'big'))

        mask = uos.urandom(4)
        header.extend(mask)

        masked_data = bytearray(length)
        for i in range(length):
            masked_data[i] = data[i] ^ mask[i % 4]

        self.sock.write(header)
        self.sock.write(masked_data)

    def recv(self):
        # 1. Check if we have unread messages in the buffer
        if self.unread_messages:
            return self.unread_messages.pop(0)

        if not self.sock:
            return None

        # Read header
        try:
            # Read 2 bytes at once
            header = self.sock.read(2)
            if not header or len(header) < 2: return None
            
            b1 = header[0]
            b2 = header[1]
            
            fin = b1 & 0x80
            opcode = b1 & 0x0f
            
            mask = b2 & 0x80
            length = b2 & 0x7f

            if length == 126:
                length_bytes = self.sock.read(2)
                if not length_bytes: return None
                length = int.from_bytes(length_bytes, 'big')
            elif length == 127:
                length_bytes = self.sock.read(8)
                if not length_bytes: return None
                length = int.from_bytes(length_bytes, 'big')

            if mask:
                mask_key = self.sock.read(4)
                if not mask_key: return None

            # Read payload
            data = bytearray(length)
            view = memoryview(data)
            pos = 0
            while pos < length:
                read_len = self.sock.readinto(view[pos:])
                if read_len == 0: 
                    return None
                pos += read_len
            
            if mask:
                unmasked = bytearray(length)
                for i in range(length):
                    unmasked[i] = data[i] ^ mask_key[i % 4]
                data = unmasked

            if opcode == 1: # Text
                return data.decode('utf-8')
            elif opcode == 2: # Binary
                return data
            elif opcode == 8: # Close
                self.close()
                return None
            elif opcode == 9: # Ping
                self.send(data, opcode=10) # Pong
                return self.recv()
            
            return data
            
        except Exception as e:
            # Don't print timeout errors as they are expected in non-blocking polling
            if "ETIMEDOUT" not in str(e) and "110" not in str(e):
                print(f"WS Recv Error: {e}")
            return None

    def close(self):
        if self.sock:
            self.sock.close()
            self.sock = None
        self.unread_messages = []