import asyncio
import gc
import pytest
import struct
from typing import Any
from unittest import mock

from kafka.protocol.metadata import (
    MetadataRequest_v0 as MetadataRequest,
    MetadataResponse_v0 as MetadataResponse)
from kafka.protocol.commit import (
    GroupCoordinatorRequest_v0 as GroupCoordinatorRequest,
    GroupCoordinatorResponse_v0 as GroupCoordinatorResponse)
from kafka.protocol.admin import (
    SaslHandShakeRequest, SaslHandShakeResponse, SaslAuthenticateRequest,
    SaslAuthenticateResponse
)
from kafka.protocol.produce import ProduceRequest_v0 as ProduceRequest

from aiokafka.conn import AIOKafkaConnection, create_conn, VersionInfo
from aiokafka.errors import (
    KafkaConnectionError, CorrelationIdError, KafkaError, NoError,
    UnknownError, UnsupportedSaslMechanismError, IllegalSaslStateError
)
from aiokafka.record.legacy_records import LegacyRecordBatchBuilder
from ._testutil import KafkaIntegrationTestCase, run_until_complete
from aiokafka.util import get_running_loop


@pytest.mark.usefixtures('setup_test_class')
class ConnIntegrationTest(KafkaIntegrationTestCase):

    @run_until_complete
    async def test_ctor(self):
        conn = AIOKafkaConnection('localhost', 1234)
        self.assertEqual('localhost', conn.host)
        self.assertEqual(1234, conn.port)
        self.assertTrue('KafkaConnection' in conn.__repr__())
        self.assertIsNone(conn._reader)
        self.assertIsNone(conn._writer)

    @run_until_complete
    async def test_global_loop_for_create_conn(self):
        loop = get_running_loop()
        host, port = self.kafka_host, self.kafka_port
        conn = await create_conn(host, port)
        self.assertIs(conn._loop, loop)
        conn.close()
        # make sure second closing does nothing and we have full coverage
        # of *if self._reader:* condition
        conn.close()

    @run_until_complete
    async def test_conn_warn_unclosed(self):
        host, port = self.kafka_host, self.kafka_port
        conn = await create_conn(
            host, port, max_idle_ms=100000)

        with self.silence_loop_exception_handler():
            with self.assertWarnsRegex(
                    ResourceWarning, "Unclosed AIOKafkaConnection"):
                del conn
                gc.collect()

    @run_until_complete
    async def test_basic_connection_load_meta(self):
        host, port = self.kafka_host, self.kafka_port
        conn = await create_conn(host, port)

        self.assertEqual(conn.connected(), True)
        request = MetadataRequest([])
        response = await conn.send(request)
        conn.close()
        self.assertIsInstance(response, MetadataResponse)

    @run_until_complete
    async def test_connections_max_idle_ms(self):
        host, port = self.kafka_host, self.kafka_port
        conn = await create_conn(
            host, port, max_idle_ms=200)
        self.assertEqual(conn.connected(), True)
        await asyncio.sleep(0.1)
        # Do some work
        request = MetadataRequest([])
        await conn.send(request)
        await asyncio.sleep(0.15)
        # Check if we're still connected after 250ms, as we were not idle
        self.assertEqual(conn.connected(), True)

        # It shouldn't break if we have a long running call either
        readexactly = conn._reader.readexactly
        with mock.patch.object(conn._reader, 'readexactly') as mocked:
            async def long_read(n):
                await asyncio.sleep(0.2)
                return (await readexactly(n))
            mocked.side_effect = long_read
            await conn.send(MetadataRequest([]))
        self.assertEqual(conn.connected(), True)

        await asyncio.sleep(0.2)
        self.assertEqual(conn.connected(), False)

    @run_until_complete
    async def test_send_without_response(self):
        """Imitate producer without acknowledge, in this case client produces
        messages and kafka does not send response, and we make sure that
        futures do not stuck in queue forever"""

        host, port = self.kafka_host, self.kafka_port
        conn = await create_conn(host, port)

        # prepare message
        builder = LegacyRecordBatchBuilder(
            magic=1, compression_type=0, batch_size=99999999)
        builder.append(offset=0, value=b"foo", key=None, timestamp=None)
        request = ProduceRequest(
            required_acks=0, timeout=10 * 1000,
            topics=[(b'foo', [(0, bytes(builder.build()))])])

        # produce messages without acknowledge
        req = []
        for i in range(10):
            req.append(conn.send(request, expect_response=False))
        # make sure futures no stuck in queue
        self.assertEqual(len(conn._requests), 0)
        for x in req:
            await x
        conn.close()

    @run_until_complete
    async def test_send_to_closed(self):
        host, port = self.kafka_host, self.kafka_port
        conn = AIOKafkaConnection(host=host, port=port)
        request = MetadataRequest([])
        with self.assertRaises(KafkaConnectionError):
            await conn.send(request)

        conn._writer = mock.MagicMock()
        conn._writer.write.side_effect = OSError('mocked writer is closed')

        with self.assertRaises(KafkaConnectionError):
            await conn.send(request)

    @run_until_complete
    async def test_invalid_correlation_id(self):
        host, port = self.kafka_host, self.kafka_port

        request = MetadataRequest([])

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = MetadataResponse(brokers=[], topics=[])
        resp = resp.encode()
        resp = int32.pack(999) + resp  # set invalid correlation id

        async def first_resp(*args: Any, **kw: Any):
            return int32.pack(len(resp))

        async def second_resp(*args: Any, **kw: Any):
            return resp

        reader.readexactly.side_effect = [first_resp(), second_resp()]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = conn._create_reader_task()

        with self.assertRaises(CorrelationIdError):
            await conn.send(request)

    @run_until_complete
    async def test_correlation_id_on_group_coordinator_req(self):
        host, port = self.kafka_host, self.kafka_port

        request = GroupCoordinatorRequest(consumer_group='test')

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = GroupCoordinatorResponse(
            error_code=0, coordinator_id=22,
            host='127.0.0.1', port=3333)
        resp = resp.encode()
        resp = int32.pack(0) + resp  # set correlation id to 0

        async def first_resp(*args: Any, **kw: Any):
            return int32.pack(len(resp))

        async def second_resp(*args: Any, **kw: Any):
            return resp

        reader.readexactly.side_effect = [first_resp(), second_resp()]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = conn._create_reader_task()

        response = await conn.send(request)
        self.assertIsInstance(response, GroupCoordinatorResponse)
        self.assertEqual(response.error_code, 0)
        self.assertEqual(response.coordinator_id, 22)
        self.assertEqual(response.host, '127.0.0.1')
        self.assertEqual(response.port, 3333)

    @run_until_complete
    async def test_osserror_in_reader_task(self):
        host, port = self.kafka_host, self.kafka_port

        async def invoke_osserror(*a, **kw):
            await asyncio.sleep(0.1)
            raise OSError('test oserror')

        request = MetadataRequest([])
        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port)

        # setup reader
        reader = mock.MagicMock()
        reader.readexactly.return_value = invoke_osserror()
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = conn._create_reader_task()

        with self.assertRaises(KafkaConnectionError):
            await conn.send(request)
        self.assertEqual(conn.connected(), False)

    @run_until_complete
    async def test_close_disconnects_connection(self):
        host, port = self.kafka_host, self.kafka_port
        conn = await create_conn(host, port)
        self.assertTrue(conn.connected())
        conn.close()
        self.assertFalse(conn.connected())

    def test_connection_version_info(self):
        # All version supported
        version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 1]
        })
        self.assertEqual(
            version_info.pick_best(SaslHandShakeRequest),
            SaslHandShakeRequest[1])

        # Broker only supports the lesser version
        version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 0]
        })
        self.assertEqual(
            version_info.pick_best(SaslHandShakeRequest),
            SaslHandShakeRequest[0])

        # We don't support any version compatible with the broker
        version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [2, 3]
        })
        with self.assertRaises(KafkaError):
            self.assertEqual(
                version_info.pick_best(SaslHandShakeRequest),
                SaslHandShakeRequest[1])

        # No information on the supported versions
        version_info = VersionInfo({})
        self.assertEqual(
            version_info.pick_best(SaslHandShakeRequest),
            SaslHandShakeRequest[0])

    @run_until_complete
    async def test__do_sasl_handshake_v0(self):
        host, port = self.kafka_host, self.kafka_port

        # setup connection with mocked send and send_bytes
        conn = AIOKafkaConnection(
            host=host, port=port,
            sasl_mechanism="PLAIN",
            sasl_plain_username="admin",
            sasl_plain_password="<PASSWORD>"
        )
        conn.close = close_mock = mock.MagicMock()

        supported_mechanisms = ["PLAIN"]
        error_class = NoError

        async def mock_send(request, expect_response=True):
            return SaslHandShakeResponse[0](
                error_code=error_class.errno,
                enabled_mechanisms=supported_mechanisms
            )

        async def mock_sasl_send(payload, expect_response):
            return b""

        conn.send = mock.Mock(side_effect=mock_send)
        conn._send_sasl_token = mock.Mock(side_effect=mock_sasl_send)
        conn._version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 0]
        })

        await conn._do_sasl_handshake()

        supported_mechanisms = ["GSSAPI"]
        with self.assertRaises(UnsupportedSaslMechanismError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)

        error_class = UnknownError
        close_mock.reset()

        with self.assertRaises(UnknownError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)

    @run_until_complete
    async def test__do_sasl_handshake_v1(self):
        host, port = self.kafka_host, self.kafka_port

        # setup connection with mocked send and send_bytes
        conn = AIOKafkaConnection(
            host=host, port=port,
            sasl_mechanism="PLAIN",
            sasl_plain_username="admin",
            sasl_plain_password="<PASSWORD>",
            security_protocol="SASL_PLAINTEXT"
        )
        conn.close = close_mock = mock.MagicMock()

        supported_mechanisms = ["PLAIN"]
        error_class = NoError
        auth_error_class = NoError

        async def mock_send(request, expect_response=True):
            if request.API_KEY == SaslHandShakeRequest[0].API_KEY:
                assert request.API_VERSION == 1
                return SaslHandShakeResponse[1](
                    error_code=error_class.errno,
                    enabled_mechanisms=supported_mechanisms
                )
            else:
                assert request.API_KEY == SaslAuthenticateRequest[0].API_KEY
                return SaslAuthenticateResponse[0](
                    error_code=auth_error_class.errno,
                    error_message="",
                    sasl_auth_bytes=b""
                )

        conn.send = mock.Mock(side_effect=mock_send)
        conn._version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 1]
        })

        await conn._do_sasl_handshake()

        supported_mechanisms = ["GSSAPI"]
        with self.assertRaises(UnsupportedSaslMechanismError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
        supported_mechanisms = ["PLAIN"]

        auth_error_class = IllegalSaslStateError
        close_mock.reset()

        with self.assertRaises(IllegalSaslStateError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
        auth_error_class = NoError

        error_class = UnknownError
        close_mock.reset()

        with self.assertRaises(UnknownError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)

    @run_until_complete
    async def test__send_sasl_token(self):
        # Before Kafka 1.0.0 SASL was performed on the wire without
        # KAFKA_HEADER in the protocol. So we needed another private
        # function to send `raw` data with only length prefixed

        # setup connection with mocked transport and protocol
        conn = AIOKafkaConnection(host="", port=9999)
        conn.close = mock.MagicMock()
        conn._writer = mock.MagicMock()
        out_buffer = []
        conn._writer.write = mock.Mock(side_effect=out_buffer.append)
        conn._reader = mock.MagicMock()
        self.assertEqual(len(conn._requests), 0)

        # Successful send
        fut = conn._send_sasl_token(b"Super data")
        self.assertEqual(b''.join(out_buffer), b"\x00\x00\x00\nSuper data")
        self.assertEqual(len(conn._requests), 1)
        out_buffer.clear()

        # Resolve the request
        conn._requests[0][2].set_result(None)
        conn._requests.clear()
        await fut

        # Broken pipe error
        conn._writer.write.side_effect = OSError
        with self.assertRaises(KafkaConnectionError):
            conn._send_sasl_token(b"<PASSWORD>")
        self.assertEqual(out_buffer, [])
        self.assertEqual(len(conn._requests), 0)
        self.assertEqual(conn.close.call_count, 1)

        conn._writer = None
        with self.assertRaises(KafkaConnectionError):
            conn._send_sasl_token(b"<PASSWORD>")
        # We don't need to close 2ce
        self.assertEqual(conn.close.call_count, 1)
