# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0.

import asyncio
from awscrt.io import *
from awscrt.websocket import *
from concurrent.futures import Future
from contextlib import closing
import gc
from io import StringIO
import logging
from os import urandom
from queue import Empty, Queue
import secrets
import socket
from test import NativeResourceTest
import threading
from time import sleep, time
from typing import Optional

# using a 3rdparty websocket library for the server
import websockets.server as websockets_server_3rdparty

TIMEOUT = 10.0  # seconds

# uncomment this for logging from the 3rdparty websockets server
# logging.basicConfig(format="%(message)s", level=logging.DEBUG)

# uncomment this for logging from our websockets client
# init_logging(LogLevel.Trace, 'stderr')


@dataclass
class RecvFrame:
    frame: IncomingFrame
    payload: bytes
    exception: Optional[Exception]


class ClientHandler:
    def __init__(self):
        self.websocket = None
        self.setup_future = Future()
        self.shutdown_future = Future()
        self.complete_frames = Queue()
        self.incoming_frame = None
        self.incoming_frame_payload = bytearray()
        self.exception = None

    def connect_sync(self, host, port, **connect_kwargs):
        connect(host=host,
                port=port,
                handshake_request=create_handshake_request(host=host),
                on_connection_setup=self._on_connection_setup,
                on_connection_shutdown=self._on_connection_shutdown,
                on_incoming_frame_begin=self._on_incoming_frame_begin,
                on_incoming_frame_payload=self._on_incoming_frame_payload,
                on_incoming_frame_complete=self._on_incoming_frame_complete,
                **connect_kwargs)
        # wait for on_connection_setup to fire
        setup_data = self.setup_future.result(TIMEOUT)
        assert setup_data.exception is None

    def close_sync(self):
        self.websocket.close()
        # wait for on_connection_shutdown to fire
        self.shutdown_future.result(TIMEOUT)

    def _raise_exception(self, msg=None):
        self.exception = RuntimeError(msg) if msg else RuntimeError()
        raise self.exception

    def _assert(self, condition, msg=None):
        if not condition:
            self._raise_exception(msg)

    def _on_connection_setup(self, data: OnConnectionSetupData):
        self._assert(not self.setup_future.done(), "setup must only fire once")
        self.websocket = data.websocket
        self.setup_future.set_result(data)

    def _on_connection_shutdown(self, data: OnConnectionShutdownData):
        self._assert(self.setup_future.done(), "setup must precede shutdown")
        self._assert(not self.shutdown_future.done(), "shutdown must only fire once")
        self._assert(self.incoming_frame is None, "incoming_frame_complete should fire before shutdown")
        self.shutdown_future.set_result(data)

    def _on_incoming_frame_begin(self, data: OnIncomingFrameBeginData):
        self._assert(self.incoming_frame is None,
                     "incoming_frame_begin cannot fire again until incoming_frame_complete")
        self.incoming_frame = data.frame

    def _on_incoming_frame_payload(self, data: OnIncomingFramePayloadData):
        self._assert(self.incoming_frame == data.frame, "frame from payload callback must match begin callback")
        self.incoming_frame_payload += data.data

    def _on_incoming_frame_complete(self, data: OnIncomingFrameCompleteData):
        self._assert(self.incoming_frame == data.frame,
                     "frame from complete callback must match begin callback")
        self.complete_frames.put(RecvFrame(self.incoming_frame, bytes(self.incoming_frame_payload), data.exception))
        self.incoming_frame = None
        self.incoming_frame_payload.clear()


class WebSocketServer:
    # Usage: Instantiate this WebSocketServer class in a `with` block.
    # On creation, it starts a thread to run the asyncio server.
    # Exiting the `with` block will cause the server to close and the thread to end.
    #
    # Details: We use the 3rdparty "websockets" library to host a local server
    # that we can test our client against.
    # This 3rdparty library has an asyncio API.
    # Our client does not have an asyncio API.
    # Since it's touch to mix code that uses asyncio with stuff that doesn't,
    # Since asyncio APIs don't mix well with non-asyncio stuff,
    # we launch a thread to run all the asyncio stuff.

    def __init__(self, host, port):
        self._host = host
        self._port = port
        # We use an old-fashioned thread-safe Event to signal the main thread
        # that the asyncio server thread has finished startup.
        self._server_started_event = threading.Event()
        self._server_thread = threading.Thread(target=self._run_server_thread)

    def __enter__(self):
        # main thread is entering the `with` block: start the server...
        self._server_thread.start()

        # don't return until the server signals that it's started up and is listening for connections
        assert self._server_started_event.wait(TIMEOUT)

        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        # main thread is exiting the `with` block: tell the server to stop...

        # asyncio events aren't actually thread-safe, so we can't simply call set() from here.
        # we need to call it from the loop.
        self._server_loop.call_soon_threadsafe(self._server_stop_event.set)

        # don't return until the server thread exits
        self._server_thread.join(TIMEOUT)
        # check whether thread really exited, or we just timed out
        assert (self._server_thread.is_alive() == False)

    def _run_server_thread(self):
        asyncio.run(self._run_asyncio_server())

    async def _run_asyncio_server(self):
        # store things the main thread will need later, so it can signal us to stop...
        self._server_loop = asyncio.get_running_loop()
        self._server_stop_event = asyncio.Event()

        # this coroutine runs the server until the _stop_event fires
        async with websockets_server_3rdparty.serve(self._run_connection, self._host, self._port, max_size=None) as server:
            # signal that server has started up
            self._server_started_event.set()
            # wait for the signal that we should stop
            await self._server_stop_event.wait()

    async def _run_connection(self, server_connection: websockets_server_3rdparty.WebSocketServerProtocol):
        # this coroutine runs once for each connection to the server
        # when this coroutine exits, the connection gets shut down
        self._current_connection = server_connection
        try:
            # await each message...
            async for msg in server_connection:
                # echo message back
                # print(f"server msg: {msg}")
                await server_connection.send(msg)

        except Exception:
            # an exception is raised when the connection ends,
            # even if the connection ends cleanly, so just swallow it
            pass

        finally:
            self._current_connection = None

    def send_async(self, msg):
        asyncio.run_coroutine_threadsafe(self._current_connection.send(msg), self._server_loop)


class TestClient(NativeResourceTest):
    def setUp(self):
        super().setUp()
        # Note: specifying IPV4 "127.0.0.1", instead of "localhost".
        # "localhost" leads to some machines hitting errors attempting to bind IPV6 addresses.
        self.host = '127.0.0.1'
        self.port = self._find_free_port()

    def _find_free_port(self):
        # find a free port by binding a temporary socket to port 0
        # and seeing what OS gives us
        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
            sock.bind(('localhost', 0))
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            return sock.getsockname()[1]

    def test_create_handshake_request(self):
        request = create_handshake_request(host='localhost', path='/mypath')
        self.assertEqual('GET', request.method)
        self.assertEqual('/mypath', request.path)
        # check that request has some of the necessary headers
        self.assertIn(('Host', 'localhost'), request.headers)
        self.assertIn(('Upgrade', 'websocket'), request.headers)

    def test_connect(self):
        # test a simple successful connection
        with WebSocketServer(self.host, self.port) as server:
            setup_future = Future()
            shutdown_future = Future()

            connect(
                host=self.host,
                port=self.port,
                handshake_request=create_handshake_request(host=self.host),
                on_connection_setup=lambda x: setup_future.set_result(x),
                on_connection_shutdown=lambda x: shutdown_future.set_result(x))

            # wait for on_connection_setup() to fire
            setup_data: OnConnectionSetupData = setup_future.result(TIMEOUT)

            # check that we had a successful connection
            self.assertIsNone(setup_data.exception)
            self.assertIsNotNone(setup_data.websocket)
            # should be a 101 response
            self.assertEqual(101, setup_data.handshake_response_status)
            # check for response header we know should be there
            self.assertIn(("Upgrade", "websocket"), setup_data.handshake_response_headers)
            # a successful handshake response has no body
            self.assertIsNone(setup_data.handshake_response_body)

            # now close the WebSocket
            setup_data.websocket.close()

            # wait for on_connection_shutdown() to fire
            shutdown_data: OnConnectionShutdownData = shutdown_future.result(TIMEOUT)
            self.assertIsNone(shutdown_data.exception, "Should have shut down cleanly")

    def test_closes_on_zero_refcount(self):
        # test that a WebSocket shuts down if it's garbage-collected while the connection is still open
        with WebSocketServer(self.host, self.port) as server:
            setup_future = Future()
            shutdown_future = Future()

            connect(
                host=self.host,
                port=self.port,
                handshake_request=create_handshake_request(host=self.host),
                on_connection_setup=lambda x: setup_future.set_result(x),
                on_connection_shutdown=lambda x: shutdown_future.set_result(x))

            # wait for on_connection_setup to fire
            websocket = setup_future.result(TIMEOUT).websocket

            # ensure the connection stays alive while we hold a reference
            gc.collect()
            sleep(0.5)
            self.assertFalse(shutdown_future.done())

            # drop any references to the WebSocket, and ensure the connection shuts down
            del setup_future
            del websocket
            gc.collect()
            shutdown_data: OnConnectionShutdownData = shutdown_future.result(TIMEOUT)
            self.assertIsNone(shutdown_data.exception, "Should have shut down cleanly")

    def test_connect_failure_without_response(self):
        # test a connection that fails before an HTTP response is received

        # the connection will fail because we're not running a server for this test
        setup_future = Future()
        shutdown_future = Future()

        connect(
            host=self.host,
            port=self.port,
            handshake_request=create_handshake_request(host=self.host),
            on_connection_setup=lambda x: setup_future.set_result(x),
            on_connection_shutdown=lambda x: shutdown_future.set_result(x))

        # wait for on_connection_setup to fire
        setup_data: OnConnectionSetupData = setup_future.result(TIMEOUT)
        self.assertIsNone(setup_data.websocket)
        self.assertIsNotNone(setup_data.exception)

        # nothing responded, so there should be no "handshake response"
        self.assertIsNone(setup_data.handshake_response_status)
        self.assertIsNone(setup_data.handshake_response_headers)
        self.assertIsNone(setup_data.handshake_response_body)

        # ensure that on_connection_shutdown does NOT fire
        sleep(0.5)
        self.assertFalse(shutdown_future.done(), "on_connection_shutdown should not have fired")

    def test_connect_failure_with_response(self):
        # test a connection that fails due an HTTP response rejecting the connection
        with WebSocketServer(self.host, self.port) as server:
            setup_future = Future()

            # remove necessary headers, so the server will reject this request
            bad_request = create_handshake_request(host=self.host)
            bad_request.headers.remove('Upgrade')  # "Upgrade: websocket"
            bad_request.headers.remove('Connection')  # "Connection: Upgrade"

            connect(
                host=self.host,
                port=self.port,
                handshake_request=bad_request,
                on_connection_setup=lambda x: setup_future.set_result(x))

            # wait for on_connection_setup to fire
            setup_data: OnConnectionSetupData = setup_future.result(TIMEOUT)
            self.assertIsNone(setup_data.websocket)
            self.assertIsNotNone(setup_data.exception)
            self.assertEqual("AWS_ERROR_HTTP_WEBSOCKET_UPGRADE_FAILURE", setup_data.exception.name)

            # check the HTTP response data
            self.assertGreaterEqual(setup_data.handshake_response_status, 400)
            self.assertIsNotNone(setup_data.handshake_response_headers)
            self.assertIsNotNone(setup_data.handshake_response_body)
            # check that body is a valid string
            self.assertGreater(len(setup_data.handshake_response_body.decode()), 0)

    def test_exception_in_setup_callback_closes_websocket(self):
        with WebSocketServer(self.host, self.port) as server:
            setup_future = Future()
            shutdown_future = Future()

            def bad_setup_callback(data: OnConnectionSetupData):
                setup_future.set_result(data)
                raise RuntimeError("Purposefully raising exception")

            connect(
                host=self.host,
                port=self.port,
                handshake_request=create_handshake_request(host=self.host),
                on_connection_setup=bad_setup_callback,
                on_connection_shutdown=lambda x: shutdown_future.set_result(x))

            # wait for on_connection_setup to fire (and raise exception)
            setup_data = setup_future.result(TIMEOUT)
            self.assertIsNotNone(setup_data.websocket)

            # wait for websocket to close due to exception
            shutdown_future.result(TIMEOUT)

    def test_exception_in_shutdown_callback_has_no_effect(self):
        with WebSocketServer(self.host, self.port) as server:
            setup_future = Future()
            shutdown_future = Future()

            def bad_shutdown_callback(data: OnConnectionShutdownData):
                shutdown_future.set_result(data)
                raise RuntimeError("Purposefully raising exception")

            connect(
                host=self.host,
                port=self.port,
                handshake_request=create_handshake_request(host=self.host),
                on_connection_setup=lambda x: setup_future.set_result(x),
                on_connection_shutdown=bad_shutdown_callback)

            # wait for on_connection_setup to fire
            websocket = setup_future.result(TIMEOUT).websocket

            # close websocket, on_connection_shutdown will fire and raise exception,
            # which should have no effect
            websocket.close()
            shutdown_future.result(TIMEOUT)

    def test_close_is_idempotent(self):
        # test that it's always safe to call WebSocket.close()
        with WebSocketServer(self.host, self.port) as server:
            setup_future = Future()
            shutdown_future = Future()

            connect(
                host=self.host,
                port=self.port,
                handshake_request=create_handshake_request(host=self.host),
                on_connection_setup=lambda x: setup_future.set_result(x),
                on_connection_shutdown=lambda x: shutdown_future.set_result(x))

            # wait for on_connection_setup to fire
            websocket = setup_future.result(TIMEOUT).websocket

            # now call close() A BUNCH of times...
            self.assertFalse(shutdown_future.done())

            # ...before shutdown has started!
            websocket.close()
            websocket.close()

            # ...while shutdown is happening!
            sleep(0.0000001)
            websocket.close()
            websocket.close()

            # wait for shutdown.
            shutdown_data: OnConnectionShutdownData = shutdown_future.result(TIMEOUT)
            self.assertIsNone(shutdown_data.exception, "Shutdown should have been clean")

            # ...and after shutdown has already occurred!
            websocket.close()
            websocket.close()

    def _send_and_receive(self, handler, opcode, payload):
        send_complete_future = Future()

        handler.websocket.send_frame(opcode, payload, on_complete=lambda x: send_complete_future.set_result(x))

        # wait for send_frame operation to complete
        send_complete_data: OnSendFrameCompleteData = send_complete_future.result(TIMEOUT)
        self.assertIsNone(send_complete_data.exception)

        # wait to receive echo message back
        recv: RecvFrame = handler.complete_frames.get(timeout=TIMEOUT)
        # assert that expected types are received (`is True` vs checking that it evaluates to True)
        self.assertIsInstance(recv.frame.opcode, Opcode)
        self.assertEqual(recv.frame.opcode, opcode)
        self.assertIs(recv.frame.fin, True)

        # check that the received payload matches what we sent in
        if isinstance(payload, str):
            self.assertEqual(recv.payload, payload.encode('utf-8'))
        elif payload is None:
            self.assertEqual(recv.payload, b'')
        else:
            self.assertEqual(recv.payload, payload)

        self.assertIsNone(handler.exception)

    def test_send_receive_data(self):
        # test sending and receiving TEXT and BINARY frames
        # (the server echos these types of messages back)
        with WebSocketServer(self.host, self.port) as server:
            handler = ClientHandler()
            handler.connect_sync(self.host, self.port)

            # try sending all kinds of "buffer protocol" types
            self._send_and_receive(handler, Opcode.TEXT, "str with ascii")
            self._send_and_receive(handler, Opcode.TEXT, "str with unicode --> 👁👄👁 <--")
            self._send_and_receive(handler, Opcode.TEXT, "str with embedded null byte --> \0 <--")
            self._send_and_receive(handler, Opcode.TEXT, b"bytes of text")
            self._send_and_receive(handler, Opcode.TEXT, bytearray(b"bytearray of text"))
            self._send_and_receive(handler, Opcode.TEXT, memoryview(b"memoryview of text"))
            self._send_and_receive(handler, Opcode.TEXT, memoryview(b"...memoryview slice...")[3: -3])
            self._send_and_receive(handler, Opcode.TEXT, "")  # empty
            self._send_and_receive(handler, Opcode.TEXT, None)

            # try sending all kinds of "buffer protocol" types
            self._send_and_receive(handler, Opcode.BINARY, "str sent binary")
            self._send_and_receive(handler, Opcode.BINARY, b"bytes sent binary")
            self._send_and_receive(handler, Opcode.BINARY, bytearray(b"bytearray sent binary"))
            self._send_and_receive(handler, Opcode.BINARY, memoryview(b"memoryview sent binary"))
            self._send_and_receive(handler, Opcode.BINARY, memoryview(b"...memoryview slice sent binary...")[3: -3])
            self._send_and_receive(handler, Opcode.BINARY, bytes())  # empty
            self._send_and_receive(handler, Opcode.BINARY, None)

            # send something very big
            self._send_and_receive(handler, Opcode.BINARY, urandom(1024 * 1024 * 4))

            handler.close_sync()
            self.assertIsNone(handler.exception)

    def test_send_frame_exceptions(self):
        with WebSocketServer(self.host, self.port) as server:
            handler = ClientHandler()
            handler.connect_sync(self.host, self.port)

            # we don't currently support sending streams
            with self.assertRaises(TypeError):
                handler.websocket.send_frame(Opcode.TEXT, StringIO("text stream"))

            # raising an exception from the completion-callback should result in websocket closing
            def bad_completion_callback(data):
                raise RuntimeError("Purposefully raising exception")

            handler.websocket.send_frame(Opcode.TEXT, "asdf", on_complete=bad_completion_callback)

            # wait for shutdown...
            handler.shutdown_future.result(TIMEOUT)

            # shouldn't be able to send frame after websocket closes
            with self.assertRaises(Exception) as raises:
                handler.websocket.send_frame(Opcode.TEXT, "asdf")
            self.assertIn("AWS_ERROR_HTTP_WEBSOCKET_CLOSE_FRAME_SENT", str(raises.exception))

            self.assertIsNone(handler.exception)

    def test_exception_from_incoming_frame_callback_closes_websocket(self):
        # loop 3 times, once for each type of on_incoming_frame_X callback
        for i in ('begin', 'payload', 'complete'):
            with WebSocketServer(self.host, self.port) as server:
                setup_future = Future()
                shutdown_future = Future()

                def bad_incoming_frame_callback(data):
                    raise RuntimeError("Purposefully raising exception")

                connect(
                    host=self.host,
                    port=self.port,
                    handshake_request=create_handshake_request(host=self.host),
                    on_connection_setup=lambda x: setup_future.set_result(x),
                    on_connection_shutdown=lambda x: shutdown_future.set_result(x),
                    on_incoming_frame_begin=bad_incoming_frame_callback if i == 'begin' else None,
                    on_incoming_frame_payload=bad_incoming_frame_callback if i == 'payload' else None,
                    on_incoming_frame_complete=bad_incoming_frame_callback if i == 'complete' else None,
                )

                # wait for on_connection_setup to fire
                websocket = setup_future.result(TIMEOUT).websocket

                # send a frame that the server will echo back
                websocket.send_frame(Opcode.TEXT, "echo")

                # wait for the frame to echo back, firing the bad callback,
                # which raises an exception, which should result in the WebSocket closing
                shutdown_future.result(TIMEOUT)

    def test_manage_read_window(self):
        # test that users can manage how much data is read by managing the read window
        with WebSocketServer(self.host, self.port) as server:
            handler = ClientHandler()
            handler.connect_sync(self.host, self.port, manage_read_window=True, initial_read_window=1000)

            # client's read window is 1000-bytes
            # have the server send 10 messages with 100-byte payloads
            # they should all get through

            for i in range(10):
                msg = secrets.token_bytes(100)  # random msg for server to send
                server.send_async(msg)
                recv: RecvFrame = handler.complete_frames.get(timeout=TIMEOUT)
                self.assertEqual(recv.payload, msg, "did not receive expected payload")

            # client window is now 0
            # have server send a 1000 byte message, NONE of its payload should arrive

            msg = secrets.token_bytes(1000)  # random msg for server to send
            server.send_async(msg)
            with self.assertRaises(Empty):
                handler.complete_frames.get(timeout=1.0)
            self.assertEqual(len(handler.incoming_frame_payload), 0, "No payload should arrive while window is 0")

            # now increment client's window to 500
            # half (500/1000) the bytes should flow in

            handler.websocket.increment_read_window(500)
            max_wait_until = time() + TIMEOUT
            while len(handler.incoming_frame_payload) < 500:
                sleep(0.001)
                self.assertLess(time(), max_wait_until, "timed out waiting for all bytes")
            sleep(1.0)  # sleep a moment to be sure we don't receive MORE than 500 bytes
            self.assertEqual(len(handler.incoming_frame_payload), 500, "received more bytes than expected")

            # client's window is 0 again, 500 bytes are still waiting to flow in
            # increment the window to let the rest in
            # let's do it by calling increment a bunch of times in a row, just to be different

            handler.websocket.increment_read_window(100)
            handler.websocket.increment_read_window(100)
            handler.websocket.increment_read_window(100)
            handler.websocket.increment_read_window(100)
            handler.websocket.increment_read_window(100)

            recv: RecvFrame = handler.complete_frames.get(timeout=TIMEOUT)
            self.assertEqual(recv.payload, msg, "did not receive expected payload")

            # done!
            handler.close_sync()