import ssl

import pytest
import trustme

import tonio
from tonio.net import socket
from tonio.net.ssl import SSLStream, open_ssl_over_tcp_stream, serve_ssl_over_tcp


_SIZE = 1024 * 1024


@pytest.fixture(scope='session')
def ssl_ca():
    return trustme.CA()


@pytest.fixture(scope='session')
def ssl_cert(ssl_ca):
    return ssl_ca.issue_server_cert('127.0.0.1')


@pytest.fixture(scope='function')
def ssl_server_ctx(ssl_cert):
    ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    ssl_cert.configure_cert(ctx)
    return ctx


@pytest.fixture(scope='function')
def ssl_client_ctx(ssl_ca):
    ctx = ssl.create_default_context()
    ssl_ca.configure_trust(ctx)
    return ctx


def _get_port():
    sock = socket.socket()

    with sock:
        yield sock.bind(('127.0.0.1', 0))
        name = sock.getsockname()
        return name[1]


def test_ssl_tcp_recv(run, ssl_server_ctx, ssl_client_ctx):
    def server():
        done = tonio.Event()
        res = []
        port = yield _get_port()

        def _server_handler(stream: SSLStream):
            buf = b''
            while len(buf) < _SIZE:
                buf += yield stream.receive_some()
            res.append(buf)
            done.set()

        with tonio.scope() as scope:
            scope.spawn(serve_ssl_over_tcp(_server_handler, host='127.0.0.1', port=port, ssl_context=ssl_server_ctx))
            scope.spawn(client(port))
            yield done.wait()
            scope.cancel()
        yield scope()

        return res[0]

    def client(port):
        yield tonio.sleep(0.5)
        stream: SSLStream = yield open_ssl_over_tcp_stream('127.0.0.1', port=port, ssl_context=ssl_client_ctx)
        yield stream.send_all(b'a' * _SIZE)

    data = run(server())
    assert data == b'a' * _SIZE


def test_streams_tcp_send(run, ssl_server_ctx, ssl_client_ctx):
    done = tonio.Event()
    state = {'data': b''}

    def server():
        port = yield _get_port()

        def _server_handler(stream: SSLStream):
            yield stream.send_all(b'a' * _SIZE)

        with tonio.scope() as scope:
            scope.spawn(serve_ssl_over_tcp(_server_handler, host='127.0.0.1', port=port, ssl_context=ssl_server_ctx))
            scope.spawn(client(port))
            yield done.wait()
            scope.cancel()
        yield scope()

    def client(port):
        yield tonio.sleep(0.5)
        stream: SSLStream = yield open_ssl_over_tcp_stream('127.0.0.1', port=port, ssl_context=ssl_client_ctx)
        while len(state['data']) < _SIZE:
            state['data'] += yield stream.receive_some()
        done.set()

    run(server())
    assert state['data'] == b'a' * _SIZE
