#!/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 re
import subprocess
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 _iface_to_ip(name):
    """Resolve a network interface name (e.g. ``"wlan0"``) to its first IPv4
    address via ``ip -4 addr show <name>``.  Returns ``"127.0.0.1"`` on failure."""
    try:
        out = subprocess.check_output(
            ["ip", "-4", "addr", "show", name],
            stderr=subprocess.DEVNULL, text=True)
        m = re.search(r"inet\s+(\d+\.\d+\.\d+\.\d+)", out)
        if m:
            return m.group(1)
    except (subprocess.CalledProcessError, FileNotFoundError, OSError):
        pass
    return "127.0.0.1"


def _mcast_interface_ip(group):
    """Return the IPv4 address of the OS-selected interface for a multicast
    group.  MARTe2's ``BasicUDPSocket::Join`` expects an IP address (passed to
    ``inet_addr``), not an interface name.

    Uses ``ip route get <group>`` to find the outgoing device, then resolves
    its IP.  Falls back to ``"127.0.0.1"`` if the route or device lookup fails.
    """
    try:
        out = subprocess.check_output(
            ["ip", "route", "get", group],
            stderr=subprocess.DEVNULL, text=True)
        m = re.search(r"dev\s+(\S+)", out)
        if m:
            return _iface_to_ip(m.group(1))
    except (subprocess.CalledProcessError, FileNotFoundError, OSError):
        pass
    return "127.0.0.1"


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)} }}"


SERVICE_NAME = "Svc"


def _service_port(scenario):
    """Port of the shared TCPStreamService (first source's port)."""
    return scenario["sources"][0]["udp_port"]


def _streamer_block(src, scenario):
    s = src["signals"]
    net = scenario["network"]
    if net == "tcp_service":
        # Port, MaxClients, ... belong to the service; the streamer only joins it.
        parts = [f'Service = "{SERVICE_NAME}"', f'StreamName = "{S.stream_name(scenario, src)}"']
    else:
        parts = [f"Port = {src['udp_port']}"]
    parts += [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']}")
        iface = src.get("interface") or _mcast_interface_ip(src["multicast_group"])
        parts.append(f'Interface = "{iface}"')
    sigs = " ".join(_streamer_sig(sig) for sig in s)
    cls = "UDPStreamer" if net in ("unicast", "multicast") else "TCPStreamer"
    return (f"        +Streamer_{src['id']} = {{ Class = {cls} "
            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)

    service_block = ""
    if scenario["network"] == "tcp_service":
        service_block = (f"+{SERVICE_NAME} = {{ Class = TCPStreamService "
                         f"Port = {_service_port(scenario)} Workers = 2 "
                         f"MaxClients = 16 MaxStreams = 16 }}\n")
        hub = scenario.get("hub")
        if hub and hub["mode"] == "join":
            names = ",".join(S.stream_name(scenario, x) for x in scenario["sources"])
            service_block += (f"+Join = {{ Class = TCPStreamJoin Service = \"{SERVICE_NAME}\" "
                              f"StreamName = \"{hub['name']}\" Sources = \"{names}\" "
                              f"Master = \"{hub['master']}\" }}\n")

    cfg = f"""{service_block}$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 = []
    hub = scenario.get("hub")
    if hub and hub["mode"] == "mux":
        names = ",".join(S.stream_name(scenario, x) for x in scenario["sources"])
        src_blocks.append(
            f'        {hub["node"]} = {{ Label = "chain" Addr = "127.0.0.1" '
            f'Port = {_service_port(scenario)} Transport = "tcp" StreamName = "{names}" }}')
    elif hub and hub["mode"] == "join":
        src_blocks.append(
            f'        {hub["name"]} = {{ Label = "chain {hub["name"]}" Addr = "127.0.0.1" '
            f'Port = {_service_port(scenario)} Transport = "tcp" StreamName = "{hub["name"]}" }}')
    for src in ([] if hub else scenario["sources"]):
        net = scenario["network"]
        port = _service_port(scenario) if net == "tcp_service" else src["udp_port"]
        parts = [f'Label = "chain {src["id"]}"', 'Addr = "127.0.0.1"',
                 f"Port = {port}"]
        if net in ("tcp", "tcp_service"):
            parts.append('Transport = "tcp"')
        if net == "tcp_service":
            parts.append(f'StreamName = "{src["id"]}"')
        if net == "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()