import os
import unittest
from shutil import rmtree

import grpc
from jina.indexer.base import BaseIndexer
from jina.indexer.doc.filesys import DirectoryIndexer
from jina.peapods.indexer import IndexerService
from jina.peapods.preprocessor import PreprocessorService
from jina.preprocessor.base import BasePreprocessor

from jina.main.parser import set_frontend_parser, set_preprocessor_parser, set_indexer_parser
from jina.peapods.base import SocketType, Pod
from jina.peapods.frontend import FrontendService
from jina.proto import jina_pb2, jina_pb2_grpc, RequestGenerator


class TestDictIndexer(unittest.TestCase):
    def setUp(self):
        self.dirname = os.path.dirname(__file__)

        self.video_path = os.path.join(self.dirname, 'videos')
        self.video_bytes = [open(os.path.join(self.video_path, _), 'rb').read()
                            for _ in os.listdir(self.video_path)]

        self.pipeline_name = 'pipe-gif'
        self.pipeline_yml_path = os.path.join(self.dirname, 'yaml/%s.yml' % self.pipeline_name)
        self.data_path = './test_chunkleveldb'
        self.dump_path = os.path.join(self.dirname, 'indexer.bin')

        self.init_db()

    def test_pymode(self):
        os.unsetenv('http_proxy')
        os.unsetenv('https_proxy')
        args = set_frontend_parser().parse_args([
            '--dump_route', 'test.json'
        ])

        p_args = set_preprocessor_parser().parse_args([
            '--port_in', str(args.port_out),
            '--port_out', '5531',
            '--socket_in', str(SocketType.PULL_CONNECT),
            '--socket_out', str(SocketType.PUSH_BIND),
            '--yaml_path', 'SentSplitPreprocessor'
        ])

        e_args = set_indexer_parser().parse_args([
            '--port_in', str(p_args.port_out),
            '--port_out', str(args.port_in),
            '--socket_in', str(SocketType.PULL_CONNECT),
            '--socket_out', str(SocketType.PUSH_CONNECT),
            '--yaml_path', '!DictIndexer {jina_config: {name: dummy_dict_indexer}}',
        ])

        with Pod(IndexerService, e_args), \
             Pod(PreprocessorService, p_args), \
             FrontendService(args), \
             grpc.insecure_channel('%s:%s' % (args.grpc_host, args.grpc_port),
                                   options=[('grpc.max_send_message_length', 70 * 1024 * 1024),
                                            ('grpc.max_receive_message_length', 70 * 1024 * 1024)]) as channel:
            stub = jina_pb2_grpc.JinaRPCStub(channel)
            all_bytes = []
            with open(os.path.join(self.dirname, '26-doc-chinese.txt'), 'r', encoding='utf8') as fp:
                for v in fp:
                    if v.strip():
                        all_bytes.append(v.encode())
            for r in stub.StreamCall(RequestGenerator.index(all_bytes)):
                print(r)

        bi = BaseIndexer.load('dummy_dict_indexer.bin')
        self.assertEqual(bi.num_docs, 26)
        print(bi.query([0]))

    def tearDown(self):
        if os.path.exists(self.data_path):
            rmtree(self.data_path)
        if os.path.exists('dummy_dict_indexer.bin'):
            os.remove('dummy_dict_indexer.bin')

    def init_db(self):
        self.db = DirectoryIndexer(self.data_path)

        self.d = jina_pb2.Document()
        self.d.doc_id = 0
        self.d.raw_bytes = self.video_bytes[0]

        preprocess = BasePreprocessor.load_yaml(self.pipeline_yml_path)
        preprocess.apply(self.d)

        self.db.add(list(range(len(self.video_bytes))), [self.d])
        self.assertEqual(self.db.num_docs, len(self.video_bytes))

    def test_add_docs(self):
        # self.init_db()
        self.assertTrue(os.path.exists(os.path.join(self.data_path, str(self.d.doc_id))))
        self.assertEqual(len(self.d.chunks), len(os.listdir(os.path.join(self.data_path, str(self.d.doc_id)))) - 1)

    def test_query_docs(self):
        # self.init_db()

        query_list = [0, 1, 2]
        res = self.db.query(query_list)
        num_non_empty = sum(1 for d in res if d)
        self.assertEqual(num_non_empty, 1)
