"""
The :mod:`websockets.framing` module implements data framing as specified in
`section 5 of RFC 6455`_.

It deals with a single frame at a time. Anything that depends on the sequence
of frames is implemented in :mod:`websockets.protocol`.

.. _section 5 of RFC 6455: http://tools.ietf.org/html/rfc6455#section-5

"""

import asyncio
import collections
import io
import random
import struct

from .exceptions import PayloadTooBig, WebSocketProtocolError


try:
    from .speedups import apply_mask
except ImportError:  # pragma: no cover
    from .utils import apply_mask


__all__ = [
    'DATA_OPCODES',
    'CTRL_OPCODES',
    'OP_CONT',
    'OP_TEXT',
    'OP_BINARY',
    'OP_CLOSE',
    'OP_PING',
    'OP_PONG',
    'Frame',
    'encode_data',
    'parse_close',
    'serialize_close',
]

DATA_OPCODES = OP_CONT, OP_TEXT, OP_BINARY = 0x00, 0x01, 0x02
CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG = 0x08, 0x09, 0x0A

# Close code that are allowed in a close frame.
# Using a list optimizes `code in EXTERNAL_CLOSE_CODES`.
EXTERNAL_CLOSE_CODES = [1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011]

FrameData = collections.namedtuple(
    'FrameData', ['fin', 'opcode', 'data', 'rsv1', 'rsv2', 'rsv3']
)


class Frame(FrameData):
    """
    WebSocket frame.

    * ``fin`` is the FIN bit
    * ``rsv1`` is the RSV1 bit
    * ``rsv2`` is the RSV2 bit
    * ``rsv3`` is the RSV3 bit
    * ``opcode`` is the opcode
    * ``data`` is the payload data

    Only these fields are needed by higher level code. The MASK bit, payload
    length and masking-key are handled on the fly by :meth:`read` and
    :meth:`write`.

    """

    def __new__(cls, fin, opcode, data, rsv1=False, rsv2=False, rsv3=False):
        return FrameData.__new__(cls, fin, opcode, data, rsv1, rsv2, rsv3)

    @classmethod
    @asyncio.coroutine
    def read(cls, reader, *, mask, max_size=None, extensions=None):
        """
        Read a WebSocket frame and return a :class:`Frame` object.

        ``reader`` is a coroutine taking an integer argument and reading
        exactly this number of bytes, unless the end of file is reached.

        ``mask`` is a :class:`bool` telling whether the frame should be masked
        i.e. whether the read happens on the server side.

        If ``max_size`` is set and the payload exceeds this size in bytes,
        :exc:`~websockets.exceptions.PayloadTooBig` is raised.

        If ``extensions`` is provided, it's a list of classes with an
        ``decode()`` method that transform the frame and return a new frame.
        They are applied in reverse order.

        This function validates the frame before returning it and raises
        :exc:`~websockets.exceptions.WebSocketProtocolError` if it contains
        incorrect values.

        """
        # Read the header.
        data = yield from reader(2)
        head1, head2 = struct.unpack('!BB', data)

        # While not Pythonic, this is marginally faster than calling bool().
        fin = True if head1 & 0b10000000 else False
        rsv1 = True if head1 & 0b01000000 else False
        rsv2 = True if head1 & 0b00100000 else False
        rsv3 = True if head1 & 0b00010000 else False
        opcode = head1 & 0b00001111

        if (True if head2 & 0b10000000 else False) != mask:
            raise WebSocketProtocolError("Incorrect masking")

        length = head2 & 0b01111111
        if length == 126:
            data = yield from reader(2)
            length, = struct.unpack('!H', data)
        elif length == 127:
            data = yield from reader(8)
            length, = struct.unpack('!Q', data)
        if max_size is not None and length > max_size:
            raise PayloadTooBig(
                "Payload length exceeds size limit ({} > {} bytes)".format(
                    length, max_size
                )
            )
        if mask:
            mask_bits = yield from reader(4)

        # Read the data.
        data = yield from reader(length)
        if mask:
            data = apply_mask(data, mask_bits)

        frame = cls(fin, opcode, data, rsv1, rsv2, rsv3)

        if extensions is None:
            extensions = []
        for extension in reversed(extensions):
            frame = extension.decode(frame, max_size=max_size)

        frame.check()

        return frame

    def write(frame, writer, *, mask, extensions=None):
        """
        Write a WebSocket frame.

        ``frame`` is the :class:`Frame` object to write.

        ``writer`` is a function accepting bytes.

        ``mask`` is a :class:`bool` telling whether the frame should be masked
        i.e. whether the write happens on the client side.

        If ``extensions`` is provided, it's a list of classes with an
        ``encode()`` method that transform the frame and return a new frame.
        They are applied in order.

        This function validates the frame before sending it and raises
        :exc:`~websockets.exceptions.WebSocketProtocolError` if it contains
        incorrect values.

        """
        # The first parameter is called `frame` rather than `self`,
        # but it's the instance of class to which this method is bound.

        frame.check()

        if extensions is None:
            extensions = []
        for extension in extensions:
            frame = extension.encode(frame)

        output = io.BytesIO()

        # Prepare the header.
        head1 = (
            (0b10000000 if frame.fin else 0)
            | (0b01000000 if frame.rsv1 else 0)
            | (0b00100000 if frame.rsv2 else 0)
            | (0b00010000 if frame.rsv3 else 0)
            | frame.opcode
        )

        head2 = 0b10000000 if mask else 0

        length = len(frame.data)
        if length < 126:
            output.write(struct.pack('!BB', head1, head2 | length))
        elif length < 65536:
            output.write(struct.pack('!BBH', head1, head2 | 126, length))
        else:
            output.write(struct.pack('!BBQ', head1, head2 | 127, length))

        if mask:
            mask_bits = struct.pack('!I', random.getrandbits(32))
            output.write(mask_bits)

        # Prepare the data.
        if mask:
            data = apply_mask(frame.data, mask_bits)
        else:
            data = frame.data
        output.write(data)

        # Send the frame.

        # The frame is written in a single call to writer in order to prevent
        # TCP fragmentation. See #68 for details. This also makes it safe to
        # send frames concurrently from multiple coroutines.
        writer(output.getvalue())

    def check(frame):
        """
        Check that this frame contains acceptable values.

        Raise :exc:`~websockets.exceptions.WebSocketProtocolError` if this
        frame contains incorrect values.

        """
        # The first parameter is called `frame` rather than `self`,
        # but it's the instance of class to which this method is bound.

        if frame.rsv1 or frame.rsv2 or frame.rsv3:
            raise WebSocketProtocolError("Reserved bits must be 0")

        if frame.opcode in DATA_OPCODES:
            return
        elif frame.opcode in CTRL_OPCODES:
            if len(frame.data) > 125:
                raise WebSocketProtocolError("Control frame too long")
            if not frame.fin:
                raise WebSocketProtocolError("Fragmented control frame")
        else:
            raise WebSocketProtocolError("Invalid opcode: {}".format(frame.opcode))


def encode_data(data):
    """
    Helper that converts :class:`str` or :class:`bytes` to :class:`bytes`.

    :class:`str` are encoded with UTF-8.

    """
    # Expect str or bytes, return bytes.
    if isinstance(data, str):
        return data.encode('utf-8')
    elif isinstance(data, bytes):
        return data
    else:
        raise TypeError("data must be bytes or str")


def parse_close(data):
    """
    Parse the data in a close frame.

    Return ``(code, reason)`` when ``code`` is an :class:`int` and ``reason``
    a :class:`str`.

    Raise :exc:`~websockets.exceptions.WebSocketProtocolError` or
    :exc:`UnicodeDecodeError` if the data is invalid.

    """
    length = len(data)
    if length >= 2:
        code, = struct.unpack('!H', data[:2])
        check_close(code)
        reason = data[2:].decode('utf-8')
        return code, reason
    elif length == 0:
        return 1005, ''
    else:
        assert length == 1
        raise WebSocketProtocolError("Close frame too short")


def serialize_close(code, reason):
    """
    Serialize the data for a close frame.

    This is the reverse of :func:`parse_close`.

    """
    check_close(code)
    return struct.pack('!H', code) + reason.encode('utf-8')


def check_close(code):
    """
    Check the close code for a close frame.

    """
    if not (code in EXTERNAL_CLOSE_CODES or 3000 <= code < 5000):
        raise WebSocketProtocolError("Invalid status code")