#!/usr/bin/python3

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

fn = sys.argv[1]

idx = index(fn)

config = parse(idx[0][3])

assert config['subtype'] == 'configuration'
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)
        
    
