# -*- coding: utf-8 -*-

import timeit
from functools import partial
import csv
import argparse
import typing

import statistics
import bz2
import zlib
import snappy
import ujson as json
import msgpack
import cbor
from google.protobuf import json_format
from proto.data_pb2 import item 

from compressors import (
    compress_bz2,
    compress_zlib,
    compress_snappy,
    no_compression,
    decompress_bz2,
    decompress_zlib,
    decompress_snappy,
    no_decompression,
)
from protocols import (
    to_json,
    to_cbor,
    to_msgpack,
    from_json,
    from_cbor,
    from_msgpack,
    to_proto,
    from_proto,
)

def load(path):
    with open(path, "rb") as f:
        return json.load(f)
    raise Exception("failed to read file")

def get_proto_obj(data):
    content = json.dumps(data)
    obj = json_format.Parse(content, item(), ignore_unknown_fields=False)
    return obj


def process(fn1, fn2, data):
    return fn2(fn1(data))

def main(number: int, repeats: int, input_file: str, output_file: str):
    serilizer = {
        'json': to_json,
        'msgpack': to_msgpack,
        'cbor': to_cbor,
        'proto': to_proto,
    }

    deserializer = {
        'json': from_json,
        'msgpack': from_msgpack,
        'cbor': from_cbor,
        'proto': from_proto,
    }

    compression = {
        'none': no_compression,
        'bz2': compress_bz2,
        'zlib': compress_zlib,
        'snappy': compress_snappy,
    }

    decompression = {
        'none': no_decompression,
        'bz2': decompress_bz2,
        'zlib': decompress_zlib,
        'snappy': decompress_snappy,
    }
    data = load(input_file)
    reports:typing.Dict[str, typing.Dict] = dict()
    for proto_name, proto_fn in serilizer.items():
        reports[proto_name] = dict()
        for comp_name, comp_fn in compression.items():
            print(proto_name, comp_name)
            reports[proto_name][comp_name] = dict()
            if proto_name == "proto":
                # convert object to proto-representation
                proto_obj = get_proto_obj(data)

                reports[proto_name][comp_name]['serialize'] = timeit.repeat(
                    partial(process, proto_fn, comp_fn, proto_obj), 
                    number=number, repeat=repeats, globals=globals())

                serialized = process(proto_fn, comp_fn, proto_obj)

                reports[proto_name][comp_name]['size']  = len(serialized)
                reports[proto_name][comp_name]['deserialize'] = timeit.repeat(
                    partial(process, decompression[comp_name], deserializer[proto_name], serialized),
                    number=number, repeat=repeats, globals=globals())
            else:
                reports[proto_name][comp_name]['serialize'] = timeit.repeat(
                    partial(process, proto_fn, comp_fn, data), 
                    number=number, repeat=repeats, globals=globals())

                serialized = process(proto_fn, comp_fn, data)

                reports[proto_name][comp_name]['size']  = len(serialized)
                reports[proto_name][comp_name]['deserialize'] = timeit.repeat(
                    partial(process, decompression[comp_name], deserializer[proto_name], serialized),
                    number=number, repeat=repeats, globals=globals())
                
    with open(output_file, 'w', newline='') as f:
        writer = csv.writer(f, delimiter=',')
        writer.writerow(['protocol', 'compression', 'size', 'serialize', 'deserialize'])
        for proto_name in reports:
            for comp_name, report in reports[proto_name].items():
                writer.writerow([
                    proto_name,
                    comp_name,
                    report["size"],
                    statistics.median(report["serialize"]),
                    statistics.median(report["deserialize"]),
                ])
    print("report generated and placed in: {}".format(output_file))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--number", type=int, default=50, help="number of execution")
    parser.add_argument("-r", "--repeat", type=int, default=5, help="number of repeats")
    parser.add_argument("-i", "--input_path", type=str, required=True, help="path to dataset")
    parser.add_argument("-o", "--output_path", type=str, required=True, help="path to output")
    args = parser.parse_args()
    main(args.number, args.repeat, args.input_path, args.output_path)