#!/usr/bin/env python3

import struct
from ektools import index, parse
import sys
import os.path

fn = sys.argv[1]
idx = index(fn)

"""Looks like Datagram 0 contains general configuration,
   datagram 1 ('initialparameter') contains more specifics, both
   have elements that containg dicts with channel-named elements."""

config = parse(idx[0][3])
assert config['subtype'] == 'configuration'
# config2 = parse(idx[1][3])
# assert config2['subtype'] == 'initialparameter'

confs = set([k.split('_')[-1] for k in config['configuration'].keys()])
print('Configurations:', confs)

def matches(predicate, msg):
    try:
        p = parse(msg)
    except:
        return False
    if 'channel_id' in p.keys():
        return predicate(p['channel_id'].decode())
    elif 'parameter' in p.keys():
        return predicate(p['parameter']['channel_id'])
    else:
        return True

def dgram_write(f, dgram):
    """Write a datagram to a file"""
    hdr = struct.pack('<l', len(dgram))
    f.write(hdr)
    f.write(dgram)
    f.write(hdr)

import lxml.etree as et

def edit_xml(predicate, xmlstr):
    """Remove all except channel ch from XML configuration"""
    # remove all channels not matching ch
    xml = et.fromstring(xmlstr)
    for trans in xml.findall('.//Transceiver'):
        for ch in trans.findall('.//Channel'):
            if not predicate(ch.get('ChannelID')):
                ch.getparent().remove(ch)
        if len(trans.findall('.//Channel')) == 0:
            trans.getparent().remove(trans)

    return xml

split = {}
for a in confs:
    split[a] = [(p, t, l, m) for p, t, l, m in index(fn) if matches(lambda c: c.split('_')[-1] == a, m)]

for a in confs:
    print(f'Config {a}, count: {len(split[a])}')
    counts = {}
    for x in split[a]:
        if x[1] not in counts:
            counts[x[1]] = 1
        else:
            counts[x[1]] += 1
    print(counts)

    outname = os.path.splitext(os.path.basename(fn))[0] + '_' + a + '.raw'
    with open(outname, 'wb') as f:
        # Output configuration datagram
        xmlconf = idx[0][3][12:-2]
        x0 = et.tostring(edit_xml(lambda c: c.split('_')[-1] == a, xmlconf))
        dgram_write(f, split[a][0][3][:12] + x0)

        for (pos, typ, length, dgram) in split[a][1:]:
            # write length (dgram already contains type)
            dgram_write(f, dgram)
