import sys
from typing import Sequence

import pytest

from jina import Request, QueryLang, Document
from jina.clients.request import request_generator
from jina.proto import jina_pb2
from jina.proto.jina_pb2 import EnvelopeProto
from jina.types.message import Message
from jina.types.request import _trigger_fields
from tests import random_docs


@pytest.mark.parametrize('field', _trigger_fields.difference({'command', 'args', 'flush'}))
def test_lazy_access(field):
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert not r.is_used

        # access r.train
        print(getattr(r, field))

        # now it is read
        assert r.is_used


def test_multiple_access():
    reqs = [Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10))]
    for r in reqs:
        assert not r.is_used
        assert r
        assert not r.is_used

    for r in reqs:
        assert not r.is_used
        assert r.index
        assert r.is_used


def test_lazy_nest_access():
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert not r.is_used
        # write access r.train
        r.docs[0].id = '1' * 16
        # now it is read
        assert r.is_used
        assert r.index.docs[0].id == '1' * 16


def test_lazy_change_message_type():
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert not r.is_used
        # write access r.train
        r.control.command = jina_pb2.RequestProto.ControlRequestProto.IDLE
        # now it is read
        assert r.is_used
        assert len(r.index.docs) == 0


def test_lazy_append_access():
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert not r.is_used
        r.request_type = 'index'
        # write access r.train
        r.docs.append(Document())
        # now it is read
        assert r.is_used


def test_lazy_clear_access():
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert not r.is_used
        # write access r.train
        r.ClearField('index')
        # now it is read
        assert r.is_used


def test_lazy_nested_clear_access():
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert not r.is_used
        # write access r.train
        r.index.ClearField('docs')
        # now it is read
        assert r.is_used


def test_lazy_msg_access():
    reqs = [Message(None, r.SerializeToString(), 'test', '123',
                    request_id='123', request_type='IndexRequest') for r in request_generator(random_docs(10))]
    for r in reqs:
        assert not r.request.is_used
        assert r.envelope
        assert len(r.dump()) == 3
        assert not r.request.is_used

    for r in reqs:
        assert not r.request.is_used
        assert r.request
        assert len(r.dump()) == 3
        assert not r.request.is_used

    for r in reqs:
        assert not r.request.is_used
        assert r.request.index.docs
        assert len(r.dump()) == 3
        assert r.request.is_used


def test_message_size():
    reqs = [Message(None, r, 'test', '123') for r in request_generator(random_docs(10))]
    for r in reqs:
        assert r.size == 0
        assert sys.getsizeof(r.envelope.SerializeToString())
        assert sys.getsizeof(r.request.SerializeToString())
        assert len(r.dump()) == 3
        assert r.size > sys.getsizeof(r.envelope.SerializeToString()) \
               + sys.getsizeof(r.request.SerializeToString())


def test_lazy_request_fields():
    reqs = (Request(r.SerializeToString(), EnvelopeProto()) for r in request_generator(random_docs(10)))
    for r in reqs:
        assert list(r.DESCRIPTOR.fields_by_name.keys())


def test_request_extend_queryset():
    q1 = {'name': 'SliceQL', 'parameters': {'start': 3, 'end': 4}}
    q2 = QueryLang({'name': 'SliceQL', 'parameters': {'start': 3, 'end': 4}, 'priority': 1})
    q3 = jina_pb2.QueryLangProto()
    q3.name = 'SliceQL'
    q3.parameters['start'] = 3
    q3.parameters['end'] = 4
    q3.priority = 2
    r = Request()
    r.queryset.extend([q1, q2, q3])
    assert isinstance(r.queryset, Sequence)
    assert len(r.queryset) == 3
    for idx, q in enumerate(r.queryset):
        assert q.priority == idx
        assert q.parameters['start'] == 3
        assert q.parameters['end'] == 4

    # q1 and q2 refer to the same
    assert len({id(q) for q in r.queryset}) == 2

    r2 = Request()
    r2.queryset.extend(r.queryset)
    assert len({id(q) for q in r2.queryset}) == 2

    r = Request()
    r.queryset.append(q1)
    r.queryset.append(q2)
    r.queryset.append(q3)
    for idx, q in enumerate(r.queryset):
        assert q.priority == idx
        assert q.parameters['start'] == 3
        assert q.parameters['end'] == 4

    with pytest.raises(TypeError):
        r.queryset.extend(1)


@pytest.mark.parametrize('typ,pb_typ', [('train', jina_pb2.RequestProto.TrainRequestProto),
                                        ('index', jina_pb2.RequestProto.IndexRequestProto),
                                        ('search', jina_pb2.RequestProto.SearchRequestProto),
                                        ('control', jina_pb2.RequestProto.ControlRequestProto)])
def test_empty_request_type(typ, pb_typ):
    r = Request()
    assert r.request_type is None
    with pytest.raises(ValueError):
        print(r.body)

    r.request_type = typ
    assert r._request_type == typ
    assert isinstance(r.body, pb_typ)


@pytest.mark.parametrize('typ,pb_typ', [('index', jina_pb2.RequestProto.IndexRequestProto),
                                        ('search', jina_pb2.RequestProto.SearchRequestProto)])
def test_add_doc_to_type(typ, pb_typ):
    r = Request()
    r.request_type = typ
    for _ in range(10):
        r.docs.append(Document())
        r.groundtruths.append(Document())
    assert len(r.docs) == 10
    assert len(r.groundtruths) == 10

[((5195, 5461), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""typ,pb_typ"""', "[('train', jina_pb2.RequestProto.TrainRequestProto), ('index', jina_pb2.\n    RequestProto.IndexRequestProto), ('search', jina_pb2.RequestProto.\n    SearchRequestProto), ('control', jina_pb2.RequestProto.ControlRequestProto)\n    ]"], {}), "('typ,pb_typ', [('train', jina_pb2.RequestProto.\n    TrainRequestProto), ('index', jina_pb2.RequestProto.IndexRequestProto),\n    ('search', jina_pb2.RequestProto.SearchRequestProto), ('control',\n    jina_pb2.RequestProto.ControlRequestProto)])\n", (5218, 5461), False, 'import pytest\n'), ((5822, 5972), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""typ,pb_typ"""', "[('index', jina_pb2.RequestProto.IndexRequestProto), ('search', jina_pb2.\n    RequestProto.SearchRequestProto)]"], {}), "('typ,pb_typ', [('index', jina_pb2.RequestProto.\n    IndexRequestProto), ('search', jina_pb2.RequestProto.SearchRequestProto)])\n", (5845, 5972), False, 'import pytest\n'), ((382, 438), 'jina.types.request._trigger_fields.difference', '_trigger_fields.difference', (["{'command', 'args', 'flush'}"], {}), "({'command', 'args', 'flush'})\n", (408, 438), False, 'from jina.types.request import _trigger_fields\n'), ((4176, 4263), 'jina.QueryLang', 'QueryLang', (["{'name': 'SliceQL', 'parameters': {'start': 3, 'end': 4}, 'priority': 1}"], {}), "({'name': 'SliceQL', 'parameters': {'start': 3, 'end': 4},\n    'priority': 1})\n", (4185, 4263), False, 'from jina import Request, QueryLang, Document\n'), ((4269, 4294), 'jina.proto.jina_pb2.QueryLangProto', 'jina_pb2.QueryLangProto', ([], {}), '()\n', (4292, 4294), False, 'from jina.proto import jina_pb2\n'), ((4407, 4416), 'jina.Request', 'Request', ([], {}), '()\n', (4414, 4416), False, 'from jina import Request, QueryLang, Document\n'), ((4779, 4788), 'jina.Request', 'Request', ([], {}), '()\n', (4786, 4788), False, 'from jina import Request, QueryLang, Document\n'), ((4883, 4892), 'jina.Request', 'Request', ([], {}), '()\n', (4890, 4892), False, 'from jina import Request, QueryLang, Document\n'), ((5619, 5628), 'jina.Request', 'Request', ([], {}), '()\n', (5626, 5628), False, 'from jina import Request, QueryLang, Document\n'), ((6055, 6064), 'jina.Request', 'Request', ([], {}), '()\n', (6062, 6064), False, 'from jina import Request, QueryLang, Document\n'), ((3436, 3467), 'jina.types.message.Message', 'Message', (['None', 'r', '"""test"""', '"""123"""'], {}), "(None, r, 'test', '123')\n", (3443, 3467), False, 'from jina.types.message import Message\n'), ((5137, 5161), 'pytest.raises', 'pytest.raises', (['TypeError'], {}), '(TypeError)\n', (5150, 5161), False, 'import pytest\n'), ((5672, 5697), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (5685, 5697), False, 'import pytest\n'), ((512, 527), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (525, 527), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((805, 820), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (818, 820), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((1132, 1147), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (1145, 1147), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((1483, 1498), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (1496, 1498), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((1863, 1878), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (1876, 1878), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((2059, 2069), 'jina.Document', 'Document', ([], {}), '()\n', (2067, 2069), False, 'from jina import Request, QueryLang, Document\n'), ((2196, 2211), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (2209, 2211), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((2499, 2514), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (2512, 2514), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((3925, 3940), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ([], {}), '()\n', (3938, 3940), False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((6136, 6146), 'jina.Document', 'Document', ([], {}), '()\n', (6144, 6146), False, 'from jina import Request, QueryLang, Document\n'), ((6178, 6188), 'jina.Document', 'Document', ([], {}), '()\n', (6186, 6188), False, 'from jina import Request, QueryLang, Document\n'), ((556, 571), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (567, 571), False, 'from tests import random_docs\n'), ((849, 864), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (860, 864), False, 'from tests import random_docs\n'), ((1176, 1191), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (1187, 1191), False, 'from tests import random_docs\n'), ((1527, 1542), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (1538, 1542), False, 'from tests import random_docs\n'), ((1907, 1922), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (1918, 1922), False, 'from tests import random_docs\n'), ((2240, 2255), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (2251, 2255), False, 'from tests import random_docs\n'), ((2543, 2558), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (2554, 2558), False, 'from tests import random_docs\n'), ((2913, 2928), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (2924, 2928), False, 'from tests import random_docs\n'), ((3495, 3510), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (3506, 3510), False, 'from tests import random_docs\n'), ((3969, 3984), 'tests.random_docs', 'random_docs', (['(10)'], {}), '(10)\n', (3980, 3984), False, 'from tests import random_docs\n')]