"""
This module provides a Pythonic wrapper around the browser's
[WebSocket API](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket),
enabling two-way communication with WebSocket servers.

Use this for real-time applications:

- Pythonic interface to browser WebSockets.
- Automatic handling of async event handlers.
- Support for receiving text (`str`) and binary (`memoryview`) data.
- Support for sending text (`str`) and binary (`bytes` and `bytearray`) data.
- Compatible with Pyodide and MicroPython.
- Works in webworker contexts.
- Naming deliberately follows the JavaScript WebSocket API closely for
  familiarity.

See the Python docs for
[an explanation of memoryview](https://docs.python.org/3/library/stdtypes.html#memoryview).


```python
from pyscript import WebSocket


def on_open(event):
    print("Connected!")
    ws.send("Hello server")

def on_message(event):
    print(f"Received: {event.data}")

def on_close(event):
    print("Connection closed")

ws = WebSocket(url="ws://localhost:8080/")
ws.onopen = on_open
ws.onmessage = on_message
ws.onclose = on_close
```
"""

import js
from pyscript.ffi import create_proxy
from pyscript.util import as_bytearray, is_awaitable


def _attach_event_handler(websocket, handler_name, handler_function):
    """
    Given a `websocket`, and `handler_name`, attach the `handler_function`
    to the `WebSocket` instance, handling both synchronous and asynchronous
    handler functions.

    Creates a JavaScript proxy for the handler and wraps async handlers
    appropriately. Handles the `WebSocketEvent` wrapping for all handlers.
    """
    if is_awaitable(handler_function):

        async def async_wrapper(event):
            await handler_function(WebSocketEvent(event))

        wrapped_handler = create_proxy(async_wrapper)
    else:
        wrapped_handler = create_proxy(
            lambda event: handler_function(WebSocketEvent(event))
        )
    # Note: Direct assignment (websocket[handler_name]) fails in Pyodide.
    setattr(websocket, handler_name, wrapped_handler)


class WebSocketEvent:
    """
    A read-only wrapper for
    [WebSocket event objects](https://developer.mozilla.org/en-US/docs/Web/API/MessageEvent).

    This class wraps browser WebSocket events and provides convenient access
    to event properties. It handles the conversion of binary data from
    JavaScript typed arrays to Python bytes-like objects.

    The most commonly used property is `event.data`, which contains the
    message data for "message" events.

    ```python
    def on_message(event):  # The event is a WebSocketEvent instance.
        # For text messages.
        if isinstance(event.data, str):
            print(f"Text: {event.data}")
        else:
            # For binary messages.
            print(f"Binary: {len(event.data)} bytes")
    ```
    """

    def __init__(self, event):
        """
        Create a WebSocketEvent wrapper from an underlying JavaScript
        `event`.
        """
        self._event = event

    def __getattr__(self, attr):
        """
        Get an attribute `attr` from the underlying event object.

        Handles special conversion of binary data from JavaScript typed
        arrays to Python `memoryview` objects.
        """
        value = getattr(self._event, attr)
        if attr == "data" and not isinstance(value, str):
            if hasattr(value, "to_py"):
                # Pyodide - convert JavaScript typed array to Python.
                return value.to_py()
            else:
                # MicroPython - manually convert JS ArrayBuffer.
                return memoryview(as_bytearray(value))
        return value


class WebSocket:
    """
    This class provides a Python-friendly interface to WebSocket connections,
    handling communication with WebSocket servers. It supports both text and
    binary data transmission.

    Access the underlying WebSocket methods and properties directly if needed.
    However, the wrapper provides a more Pythonic API. If you need to work
    with the raw JavaScript WebSocket instance, you can access it via the
    `_js_websocket` attribute.

    Using textual (`str`) data:

    ```python
    from pyscript import WebSocket


    # Create WebSocket with handlers as arguments.
    def handle_message(event):
        print(f"Got: {event.data}")

    ws = WebSocket(
        url="ws://echo.websocket.org/",
        onmessage=handle_message
    )

    # Or assign handlers after creation.
    def handle_open(event):
        ws.send("Hello!")

    ws.onopen = handle_open
    ```

    Using binary (`memoryview`) data:

    ```python
    def handle_message(event):
        if isinstance(event.data, str):
            print(f"Text: {event.data}")
        else:
            # Binary data as memoryview.
            print(f"Binary: {len(event.data)} bytes")

    ws = WebSocket(url="ws://example.com/", onmessage=handle_message)

    # Send binary data.
    data = bytearray([0x01, 0x02, 0x03])
    ws.send(data)
    ```

    Read more about Python's
    [`memoryview` here](https://docs.python.org/3/library/stdtypes.html#memoryview).
    """

    # WebSocket ready state constants.
    CONNECTING = 0
    OPEN = 1
    CLOSING = 2
    CLOSED = 3

    def __init__(self, url, protocols=None, **handlers):
        """
        Create a new WebSocket connection from the given `url` (`ws://` or
        `wss://`). Optionally specify `protocols` (a string or a list of
        protocol strings) and event handlers (`onopen`, `onmessage`, etc.) as
        keyword arguments.

        These arguments and naming conventions mirror those of the
        [underlying JavaScript WebSocket API](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket)
        for familiarity.

        If you need access to the underlying JavaScript WebSocket instance,
        you can get it via the `_js_websocket` attribute.

        ```python
        # Basic connection.
        ws = WebSocket(url="ws://localhost:8080/")

        # With protocol.
        ws = WebSocket(
            url="wss://example.com/socket",
            protocols="chat"
        )

        # With handlers.
        ws = WebSocket(
            url="ws://localhost:8080/",
            onopen=lambda e: print("Connected"),
            onmessage=lambda e: print(e.data)
        )
        ```
        """
        # Create underlying JavaScript WebSocket.
        if protocols:
            js_websocket = js.WebSocket.new(url, protocols)
        else:
            js_websocket = js.WebSocket.new(url)
        # Set binary type to arraybuffer for easier Python handling.
        js_websocket.binaryType = "arraybuffer"
        # Store the underlying WebSocket.
        # Use object.__setattr__ to bypass our custom __setattr__.
        object.__setattr__(self, "_js_websocket", js_websocket)
        # Attach any event handlers passed as keyword arguments.
        for handler_name, handler in handlers.items():
            setattr(self, handler_name, handler)

    def __getattr__(self, attr):
        """
        Get an attribute `attr` from the underlying WebSocket.

        This allows transparent access to WebSocket properties like
        `readyState`, `url`, `bufferedAmount`, etc.
        """
        return getattr(self._js_websocket, attr)

    def __setattr__(self, attr, value):
        """
        Set an attribute `attr` on the WebSocket to the given `value`.

        Event handler attributes (`onopen`, `onmessage`, etc.) are specially
        handled to create proper proxies. Other attributes are set on the
        underlying WebSocket directly.
        """
        if attr in ["onclose", "onerror", "onmessage", "onopen"]:
            _attach_event_handler(self._js_websocket, attr, value)
        else:
            setattr(self._js_websocket, attr, value)

    def send(self, data):
        """
        Send `data` through the WebSocket.

        Accepts both text (`str`) and binary data (`bytes`, `bytearray`, etc.).
        Binary data is automatically converted to a JavaScript `Uint8Array`.

        ```python
        # Send text.
        ws.send("Hello server!")

        # Send binary.
        ws.send(bytes([1, 2, 3, 4]))
        ws.send(bytearray([5, 6, 7, 8]))
        ```

        !!! warning

            The WebSocket **must be in the OPEN state to send data**.
        """
        if isinstance(data, str):
            self._js_websocket.send(data)
        else:
            buffer = js.Uint8Array.new(len(data))
            for index, byte_value in enumerate(data):
                buffer[index] = byte_value
            self._js_websocket.send(buffer)

    def close(self, code=None, reason=None):
        """
        Close the WebSocket connection. Optionally specify a `code` (`int`)
        and a `reason` (`str`) for closing the connection.

        ```python
        # Normal close.
        ws.close()

        # Close with code and reason.
        ws.close(code=1000, reason="Task completed")
        ```

        Usage and values for `code` and `reasons`
        [are explained here](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close).
        """
        if code and reason:
            self._js_websocket.close(code, reason)
        elif code:
            self._js_websocket.close(code)
        else:
            self._js_websocket.close()
