#!/usr/bin/env python3
import os
from datetime import datetime
import subprocess
from argparse import ArgumentParser
from io import StringIO

import pandas as pd


def format_udp_pcap_to_csv(folder, sender_pcap_name, receiver_pcap_name, output=None):
    tmp = list()
    for pcap in [sender_pcap_name, receiver_pcap_name]:
        tshark_command = [
            "tshark",
            "-r",
            "{}{}".format(folder, pcap),
            "-O",
            "udp",
            "-T",
            "fields",
            "-e",
            "frame.time_epoch",
            "-e",
            "data.len",
            "-e",
            "data",
            "-E",
            "header=y",
            "-E",
            "separator=,",
        ]
        tshark_out = subprocess.check_output(tshark_command)
        tmp.append(
            pd.read_csv(
                StringIO(tshark_out.decode("utf-8")),
                converters=dict(data=lambda x: hash(x)),
            )
        )
    sender_df = tmp[0].rename(
        columns={"frame.time_epoch": "departure_time", "data": "data_hash", "data.len": "data.len_send"}
    )
    receiver_df = tmp[1].rename(
        columns={"frame.time_epoch": "arrival_time", "data": "data_hash", "data.len": "data.len_received"}
    )


    transmission_df = pd.merge(sender_df, receiver_df, on="data_hash", how="outer")
    transmission_df["delay"] = transmission_df.arrival_time - transmission_df.departure_time

    direction = "uplink"
    if "server" in sender_pcap_name and "client" in receiver_pcap_name:
        direction = "downlink"

    filename = sender_pcap_name.replace("_sender_", "_{}_".format(direction))
    filename = filename.replace("client_", "")
    filename = filename.replace("server_", "")
    filename = filename.replace(".pcap", ".csv")

    transmission_df.to_csv("{}{}".format(folder if output is None else output, filename))


if __name__ == "__main__":
    start_time = datetime.now()
    parser = ArgumentParser()

    parser.add_argument("-f", "--folder", required=True, help="Folder with pcap files.")
    parser.add_argument(
        "-o", "--output", required=False, help="Folder for created csv files."
    )
    parser.add_argument(
        "--name_contains", required=False, default=".pcap", help="Substring in filename."
    )
    parser.add_argument(
        "--use_subdirs", required=False, default=False, help="Use pcaps in subdirs."
    )

    args = parser.parse_args()

    sender_pcap_list = list()
    counter = 1

    if args.use_subdirs:
        sub_dirs = os.listdir(args.folder)
        for d in sub_dirs:
            filenames = os.listdir("{}{}".format(args.folder, d))
            for filename in filenames:
                if filename.endswith(".pcap") and args.name_contains in filename and "sender" in filename:
                    sender_pcap_list.append("{}/{}".format(d, filename))
    else:
        filenames = os.listdir(args.folder)
        for filename in filenames:
            if filename.endswith(".pcap") and args.name_contains in filename and "sender" in filename:
                sender_pcap_list.append(filename)

    for sender_pcap in sender_pcap_list:
        receiver_pcap = sender_pcap.replace("sender", "receiver")
        if "client" in receiver_pcap:
            receiver_pcap = receiver_pcap.replace("client", "server")
        else:
            receiver_pcap = receiver_pcap.replace("server", "client")

        print("\rProcessing file pair no: {} out of {}    ({}%)    ".format(
            counter,
            len(sender_pcap_list),
            round((counter / len(sender_pcap_list)) * 100, 2)
        ), end="")

        format_udp_pcap_to_csv(
            args.folder,
            sender_pcap,
            receiver_pcap,
            output=args.folder if args.output is None else args.output
        )
        counter += 1

    print("\nWrote CSV files to: {}".format(args.folder if args.output is None else args.output))
    print("Runtime: {}".format(datetime.now() - start_time))