#!/usr/bin/env python3
"""
gen_cfg.py — MARTe2 app + StreamHub config generator for the streaming-chain E2E.

Given a scenario (scenarios.py) it produces two config files:

  * MARTe app cfg: ``LinuxTimer + FileReader(input) -> IOGAM -> UDPStreamer`` per
    source. When ``oracle in {fed,both}`` a second IOGAM branch taps the same fed
    signals into a ``FileWriter`` (the "fed reference").
  * StreamHub cfg: one ``Source`` per UDPStreamer (unicast or multicast) on the
    scenario's ``ws_port``.

The producer rate (LinuxTimer Frequency) defaults to 1 kHz; the FileReader uses
``EOF = "Rewind"`` so the NUM_ROWS-row input loops for the whole run. The
validator reconstructs truth as the cyclic ground-truth sequence, so the wrap is
expected, not an error.
"""
import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import scenarios as S  # noqa: E402

PRODUCER_HZ = 1000      # LinuxTimer frequency (Hz)


def _ndims(elements):
    return 0 if elements == 1 else 1


def _gam_sig(sig, datasource):
    """A GAM signal entry referencing a DataSource (no UDPStreamer extras)."""
    return _gam_sig_named(sig["name"], sig, datasource)


def _gam_sig_named(name, sig, datasource):
    """A GAM signal entry with an explicit (possibly renamed) signal name.

    IOGAM copies inputs to outputs positionally, so the input and output names
    may differ; we exploit that to route through the DDB with source-prefixed
    names that never clash with the TimerGAM's Counter/Time or across sources.
    """
    return (f"{name} = {{ Type = {sig['type']} "
            f"NumberOfDimensions = {_ndims(sig['elements'])} "
            f"NumberOfElements = {sig['elements']} DataSource = {datasource} }}")


def _streamer_sig(sig):
    """A UDPStreamer DataSource signal entry (carries time/quant/unit options)."""
    parts = [f"Type = {sig['type']}",
             f"NumberOfDimensions = {_ndims(sig['elements'])}",
             f"NumberOfElements = {sig['elements']}"]
    if sig["time_mode"] and sig["time_mode"] != "PacketTime":
        parts.append(f'TimeMode = "{sig["time_mode"]}"')
        if sig["time_signal"]:
            parts.append(f'TimeSignal = "{sig["time_signal"]}"')
        if sig["sampling_rate"]:
            parts.append(f"SamplingRate = {sig['sampling_rate']}")
    if sig["quant"] and sig["quant"] != "none":
        parts.append(f'QuantizedType = "{sig["quant"]}"')
        parts.append(f"RangeMin = {sig['range_min']}")
        parts.append(f"RangeMax = {sig['range_max']}")
    if sig["unit"]:
        parts.append(f'Unit = "{sig["unit"]}"')
    return f"{sig['name']} = {{ {' '.join(parts)} }}"


def _streamer_block(src, scenario):
    s = src["signals"]
    parts = [f"Port = {src['udp_port']}",
             f"MaxPayloadSize = {scenario['max_payload']}",
             f'PublishingMode = "{scenario["publishing"]}"']
    if scenario["publishing"] == "Accumulate":
        parts.append(f"MinRefreshRate = {scenario['min_refresh_hz']}")
    if scenario["publishing"] == "Decimate":
        parts.append(f"Ratio = {scenario['ratio']}")
    if scenario["network"] == "multicast":
        parts.append(f'MulticastGroup = "{src["multicast_group"]}"')
        parts.append(f"DataPort = {src['data_port']}")
    sigs = " ".join(_streamer_sig(sig) for sig in s)
    return (f"        +Streamer_{src['id']} = {{ Class = UDPStreamer "
            f"{' '.join(parts)} Signals = {{ {sigs} }} }}")


def write_marte_cfg(scenario, path, input_bin, tap_bin=None):
    """Write the MARTe app cfg. ``input_bin`` is the src[0] file; additional
    sources read ``<input_bin>.<srcid>`` (matching gen_data.write_input)."""
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    srcs = scenario["sources"]
    want_tap = scenario["oracle"] in ("fed", "both") and tap_bin is not None

    _row_dt, _num_rows, producer_hz, _loop_hz = S.geometry(scenario)

    gams = []          # +Functions entries
    datas = []         # +Data entries
    thread_funcs = ["TimerGAM"]

    # Timer GAM (drives the schedule)
    gams.append(
        "        +TimerGAM = { Class = IOGAM "
        "InputSignals = { Counter = { DataSource = ReaderTimer Type = uint32 } "
        f"Time = {{ Frequency = {producer_hz} DataSource = ReaderTimer Type = uint32 }} }} "
        "OutputSignals = { Counter = { DataSource = DDB Type = uint32 } "
        "Time = { DataSource = DDB Type = uint32 } } }")

    for i, src in enumerate(srcs):
        sid = src["id"]
        fpath = input_bin if i == 0 else f"{input_bin}.{sid}"
        rds = f"FileReaderDS_{sid}"
        # The FileReader DataSource allows exactly one consuming Function, so a
        # tapped source must route through the DDB: ReaderGAM copies FileReader
        # -> DDB (source-prefixed names), then StreamGAM and TapGAM both read DDB.
        tap_here = want_tap and i == 0
        if tap_here:
            in_sigs = " ".join(_gam_sig(sig, rds) for sig in src["signals"])
            ddb_out = " ".join(_gam_sig_named(f"{sid}_{sig['name']}", sig, "DDB")
                               for sig in src["signals"])
            gams.append(
                f"        +ReaderGAM_{sid} = {{ Class = IOGAM "
                f"InputSignals = {{ {in_sigs} }} OutputSignals = {{ {ddb_out} }} }}")
            ddb_in = " ".join(_gam_sig_named(f"{sid}_{sig['name']}", sig, "DDB")
                              for sig in src["signals"])
            stream_out = " ".join(_gam_sig(sig, f"Streamer_{sid}")
                                  for sig in src["signals"])
            gams.append(
                f"        +StreamGAM_{sid} = {{ Class = IOGAM "
                f"InputSignals = {{ {ddb_in} }} OutputSignals = {{ {stream_out} }} }}")
            thread_funcs.append(f"ReaderGAM_{sid}")
            thread_funcs.append(f"StreamGAM_{sid}")
        else:
            in_sigs = " ".join(_gam_sig(sig, rds) for sig in src["signals"])
            out_sigs = " ".join(_gam_sig(sig, f"Streamer_{sid}") for sig in src["signals"])
            gams.append(
                f"        +ReaderGAM_{sid} = {{ Class = IOGAM "
                f"InputSignals = {{ {in_sigs} }} OutputSignals = {{ {out_sigs} }} }}")
            thread_funcs.append(f"ReaderGAM_{sid}")
        datas.append(
            f'        +{rds} = {{ Class = FileReader Filename = "{fpath}" '
            f'Interpolate = "no" FileFormat = "binary" EOF = "Rewind" }}')
        datas.append(_streamer_block(src, scenario))

    if want_tap:
        # tap the first source's signals (now in the DDB) to a FileWriter.
        src = srcs[0]
        sid = src["id"]
        tap_in = " ".join(_gam_sig_named(f"{sid}_{sig['name']}", sig, "DDB")
                          for sig in src["signals"])
        tap_out = " ".join(_gam_sig(sig, "TapWriterDS") for sig in src["signals"])
        gams.append(
            f"        +TapGAM = {{ Class = IOGAM "
            f"InputSignals = {{ {tap_in} }} OutputSignals = {{ {tap_out} }} }}")
        thread_funcs.append("TapGAM")
        tap_sigs = " ".join(_gam_sig(sig, "TapWriterDS").replace(
            " DataSource = TapWriterDS", "") for sig in src["signals"])
        datas.append(
            f'        +TapWriterDS = {{ Class = FileWriter NumberOfBuffers = 10 '
            f'CPUMask = 0x4 StackSize = 10000000 Filename = "{tap_bin}" '
            f'Overwrite = "yes" StoreOnTrigger = 0 FileFormat = "binary" '
            f"Signals = {{ {tap_sigs} }} }}")

    funcs_block = "\n".join(gams)
    data_block = "\n".join(datas)
    thread_list = " ".join(thread_funcs)

    cfg = f"""$ChainE2E = {{
    Class = RealTimeApplication
    +Functions = {{
        Class = ReferenceContainer
{funcs_block}
    }}
    +Data = {{
        Class = ReferenceContainer DefaultDataSource = DDB
        +DDB = {{ Class = GAMDataSource }}
        +ReaderTimer = {{ Class = LinuxTimer SleepNature = "Default" Signals = {{ Counter = {{ Type = uint32 }} Time = {{ Type = uint32 }} }} }}
{data_block}
        +Timings = {{ Class = TimingDataSource }}
    }}
    +States = {{ Class = ReferenceContainer +Running = {{ Class = RealTimeState +Threads = {{ Class = ReferenceContainer +ReaderThread = {{ Class = RealTimeThread CPUs = 0x1 Functions = {{{thread_list}}} }} }} }} }}
    +Scheduler = {{ Class = GAMScheduler TimingDataSource = Timings }}
}}
"""
    with open(path, "w") as f:
        f.write(cfg)
    return path


def write_hub_cfg(scenario, path):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    src_blocks = []
    for src in scenario["sources"]:
        parts = [f'Label = "chain {src["id"]}"', 'Addr = "127.0.0.1"',
                 f"Port = {src['udp_port']}"]
        if scenario["network"] == "multicast":
            parts.append(f'MulticastGroup = "{src["multicast_group"]}"')
            parts.append(f"DataPort = {src['data_port']}")
        src_blocks.append(f"        {src['id']} = {{ {' '.join(parts)} }}")
    sources = "\n".join(src_blocks)
    cfg = f"""Hub = {{
    WSPort        = {scenario['ws_port']}
    MaxPoints     = 200000
    PushRate      = 30
    MaxPushPoints = 2000
    RingTemporal  = 1000000
    RingScalar    = 100000
    Sources = {{
{sources}
    }}
}}
"""
    with open(path, "w") as f:
        f.write(cfg)
    return path


def main():
    p = argparse.ArgumentParser(description="Generate MARTe + StreamHub cfgs")
    p.add_argument("--scenario", required=True)
    p.add_argument("--input", required=True, help="input_<id>.bin path")
    p.add_argument("--marte-out", required=True)
    p.add_argument("--hub-out", required=True)
    p.add_argument("--tap", default=None)
    args = p.parse_args()
    sc = next((s for s in S.SCENARIOS if s["id"] == args.scenario), None)
    if sc is None:
        print(f"unknown scenario {args.scenario}", file=sys.stderr)
        sys.exit(2)
    write_marte_cfg(sc, args.marte_out, args.input, args.tap)
    write_hub_cfg(sc, args.hub_out)
    print(f"marte: {args.marte_out}")
    print(f"hub:   {args.hub_out}")


if __name__ == "__main__":
    main()
