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

from awscrt.eventstream import *
from awscrt.eventstream.rpc import *
from awscrt.io import (ClientBootstrap, DefaultHostResolver, EventLoopGroup, init_logging, LogLevel)
import gc
import os
from queue import Queue
from test import NativeResourceTest, TIMEOUT
from threading import Event
import time
from unittest import skipUnless
from uuid import UUID, uuid4
import weakref

# TODO: setup permanent online echo server we can hit from tests
RUN_LOCALHOST_TESTS = os.getenv('EVENTSTREAM_ECHO_TEST')

#init_logging(LogLevel.Trace, 'stderr')


class ConnectionRecord:
    def __init__(self):
        self.setup_call = None
        self.shutdown_call = None
        self.message_calls = Queue()


class ConnectionHandler(ClientConnectionHandler):
    def __init__(self, fail_test_fn):
        self.record = ConnectionRecord()
        self.fail_test = fail_test_fn
        self.connection = None
        # whether to purposefully raise an exception in a callback
        self.raise_during_setup = False

    def on_connection_setup(self, connection, error, **kwargs):
        if self.record.setup_call is not None:
            self.fail_test("setup can only fire once")
        self.record.setup_call = {'connection': connection, 'error': error}
        self.connection = connection
        if self.raise_during_setup:
            raise RuntimeError("Purposefully raising error in callback")

    def on_connection_shutdown(self, reason, **kwargs):
        if self.record.setup_call is None:
            self.fail_test("if setup doesn't fire, shutdown shouldn't fire")
        if self.record.setup_call['error'] is not None:
            self.fail_test("shutdown shouldn't fire if setup had an error")
        if self.record.shutdown_call is not None:
            self.fail_test("shutdown can only fire once")
        self.record.shutdown_call = {'reason': reason}

    def on_protocol_message(self, headers, payload, message_type, flags, **kwargs):
        if self.record.setup_call is None:
            self.fail_test("setup must fire before messages")
        if self.record.shutdown_call is not None:
            self.fail_test("messages cannot fire after shutdown")
        self.record.message_calls.put(
            {'headers': headers, 'payload': payload, 'message_type': message_type, 'flags': flags})


class ContinuationRecord:
    def __init__(self):
        self.message_calls = Queue()
        self.close_call = Event()


class ContinuationHandler(ClientContinuationHandler):
    def __init__(self, fail_test_fn):
        self.record = ContinuationRecord()
        self.fail_test = fail_test_fn

    def on_continuation_message(
            self,
            headers,
            payload,
            message_type,
            flags,
            **kwargs):
        if self.record.close_call.is_set():
            self.fail_test("messages should not fire after close")
        self.record.message_calls.put(
            {'headers': headers, 'payload': payload, 'message_type': message_type, 'flags': flags})

    def on_continuation_closed(self, **kwargs):
        if self.record.close_call.is_set():
            self.fail_test("shutdown can only fire once")
        self.record.close_call.set()


@skipUnless(RUN_LOCALHOST_TESTS, "Skipping until we have permanent echo server")
class TestClient(NativeResourceTest):

    def _fail_test_from_callback(self, msg):
        print("ERROR FROM CALLBACK", msg)
        self._failure_from_callback = msg

    def _assertNoFailuresFromCallbacks(self):
        self.assertIsNone(getattr(self, '_failure_from_callback', None))

    def test_connect_failure(self):
        elg = EventLoopGroup()
        resolver = DefaultHostResolver(elg)
        bootstrap = ClientBootstrap(elg, resolver)
        handler = ConnectionHandler(self._fail_test_from_callback)
        future = ClientConnection.connect(
            handler=handler,
            host_name="fake",
            port=99,
            bootstrap=bootstrap)

        failure = future.exception(TIMEOUT)
        self.assertTrue(isinstance(failure, Exception))
        self.assertIsNotNone(handler.record.setup_call)
        self.assertIsNone(handler.record.setup_call['connection'])
        self.assertTrue(isinstance(handler.record.setup_call['error'], Exception))
        self.assertIsNone(handler.record.shutdown_call)
        self._assertNoFailuresFromCallbacks()

    def test_connect_success(self):
        elg = EventLoopGroup()
        resolver = DefaultHostResolver(elg)
        bootstrap = ClientBootstrap(elg, resolver)
        handler = ConnectionHandler(self._fail_test_from_callback)
        future = ClientConnection.connect(
            handler=handler,
            host_name="127.0.0.1",
            port=8033,
            bootstrap=bootstrap)

        self.assertIsNone(future.exception(TIMEOUT))
        self.assertIsNotNone(handler.record.setup_call)
        self.assertTrue(isinstance(handler.record.setup_call['connection'], ClientConnection))
        self.assertIsNone(handler.record.setup_call['error'])
        self.assertTrue(handler.connection.is_open())
        self.assertIsNone(handler.record.shutdown_call)

        # close
        shutdown_future = handler.connection.close()
        self.assertIsNone(shutdown_future.exception(TIMEOUT))
        self.assertIsNotNone(handler.record.shutdown_call)
        self.assertIsNone(handler.record.shutdown_call['reason'])

        self._assertNoFailuresFromCallbacks()

    def _connect_and_return_weakref(self):
        elg = EventLoopGroup()
        resolver = DefaultHostResolver(elg)
        bootstrap = ClientBootstrap(elg, resolver)
        handler = ConnectionHandler(self._fail_test_from_callback)
        future = ClientConnection.connect(
            handler=handler,
            host_name="127.0.0.1",
            port=8033,
            bootstrap=bootstrap)

        # connection should succeed
        self.assertIsNone(future.exception(TIMEOUT))

        return weakref.ref(handler.connection)

    def test_connection_stays_alive_until_close(self):
        # Currently, eventstream connections do NOT close just because their refcount reaches zero.
        # We *could* make it work this way, but it would take some combo of:
        # 1) handler shutdown callbacks aren't guaranteed to fire,
        #    because handler likely has reference to connection, and we want
        #    them to get GC'd together.
        # 2) OR guarantee shutdown callbacks, but explain the subtleties of
        #    avoiding a strong reference cycle involving the handler.
        #    This is not a typical thing Python programmers need to think about.
        # Seems best to keep it simple and just remind the user at every
        # opportunity that connection must be closed to avoid resource leaks.
        connection_weakref = self._connect_and_return_weakref()

        # wait a moment to see if connection gets cleaned up
        time.sleep(1)
        gc.collect()

        connection = connection_weakref()
        self.assertIsNotNone(connection, "open connection should not have been GC'd")
        self.assertTrue(connection.is_open())

        # ok, now shut it down, (and nuke local reference to connection
        # just to be sure we don't crash if both happen simultaneously)
        shutdown_future = connection.close()
        del connection
        self.assertIsNone(shutdown_future.exception(TIMEOUT))

        # GC should have cleaned it up, now that the connection is closed
        # and no strong references remain
        gc.collect()
        self.assertIsNone(connection_weakref())

        self._assertNoFailuresFromCallbacks()

    def test_setup_callback_exception_closes_connection(self):
        elg = EventLoopGroup()
        resolver = DefaultHostResolver(elg)
        bootstrap = ClientBootstrap(elg, resolver)
        handler = ConnectionHandler(self._fail_test_from_callback)

        # tell handler to explode during setup callback
        handler.raise_during_setup = True
        connect_future = ClientConnection.connect(
            handler=handler,
            host_name="127.0.0.1",
            port=8033,
            bootstrap=bootstrap)

        # the connect future does actually succeed, since the user's setup
        # callback got a chance to run.
        connect_future.exception(TIMEOUT)

        # but the unhandled exception should cause the connection to die
        shutdown_reason = handler.connection.shutdown_future.exception(TIMEOUT)
        self.assertTrue("CALLBACK_EXCEPTION" in shutdown_reason.name)
        self.assertTrue("CALLBACK_EXCEPTION" in handler.record.shutdown_call['reason'].name)

    def _connect_fully(self):
        # connect, send CONNECT message, receive CONNECT_ACK
        elg = EventLoopGroup()
        resolver = DefaultHostResolver(elg)
        bootstrap = ClientBootstrap(elg, resolver)
        handler = ConnectionHandler(self._fail_test_from_callback)
        connect_future = ClientConnection.connect(
            handler=handler,
            host_name="127.0.0.1",
            port=8033,
            bootstrap=bootstrap)

        self.assertIsNone(connect_future.exception(TIMEOUT))

        # send CONNECT msg
        msg_future = handler.connection.send_protocol_message(
            message_type=MessageType.CONNECT,
            headers=[Header.from_string(':version', '0.1.0'),
                     Header.from_string('client-name', 'accepted.testy_mc_testerson')])

        self.assertIsNone(msg_future.exception(TIMEOUT))

        # wait to receive CONNECT_ACK
        msg = handler.record.message_calls.get(timeout=TIMEOUT)
        self.assertIs(MessageType.CONNECT_ACK, msg['message_type'])
        self.assertTrue(MessageFlag.CONNECTION_ACCEPTED & msg['flags'])

        return handler

    def test_stream_cleans_up_if_never_activated(self):
        # check that there are no resource leaks if a stream/continuation is never activated
        handler = self._connect_fully()

        # create stream, but do not activate it
        stream_handler = ContinuationHandler(self._fail_test_from_callback)
        continuation = handler.connection.new_stream(stream_handler)

        # close connection
        handler.connection.close()
        self.assertIsNone(handler.connection.shutdown_future.exception(TIMEOUT))

        # stream callbacks should not fire because it was never activated
        self.assertFalse(stream_handler.record.close_call.isSet())

        self._assertNoFailuresFromCallbacks()

    def test_stream_message_echo(self):
        # test sending and receiving messages
        handler = self._connect_fully()

        # send PING msg, server will echo back its headers and payload in the PING_RESPONSE.
        # test every single header type
        echo_headers = [
            Header.from_bool('echo-true', True),
            Header.from_bool('echo-false', False),
            Header.from_byte('echo-byte', 127),
            Header.from_int16('echo-int16', 32000),
            Header.from_int32('echo-int32', 2000000000),
            Header.from_int64('echo-int64', 999999999999),
            Header.from_byte_buf('echo-byte-buf', b'\x00\xff\x0f\xf0'),
            Header.from_string('echo-string', 'noodles'),
            # utf-8 breaks echo test. don't get response.
            #Header.from_string('echo-string-utf8', '--\u1234--'),
            Header.from_timestamp('echo-timestamp', time.time()),
            Header.from_uuid('echo-uuid', UUID('01234567-89ab-cdef-0123-456789abcdef')),
        ]
        echo_payload = b'\x00\xDE\xAD\xBE\xEF'
        msg_future = handler.connection.send_protocol_message(
            message_type=MessageType.PING,
            headers=echo_headers,
            payload=echo_payload)

        self.assertIsNone(msg_future.exception(TIMEOUT))

        # wait to receive PING_RESPONSE, which should echo what we sent
        msg = handler.record.message_calls.get(timeout=TIMEOUT)
        self.assertIs(MessageType.PING_RESPONSE, msg['message_type'])
        for sent_header in echo_headers:
            recv_header = next(x for x in msg['headers'] if x.name.lower() == sent_header.name.lower())
            self.assertEqual(sent_header.type, recv_header.type)
            self.assertEqual(sent_header.value, recv_header.value)
        self.assertEqual(echo_payload, msg['payload'])

        # use a stream to execute the "EchoMessage" operation,
        # which takes 1 message and responds with 1 message
        stream_handler = ContinuationHandler(self._fail_test_from_callback)
        continuation = handler.connection.new_stream(stream_handler)
        msg_flushed = Event()

        def on_msg_flush(error, **kwargs):
            msg_flushed.set()

        msg_future = continuation.activate(
            operation='awstest#EchoMessage',
            headers=[],
            payload=b'{}',
            message_type=MessageType.APPLICATION_MESSAGE,
            flags=MessageFlag.NONE,
            on_flush=on_msg_flush)

        self.assertIsNone(msg_future.exception(TIMEOUT))
        msg_flushed.wait(TIMEOUT)

        # wait to receive response, which should end the stream
        msg = stream_handler.record.message_calls.get(timeout=TIMEOUT)
        self.assertEqual(MessageType.APPLICATION_MESSAGE, msg['message_type'])
        self.assertTrue(MessageFlag.TERMINATE_STREAM & msg['flags'])
        self.assertIsNone(continuation.closed_future.exception(TIMEOUT))

        # use a stream to execute the "EchoStreamMessages" operation,
        # which has an empty initial request/response, but then allows further messages
        stream_handler = ContinuationHandler(self._fail_test_from_callback)
        continuation = handler.connection.new_stream(stream_handler)
        msg_future = continuation.activate(
            operation='awstest#EchoStreamMessages',
            headers=[],
            payload=b'{}',
            message_type=MessageType.APPLICATION_MESSAGE,
            flags=MessageFlag.NONE)

        self.assertIsNone(msg_future.exception(TIMEOUT))

        # wait to initial response, which should end the stream
        msg = stream_handler.record.message_calls.get(timeout=TIMEOUT)
        self.assertEqual(MessageType.APPLICATION_MESSAGE, msg['message_type'])
        self.assertFalse(MessageFlag.TERMINATE_STREAM & msg['flags'])

        # send a 2nd message on the stream
        msg_future = continuation.send_message(
            headers=[],
            payload=b'{}',
            message_type=MessageType.APPLICATION_MESSAGE,
            flags=MessageFlag.NONE)

        self.assertIsNone(msg_future.exception(TIMEOUT))

        # wait for 2nd response
        msg = stream_handler.record.message_calls.get(timeout=TIMEOUT)
        self.assertEqual(MessageType.APPLICATION_MESSAGE, msg['message_type'])
        self.assertFalse(MessageFlag.TERMINATE_STREAM & msg['flags'])

        # close connection
        handler.connection.close()
        self.assertIsNone(handler.connection.shutdown_future.exception(TIMEOUT))
        self.assertIsNone(continuation.closed_future.exception(TIMEOUT))

        self._assertNoFailuresFromCallbacks()

    def _start_request_and_drop_local_references(self, connection):
        # use a stream to execute the "EchoMessage" operation,
        # which takes 1 message and responds with 1 message
        stream_handler = ContinuationHandler(self._fail_test_from_callback)
        continuation = connection.new_stream(stream_handler)
        msg_future = continuation.activate(
            operation='awstest#EchoMessage',
            headers=[],
            payload=b'{}',
            message_type=MessageType.APPLICATION_MESSAGE,
            flags=MessageFlag.NONE)

        return (continuation.closed_future, stream_handler.record)

    def test_stream_stays_alive_until_close(self):
        # Test checks that stream stays alive, and all callbacks fire,
        # even if there are no local references to stream or its handler.
        handler = self._connect_fully()

        # create new stream and handler, and send request, but don't get back
        # anything that keeps a reference to these things.
        # function returns items that should show side-effects of the
        # request/response continuation to operate to completion.
        # If this test fails, it's because something got GC'd too early and all
        # the callbacks couldn't fire
        stream_closed_future, stream_record = self._start_request_and_drop_local_references(handler.connection)

        # wait to receive response, which should end the stream
        msg = stream_record.message_calls.get(timeout=TIMEOUT)

        # close connection
        handler.connection.close()
        self.assertIsNone(handler.connection.shutdown_future.exception(TIMEOUT))
        self.assertTrue(stream_record.close_call.wait(timeout=TIMEOUT))
        self.assertIsNone(stream_closed_future.result(timeout=TIMEOUT))

        self._assertNoFailuresFromCallbacks()

    def test_on_closed_deadlock_regression(self):
        # ensure that during the on_closed() callback of the first stream,
        # we can activate a second stream without deadlocking
        handler = self._connect_fully()

        # create 2 streams but don't activate yet
        first_stream_handler = DeadlockStreamHandler(handler.connection)
        first_stream = handler.connection.new_stream(first_stream_handler)

        second_stream_handler = ContinuationHandler(self._fail_test_from_callback)
        second_stream = handler.connection.new_stream(second_stream_handler)

        # activate first stream with message that closes the stream
        # from the on_closed() callback it will activate the second stream
        first_stream_handler.second_stream = second_stream
        first_stream.activate(operation="first",
                              message_type=MessageType.APPLICATION_MESSAGE)

        first_stream_handler.activated_second_stream.wait(TIMEOUT)

        handler.connection.close().result()


class DeadlockStreamHandler(ClientContinuationHandler):
    def __init__(self, connection):
        super().__init__()
        self.connection = connection
        self.activated_second_stream = Event()

    def on_continuation_message(self, headers, payload, message_type, flags, **kwargs) -> None:
        pass

    def on_continuation_closed(self, **kwargs) -> None:
        print("ACTIVATING")
        self.second_stream.activate(operation="next",
                                    message_type=MessageType.APPLICATION_MESSAGE)
        print("EVENTING")
        self.activated_second_stream.set()
