import struct
from nacl.secret import SecretBox

from .util import bytes_to_long, long_to_bytes

NONCE_SIZE = 24
HEADER_LENGTH = 2 + 16 + 16
MAX_SEGMENT_SIZE = 4 * 1024
TERMINATION_HEADER = (b'\x00' * 18)
MAX_NONCE = (8 * NONCE_SIZE)

# TODO: Implement handling of messages > 4k


def inc_nonce(nonce):
    num = bytes_to_long(nonce) + 1
    if num > 2 ** MAX_NONCE:
        num = 0
    bnum = long_to_bytes(num)
    bnum = b'\x00' * (NONCE_SIZE - len(bnum)) + bnum
    return bnum


def get_stream_pair(reader, writer, **kwargs):
    """Return a tuple with `(unbox_stream, box_stream)` (reader/writer).

    :return: (:class:`secret_handshake.boxstream.UnboxStream`,
              :class:`secret_handshake.boxstream.BoxStream`) """
    box_args = {
        'key': kwargs['encrypt_key'],
        'nonce': kwargs['encrypt_nonce'],
    }
    unbox_args = {
        'key': kwargs['decrypt_key'],
        'nonce': kwargs['decrypt_nonce'],
    }
    return UnboxStream(reader, **unbox_args), BoxStream(writer, **box_args)


class UnboxStream(object):
    def __init__(self, reader, key, nonce):
        self.reader = reader
        self.key = key
        self.nonce = nonce
        self.closed = False

    async def read(self):
        data = await self.reader.read(HEADER_LENGTH)

        if not data:
            self.closed = True
            return None

        box = SecretBox(self.key)

        header = box.decrypt(data, self.nonce)

        if header == TERMINATION_HEADER:
            self.closed = True
            return None

        length = struct.unpack('>H', header[:2])[0]
        mac = header[2:]

        data = await self.reader.read(length)

        body = box.decrypt(mac + data, inc_nonce(self.nonce))

        self.nonce = inc_nonce(inc_nonce(self.nonce))
        return body

    async def __aiter__(self):
        while True:
            data = await self.read()
            if data is None:
                return
            yield data


class BoxStream(object):
    def __init__(self, writer, key, nonce):
        self.writer = writer
        self.key = key
        self.box = SecretBox(self.key)
        self.nonce = nonce

    def write(self, data):

        # XXX: This nonce logic is almost for sure wrong

        body = self.box.encrypt(data, inc_nonce(self.nonce))[24:]
        header = struct.pack('>H', len(body) - 16) + body[:16]

        hdrbox = self.box.encrypt(header, self.nonce)[24:]
        self.writer.write(hdrbox)

        self.nonce = inc_nonce(inc_nonce(self.nonce))
        self.writer.write(body[16:])

    def close(self):
        self.writer.write(self.box.encrypt(b'\x00' * 18, self.nonce)[24:])