# Copyright 2020 The FedLearner Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# coding: utf-8

import unittest
import os

import tensorflow.compat.v1 as tf
tf.enable_eager_execution()
from google.protobuf import text_format, timestamp_pb2
import tensorflow_io
from tensorflow.compat.v1 import gfile

from fedlearner.common import db_client
from fedlearner.common import common_pb2 as common_pb
from fedlearner.common import data_join_service_pb2 as dj_pb
from fedlearner.data_join import (
    data_block_manager, common, data_block_dumper,
    raw_data_manifest_manager, raw_data_visitor, visitor
)
from fedlearner.data_join.data_block_manager import DataBlockBuilder
from fedlearner.data_join.raw_data_iter_impl.tf_record_iter import TfExampleItem

class TestDataBlockDumper(unittest.TestCase):
    def setUp(self):
        data_source_f = common_pb.DataSource()
        data_source_f.data_source_meta.name = "milestone"
        data_source_f.data_source_meta.partition_num = 1
        data_source_f.output_base_dir = "./output-f"
        self.data_source_f = data_source_f
        if gfile.Exists(self.data_source_f.output_base_dir):
            gfile.DeleteRecursively(self.data_source_f.output_base_dir)
        data_source_l = common_pb.DataSource()
        data_source_l.data_source_meta.name = "milestone"
        data_source_l.data_source_meta.partition_num = 1
        data_source_l.output_base_dir = "./output-l"
        self.raw_data_dir_l = "./raw_data-l"
        self.data_source_l = data_source_l
        if gfile.Exists(self.data_source_l.output_base_dir):
            gfile.DeleteRecursively(self.data_source_l.output_base_dir)
        if gfile.Exists(self.raw_data_dir_l):
            gfile.DeleteRecursively(self.raw_data_dir_l)
        self.kvstore = db_client.DBClient('etcd', True)
        self.kvstore.delete_prefix(common.data_source_kvstore_base_dir(self.data_source_l.data_source_meta.name))
        self.manifest_manager = raw_data_manifest_manager.RawDataManifestManager(
            self.kvstore, self.data_source_l)

    def generate_follower_data_block(self):
        dbm = data_block_manager.DataBlockManager(self.data_source_f, 0)
        self.assertEqual(dbm.get_dumped_data_block_count(), 0)
        self.assertEqual(dbm.get_lastest_data_block_meta(), None)
        leader_index = 0
        follower_index = 65536
        self.dumped_metas = []
        for i in range(5):
            builder = DataBlockBuilder(
                    common.data_source_data_block_dir(self.data_source_f),
                    self.data_source_f.data_source_meta.name,
                    0, i, dj_pb.WriterOptions(output_writer='TF_RECORD'), None
                )
            builder.set_data_block_manager(dbm)
            for j in range(1024):
                feat = {}
                example_id = '{}'.format(i * 1024 + j).encode()
                feat['example_id'] = tf.train.Feature(
                        bytes_list=tf.train.BytesList(value=[example_id]))
                event_time = 150000000 + i * 1024 + j
                feat['event_time'] = tf.train.Feature(
                        int64_list=tf.train.Int64List(value=[event_time]))
                feat['leader_index'] = tf.train.Feature(
                        int64_list=tf.train.Int64List(value=[leader_index]))
                feat['follower_index'] = tf.train.Feature(
                        int64_list=tf.train.Int64List(value=[follower_index]))
                example = tf.train.Example(features=tf.train.Features(feature=feat))
                builder.append_item(TfExampleItem(example.SerializeToString()),
                                    leader_index, follower_index)
                leader_index += 3
                follower_index += 1
            meta = builder.finish_data_block()
            self.dumped_metas.append(meta)
        self.leader_start_index = 0
        self.leader_end_index = leader_index
        self.assertEqual(dbm.get_dumped_data_block_count(), 5)
        for (idx, meta) in enumerate(self.dumped_metas):
            self.assertEqual(dbm.get_data_block_meta_by_index(idx), meta)

    def generate_leader_raw_data(self):
        dbm = data_block_manager.DataBlockManager(self.data_source_l, 0)
        raw_data_dir = os.path.join(self.raw_data_dir_l, common.partition_repr(0))
        if gfile.Exists(raw_data_dir):
            gfile.DeleteRecursively(raw_data_dir)
        gfile.MakeDirs(raw_data_dir)
        rdm = raw_data_visitor.RawDataManager(self.kvstore, self.data_source_l, 0)
        block_index = 0
        builder = DataBlockBuilder(
                    self.raw_data_dir_l,
                    self.data_source_l.data_source_meta.name,
                    0, block_index, dj_pb.WriterOptions(output_writer='TF_RECORD'), None
            )
        process_index = 0
        start_index = 0
        for i in range(0, self.leader_end_index + 3):
            if (i > 0 and i % 2048 == 0) or (i == self.leader_end_index + 2):
                meta = builder.finish_data_block()
                if meta is not None:
                    ofname = common.encode_data_block_fname(
                            self.data_source_l.data_source_meta.name,
                            meta
                        )
                    fpath = os.path.join(raw_data_dir, ofname)
                    self.manifest_manager.add_raw_data(
                            0,
                            [dj_pb.RawDataMeta(file_path=fpath,
                                               timestamp=timestamp_pb2.Timestamp(seconds=3))],
                            False)
                    process_index += 1
                    start_index += len(meta.example_ids)
                block_index += 1
                builder = DataBlockBuilder(
                        self.raw_data_dir_l,
                        self.data_source_l.data_source_meta.name,
                        0, block_index, dj_pb.WriterOptions(output_writer='TF_RECORD'), None
                    )
            feat = {}
            pt = i + 1 << 30
            if i % 3 == 0:
                pt = i // 3
            example_id = '{}'.format(pt).encode()
            feat['example_id'] = tf.train.Feature(
                    bytes_list=tf.train.BytesList(value=[example_id]))
            event_time = 150000000 + i
            feat['event_time'] = tf.train.Feature(
                    int64_list=tf.train.Int64List(value=[event_time]))
            example = tf.train.Example(features=tf.train.Features(feature=feat))
            builder.append_item(TfExampleItem(example.SerializeToString()), i, i)
        fpaths = [os.path.join(raw_data_dir, f)
                    for f in gfile.ListDirectory(raw_data_dir)
                    if not gfile.IsDirectory(os.path.join(raw_data_dir, f))]
        for fpath in fpaths:
            if not fpath.endswith(common.DataBlockSuffix):
                gfile.Remove(fpath)
        
    def test_data_block_dumper(self):
        self.generate_follower_data_block()
        self.generate_leader_raw_data()
        dbd = data_block_dumper.DataBlockDumperManager(
                self.kvstore, self.data_source_l, 0,
                dj_pb.RawDataOptions(raw_data_iter='TF_RECORD', read_ahead_size=1<<20, read_batch_size=128),
                dj_pb.WriterOptions(output_writer='TF_RECORD')
            )
        self.assertEqual(dbd.get_next_data_block_index(), 0)
        for (idx, meta) in enumerate(self.dumped_metas):
            success, next_index = dbd.add_synced_data_block_meta(meta)
            self.assertTrue(success)
            self.assertEqual(next_index, idx + 1)
        self.assertTrue(dbd.need_dump())
        self.assertEqual(dbd.get_next_data_block_index(), len(self.dumped_metas))
        with dbd.make_data_block_dumper() as dumper:
            dumper()
        dbm_f = data_block_manager.DataBlockManager(self.data_source_f, 0)
        dbm_l = data_block_manager.DataBlockManager(self.data_source_l, 0)
        self.assertEqual(dbm_f.get_dumped_data_block_count(), len(self.dumped_metas))
        self.assertEqual(dbm_f.get_dumped_data_block_count(),
                            dbm_l.get_dumped_data_block_count())
        for (idx, meta) in enumerate(self.dumped_metas):
            self.assertEqual(meta.data_block_index, idx)
            self.assertEqual(dbm_l.get_data_block_meta_by_index(idx), meta)
            self.assertEqual(dbm_f.get_data_block_meta_by_index(idx), meta)
            meta_fpth_l = os.path.join(
                    common.data_source_data_block_dir(self.data_source_l),
                    common.partition_repr(0),
                    common.encode_data_block_meta_fname(
                        self.data_source_l.data_source_meta.name,
                        0, meta.data_block_index
                    )
                )
            mitr = tf.io.tf_record_iterator(meta_fpth_l)
            meta_l = text_format.Parse(next(mitr), dj_pb.DataBlockMeta())
            self.assertEqual(meta_l, meta)
            meta_fpth_f = os.path.join(
                    common.data_source_data_block_dir(self.data_source_f),
                    common.partition_repr(0),
                    common.encode_data_block_meta_fname(
                        self.data_source_f.data_source_meta.name,
                        0, meta.data_block_index
                    )
                )
            mitr = tf.io.tf_record_iterator(meta_fpth_f)
            meta_f = text_format.Parse(next(mitr), dj_pb.DataBlockMeta())
            self.assertEqual(meta_f, meta)
            data_fpth_l = os.path.join(
                    common.data_source_data_block_dir(self.data_source_l),
                    common.partition_repr(0),
                    common.encode_data_block_fname(
                        self.data_source_l.data_source_meta.name,
                        meta_l
                    )
                )
            for (iidx, record) in enumerate(tf.io.tf_record_iterator(data_fpth_l)):
                example = tf.train.Example()
                example.ParseFromString(record)
                feat = example.features.feature
                self.assertEqual(feat['example_id'].bytes_list.value[0],
                                 meta.example_ids[iidx])
            self.assertEqual(len(meta.example_ids), iidx + 1)
            data_fpth_f = os.path.join(
                    common.data_source_data_block_dir(self.data_source_f),
                    common.partition_repr(0),
                    common.encode_data_block_fname(
                        self.data_source_l.data_source_meta.name,
                        meta_f
                    )
                )
            for (iidx, record) in enumerate(tf.io.tf_record_iterator(data_fpth_f)):
                example = tf.train.Example()
                example.ParseFromString(record)
                feat = example.features.feature
                self.assertEqual(feat['example_id'].bytes_list.value[0],
                                 meta.example_ids[iidx])
            self.assertEqual(len(meta.example_ids), iidx +1)

    def tearDown(self):
        if gfile.Exists(self.data_source_f.output_base_dir):
            gfile.DeleteRecursively(self.data_source_f.output_base_dir)
        if gfile.Exists(self.data_source_l.output_base_dir):
            gfile.DeleteRecursively(self.data_source_l.output_base_dir)
        if gfile.Exists(self.raw_data_dir_l):
            gfile.DeleteRecursively(self.raw_data_dir_l)
        self.kvstore.delete_prefix(common.data_source_kvstore_base_dir(self.data_source_l.data_source_meta.name))

if __name__ == '__main__':
    unittest.main()
