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

[((170, 1, 173, 96), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(170, 25, 170, 37): '"""typ,pb_typ"""', (170, 39, 173, 95): "[('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)])", False, 'import pytest\n'), ((185, 1, 186, 94), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(185, 25, 185, 37): '"""typ,pb_typ"""', (185, 39, 186, 93): "[('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)])", False, 'import pytest\n'), ((15, 34, 15, 90), 'jina.types.request._trigger_fields.difference', '_trigger_fields.difference', ({(15, 61, 15, 89): "{'command', 'args', 'flush'}"}, {}), "({'command', 'args', 'flush'})", False, 'from jina.types.request import _trigger_fields\n'), ((135, 9, 135, 92), 'jina.QueryLang', 'QueryLang', ({(135, 19, 135, 91): "{'name': 'SliceQL', 'parameters': {'start': 3, 'end': 4}, 'priority': 1}"}, {}), "({'name': 'SliceQL', 'parameters': {'start': 3, 'end': 4},\n    'priority': 1})", False, 'from jina import Request, QueryLang, Document\n'), ((136, 9, 136, 34), 'jina.proto.jina_pb2.QueryLangProto', 'jina_pb2.QueryLangProto', ({}, {}), '()', False, 'from jina.proto import jina_pb2\n'), ((141, 8, 141, 17), 'jina.Request', 'Request', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((153, 9, 153, 18), 'jina.Request', 'Request', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((157, 8, 157, 17), 'jina.Request', 'Request', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((175, 8, 175, 17), 'jina.Request', 'Request', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((188, 8, 188, 17), 'jina.Request', 'Request', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((117, 12, 117, 43), 'jina.types.message.Message', 'Message', ({(117, 20, 117, 24): 'None', (117, 26, 117, 27): 'r', (117, 29, 117, 35): '"""test"""', (117, 37, 117, 42): '"""123"""'}, {}), "(None, r, 'test', '123')", False, 'from jina.types.message import Message\n'), ((166, 9, 166, 33), 'pytest.raises', 'pytest.raises', ({(166, 23, 166, 32): 'TypeError'}, {}), '(TypeError)', False, 'import pytest\n'), ((177, 9, 177, 34), 'pytest.raises', 'pytest.raises', ({(177, 23, 177, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n'), ((17, 43, 17, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((29, 43, 29, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((42, 43, 42, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((53, 43, 53, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((64, 43, 64, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((69, 22, 69, 32), 'jina.Document', 'Document', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((75, 43, 75, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((85, 43, 85, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((128, 43, 128, 58), 'jina.proto.jina_pb2.EnvelopeProto', 'EnvelopeProto', ({}, {}), '()', False, 'from jina.proto.jina_pb2 import EnvelopeProto\n'), ((191, 22, 191, 32), 'jina.Document', 'Document', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((192, 30, 192, 40), 'jina.Document', 'Document', ({}, {}), '()', False, 'from jina import Request, QueryLang, Document\n'), ((17, 87, 17, 102), 'tests.random_docs', 'random_docs', ({(17, 99, 17, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((29, 87, 29, 102), 'tests.random_docs', 'random_docs', ({(29, 99, 29, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((42, 87, 42, 102), 'tests.random_docs', 'random_docs', ({(42, 99, 42, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((53, 87, 53, 102), 'tests.random_docs', 'random_docs', ({(53, 99, 53, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((64, 87, 64, 102), 'tests.random_docs', 'random_docs', ({(64, 99, 64, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((75, 87, 75, 102), 'tests.random_docs', 'random_docs', ({(75, 99, 75, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((85, 87, 85, 102), 'tests.random_docs', 'random_docs', ({(85, 99, 85, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((96, 94, 96, 109), 'tests.random_docs', 'random_docs', ({(96, 106, 96, 108): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((117, 71, 117, 86), 'tests.random_docs', 'random_docs', ({(117, 83, 117, 85): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n'), ((128, 87, 128, 102), 'tests.random_docs', 'random_docs', ({(128, 99, 128, 101): '(10)'}, {}), '(10)', False, 'from tests import random_docs\n')]