#!/usr/bin/env python3
# Copyright 2019 Oticon A/S
# SPDX-License-Identifier: Apache-2.0

import time
import argparse
import struct
from csv_common import *

# Wireshark constants
WS_FLAGS_PHY_1M = 0b00<<14
WS_FLAGS_PHY_2M = 0b01<<14
WS_FLAGS_PHY_CODED = 0b10<<14
WS_FLAGS_SIGNAL_POWER = 0b1<<1

# BabbleSim constants (see ext_2G4_libPhyComv1/src/bs_pc_2G4_modulations.h)
P2G4_MOD_BLE = 0x10
P2G4_MOD_BLE2M = 0x20
P2G4_MOD_BLE_CODED = 0x50

# default max packet length (phdr + access address + pdu + crc)
SNAPLEN = 512

def get_data_and_len(row):
	pdu_crc = bytes.fromhex("00")

	orig_len = int(row['packet_size'], 10)

	if orig_len != 0:
		try:
			pdu_crc = bytes.fromhex(row['packet'])
		except: #In case the packet is broken mid byte
			pdu_crc = bytes.fromhex("00")

		if len(pdu_crc) != orig_len:  # Let's handle this somehow gracefully
			print("Truncated input file (partial packet), writing partial packet in output")
			orig_len = len(pdu_crc)
		orig_len += 14; # 10 bytes phdr + 4 bytes access address

	return (orig_len, pdu_crc)

def is_coded_FEC2(row):
	if not row:
		return False
	if int(row['modulation'], 10) != P2G4_MOD_BLE_CODED:
		return False
	if int(row['packet_size'], 10) < 5: # It must at least have a header and CRC (2+3 bytes)
		return False
	return True

def write(outfile, *inputs, snaplen = SNAPLEN, basetime=None):
	#For information on the pcap format see https://wiki.wireshark.org/Development/LibpcapFileFormat

	buf = bytearray(30)

	if basetime == None:
		basetime = int(time.time() * 1000000)

	# write pcap header
	struct.pack_into('<IHHiIII', buf, 0,
			0xa1b2c3d4, # magic_number
			2,          # version_major
			4,          # version_minor
			0,          # thiszone
			0,          # sigfigs
			snaplen,    # snaplen
			256)        # network, 256 = BLUETOOTH_LE_LL_WITH_PHDR
	outfile.write(buf[:24])

	inputs = [ CSVFile(f) for f in inputs ]
	rows = [ next(cf, None) for cf in inputs ]

	while True:
		min_ts = None
		min_idx = None
		for idx, row in enumerate(rows):
			if not row:
				continue
			ts = row['start_time']
			if min_ts == None or ts < min_ts:
				min_ts = ts
				min_idx = idx

		if min_idx == None:
			break

		row = rows[min_idx]
		rows[min_idx] = next(inputs[min_idx], None)

		(orig_len, pdu_crc) = get_data_and_len(row)

		freq = float(row['center_freq'])
		if freq >= 1.0 and freq < 81.0:
			rf_channel = int((freq - 1.0) / 2)
		elif freq >= 2401.0 and freq < 2481.0:
			rf_channel = int((freq - 2401.0) / 2)
		else:
			raise ValueError

		access_address = int(row['phy_address'], 16)

		flags = 0
		modulation = int(row['modulation'], 10)
		if modulation == P2G4_MOD_BLE:
			flags |= WS_FLAGS_PHY_1M
		elif modulation == P2G4_MOD_BLE2M:
			flags |= WS_FLAGS_PHY_2M
		elif modulation == P2G4_MOD_BLE_CODED:
			flags |= WS_FLAGS_PHY_CODED
			if (int(row['packet_size'], 10) != 1): # If this is a FEC1, it has only the CI byte
				continue

			# The next row likely contains this packet FEC2:
			row = rows[min_idx]
			if not is_coded_FEC2(row):
				# Otherwise the coded phy packet was aborted before the FEC2, so we just ignore it
				continue

			rows[min_idx] = next(inputs[min_idx], None)

			(orig_len, pdu_crc_FEC2) = get_data_and_len(row)

			pdu_crc = bytes(pdu_crc) + pdu_crc_FEC2
			orig_len += 1

		# Transmission power (dBm)
		flags |= WS_FLAGS_SIGNAL_POWER
		signal_power = int(float(row['power_level']))

		incl_len = min(orig_len, snaplen)

		ts = basetime + min_ts
		ts_sec = ts // 1000000
		ts_usec = ts % 1000000

		struct.pack_into('<IIIIBbbBIHI', buf, 0,
				# pcap record header, 16 bytes
				ts_sec,
				ts_usec,
				incl_len,
				orig_len,
				# packet data, incl_len bytes
				# - phdr, 10 bytes
				rf_channel,
				signal_power,
				0, # noise power
				0, # access address offenses
				0, # reference access address
				flags,
				# - le packet (access address + pdu + crc, no preamble)
				access_address)
		outfile.write(buf)
		outfile.write(pdu_crc[:(incl_len - 14)])

def parse_args():
	parser = argparse.ArgumentParser(description='Convert BabbleSim Phy csv files to pcap')
	parser.add_argument(
			'-er', '--epoch_real',
			action='store_true',
			dest='epoch_real',
			required=False,
			default=True,
			help='If set, the pcap timestamps will be offset based on the host time when the simulation was run. Otherwise they are directly the simulation timestamps (default)',
			)
	parser.add_argument(
			'-es', '--epoch_simu',
			action='store_false',
			dest='epoch_real',
			required=False,
			help='If set, pcap timestamp are directly the simulation timestamps',
			)
	parser.add_argument('-o', '--output',
			dest='output',
			metavar='OUTFILE',
			help='Write to this pcap file (required)',
			required=True,
			type=argparse.FileType(mode='wb'))
	parser.add_argument(
			dest='inputs',
			metavar='INFILE',
			help='Input csv file(s) (at least one is required)',
			nargs='+',
			type=open_input)
	return parser.parse_args()

args = parse_args()

if (args.epoch_real == True):
	basetime = int(min([ p[0] for p in args.inputs ]) * 1000000)
else:
	basetime = 0

write(args.output, *[ p[1] for p in args.inputs ], basetime=basetime)

# vim: set ts=4 sw=4 noet: