dromap_editor

#!/usr/bin/env python

import daqconf.detreadoutmap as dromap
import detdataformats 


import click
import json
from rich import print
from rich.table import Table
from rich.console import Console


import click_shell

console = Console()

# ------------------------------------------------------------------------------
# Add -h as default help option
CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])
# ------------------------------------------------------------------------------

@click_shell.shell(prompt='dro> ', chain=True, context_settings=CONTEXT_SETTINGS)
@click.pass_obj
def cli(obj):
    pass


@cli.command('load', help="Load map from file")
@click.argument('path', type=click.Path(exists=True))
@click.option('-m', '--merge', is_flag=True, type=bool, default=False)
@click.option('-o', '--offset', type=int, default=0)
@click.pass_obj
def load(obj, path, merge, offset):
    m = obj
    m.load(path, merge, offset)

    console.print(m.as_table())

@cli.command('rm', help="Remove a stream")
@click.argument('src_id', type=int)
@click.pass_obj
def rm(obj, src_id):
    console.log(f"Removing {src_id}")
    m = obj
    m.remove_srcid(src_id)
    console.print(m.as_table())


@cli.command("add-flx", help="Add a felix stream")
@click.option('--force', type=bool, is_flag=True, default=False)
@click.option('--src-id', type=int, default=None)
@click.option('--geo-det-id', type=int, default=0)
@click.option('--geo-crate-id', type=int, default=0)
@click.option('--geo-slot-id', type=int, default=0)
@click.option('--geo-stream-id', type=int, default=0)
@click.option('--flx-host')
@click.option('--flx-protocol')
@click.option('--flx-mode')
@click.option('--flx-card')
@click.option('--flx-slr')
@click.option('--flx-link')
@click.pass_obj
def add_flx(obj, force, src_id, **kwargs):
    m = obj

    if src_id is None:
        src_id = max(m.get())+1 if m.get() else 0
    elif src_id in m.get_src_ids() and force:
        m.remove_srcid(src_id)
        
    console.log(f"Adding felix entry {src_id}")

    geo_args = {k.removeprefix('geo_'):v for k,v in kwargs.items() if k.startswith('geo_') and v is not None}
    flx_args = {k.removeprefix('flx_'):v for k,v in kwargs.items() if k.startswith('flx_') and v is not None}

    try:
        m.add_srcid(src_id, dromap.GeoID(**geo_args), 'flx', **flx_args )
    except ValueError:
        print(f"ERROR: {e}")
        raise click.ClickException(f"Failed to insert src {src_id}")
    print(m.as_table())

    
@cli.command("add-eth", help="Add an Ethernet stream")
@click.option('--force', type=bool, is_flag=True, default=False)
@click.option('--src-id', type=int, default=None)
@click.option('--geo-det-id', type=int, default=0)
@click.option('--geo-crate-id', type=int, default=0)
@click.option('--geo-slot-id', type=int, default=0)
@click.option('--geo-stream-id', type=int, default=0)
@click.option('--eth-protocol')
@click.option('--eth-mode')
@click.option('--eth-tx-host')
@click.option('--eth-tx-mac')
@click.option('--eth-tx-ip')
@click.option('--eth-rx-host')
@click.option('--eth-rx-mac')
@click.option('--eth-rx-pcie-dev')
@click.option('--eth-rx-ip')
@click.option('--eth-rx-iface')
@click.pass_obj
def add_eth(obj, force, src_id, **kwargs):
    m = obj


    if src_id is None:
        src_id = max(m.get())+1 if m.get() else 0
    elif src_id in m.get_src_ids() and force:
        m.remove_srcid(src_id)


    console.log(f"Adding felix entry {src_id}")

    geo_args = {k.removeprefix('geo_'):v for k,v in kwargs.items() if k.startswith('geo_') and v is not None}
    eth_args = {k.removeprefix('eth_'):v for k,v in kwargs.items() if k.startswith('eth_') and v is not None}

    try:
        m.add_srcid(src_id, dromap.GeoID(**geo_args), 'eth', **eth_args )
    except ValueError as e:
        print(f"ERROR: {e}")
        raise click.ClickException(f"Failed to insert src {src_id}: {e}")
    print(m.as_table())



@cli.command("add-np-wib-crate")
@click.argument('addrbook_path', type=click.Path(exists=True))
@click.argument('wib_filter', type=str)
@click.argument('ru_interface', type=str)
@click.option('--rx-iface', type=int, default=0, help="Interface id on the receiver host")
@click.pass_obj
def add_np_wib_crate(obj, addrbook_path, wib_filter, ru_interface, rx_iface):
    """Adds collections of wibs to the readout map and routes them to a destination"""
    m = obj

    with (open(addrbook_path, 'r')) as f:
        addrbook = json.load(f)

    import re

    wib_re = re.compile(wib_filter)
    wib_sources = { k:v.copy() for k,v in addrbook.items() if wib_re.match(k)}
    if not wib_sources:
        raise RuntimeError(f'No sources selected by {wib_filter}')
    
    for host in wib_sources:
        del wib_sources[host][host]

    

    ru_hosts = [k for k,v in addrbook.items() if ru_interface in v]
    if not ru_hosts:
        raise RuntimeError(f"Readout unit interface '{ru_interface}' not found")
    elif len(ru_hosts) > 1:
        raise RuntimeError(f"Readout unit interface '{ru_interface}' found on multiple hosts {ru_hosts}")

    ru_host = ru_hosts[0]
    ru_rx = addrbook[ru_host][ru_interface]

    # Start from the next available src id
    src_id = max(m.get())+1 if m.get() else 0
    
    # Constant
    link_stream_offset = 0x40

    for name, ifaces in wib_sources.items():

        # Recover detector, crate, slot from NP wib name
        name_tokens = name.split('-')
        print(name_tokens)

        match name_tokens[0]:
            case 'np04':
                det_id =  detdataformats.DetID.kHD_TPC
            case 'np02':
                det_id =   detdataformats.DetID.kVD_BottomTPC
            case other:
                raise ValueError(f'Detector {name_tokens[0]} Unknown')

        wib_id = int(name_tokens[-1])

        crate_id = (wib_id % 1000)//100
        slot_id = (wib_id % 100)-1

        for ifname, ifdata in ifaces.items():

            link = int(ifname.removeprefix(name+'-d'))
            if link not in (0,1):
                raise ValueError(f"Recovered link id {link} from {ifname} is not 0 or 1 as expected")

            for s in range(4):
                m.add_srcid(
                    src_id, 
                    dromap.GeoID(
                        det_id=det_id.value,
                        crate_id=crate_id,
                        slot_id=slot_id,
                        stream_id=link_stream_offset*link+s
                    ),
                    'eth',
                    protocol='udp',
                    mode='fix_rate',
                    tx_host=name,
                    tx_mac=ifdata['mac'],
                    tx_ip=ifdata['ip'],
                    rx_host=ru_host,
                    rx_mac=ru_rx['mac'],
                    rx_ip=ru_rx['ip'],
                    rx_iface=rx_iface,
                )
                src_id += 1

    print(m.as_table())


@cli.command("save", help="Save the map to json file")
@click.argument('path', type=click.Path())
@click.pass_obj
def save(obj, path):

    m = obj
    with open(path,"w") as f:
        json.dump(m.as_json(), f, indent=4)
    console.log(f"Map saved to '{path}'")
    

@cli.command("ipy", help="Start IPython")
@click.pass_obj
def ipy(obj):

    m = obj

    try:
        import IPython
        IPython.embed(colors="neutral")
    except ModuleNotFoundError:
        print("[red]Error: IPython is not installed[/red]")
        raise SystemExit(-1)

if __name__ == "__main__":
    cli(obj=dromap.DetReadoutMapService())