#!/usr/bin/python3
# @ingroup ifw-rtmslib_receiver
# @copyright
#   SPDX-FileCopyrightText: 2026 European Southern Observatory (ESO)
#   SPDX-License-Identifier: LGPL-3.0-only

"""
Simple RTMS receiver using the ifw.rtmslib Python bindings.

Example:
  rtmslibReceiver --receiver-ip 239.10.0.1 --interface eth0 --port 50000 --data-type Int16 --count 5
"""

import argparse
import sys
import logging
import time
import math

import ifw.rtmslib as rtmslib

DATA_TYPES = {
    "Int8": 64,       # BYTE
    "Int16": 256,
    "UInt16": 512,
    "Int32": 1024,
    "UInt32": 2048,
    "Float": 65536,
    "Double": 131072,
}


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="RTMS receiver (Python bindings)")
    p.add_argument("-a", "--address", required=True, help="RTMS source IP (unicast or multicast).")
    p.add_argument("-n", "--netif", required=True, help="Network interface to use (for multicast).")
    p.add_argument("-p", "--port", type=int, required=True, help="UDP port to listen on.")
    p.add_argument("-w", "--width", type=int, help="Image width (optional; can be inferred from ext-info).")
    p.add_argument("-e", "--height", type=int, help="Image height (optional; can be inferred from ext-info).")
    p.add_argument("--samples", type=int, default=0, help="Number of samples to receive (0 = run until Ctrl-C).")
    p.add_argument("--topic-id", type=int, default=10, help="RTMS topic identifier.")
    p.add_argument("-m", "--mtu-size", type=int, default=1500, help="MTU size (1500 or 9000 bytes).")
    p.add_argument("-x", "--ext-info", choices=["none", "timestamp", "basic", "TS", "BII"],
                   default="basic", help="Expected extended info layout.")
    p.add_argument("--ignore-checksum", action="store_true",
                   help="Ignore MUDPI checksum mismatches.")
    p.add_argument("-t", "--data-type", choices=DATA_TYPES.keys(), default="Int16",
                   help="Image data type: Byte/Int16/UInt16/Int32/UInt32/Float/Double.")
    p.add_argument("--timeout", type=float, default=2.0, help="Receive timeout in seconds.")
    p.add_argument("--total-timeout", type=float, default=0.0,
                   help="Total time limit in seconds (0 = no limit). Exit after this time even if not all samples received.")
    p.add_argument("-l", "--log-level", default="INFO", help="Log level (TRACE/DEBUG/INFO/WARN/ERROR/FATAL).")
    return p.parse_args()


def extinfo_mode(name: str) -> rtmslib.ExtendedInfoMode:
    name = name.lower()
    if name == "none":
        return getattr(rtmslib.ExtendedInfoMode, "None")
    if name in ("timestamp", "ts"):
        return rtmslib.ExtendedInfoMode.TimestampOnly
    return rtmslib.ExtendedInfoMode.BasicImageInfo


def main() -> None:
    args = parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s %(levelname)s: %(message)s",
    )

    cfg = rtmslib.ReceiverConfig()
    cfg.ext_info_mode = extinfo_mode(args.ext_info)
    cfg.topic_id = args.topic_id
    cfg.mtu_size = args.mtu_size
    cfg.checksum_force_ignore = args.ignore_checksum
    cfg.rxbuf_size = 64 * 1024 * 1024  # 64 MB receive buffer to avoid packet loss
    if args.width:
        cfg.width = args.width
    if args.height:
        cfg.height = args.height

    datatype_code = DATA_TYPES[args.data_type]
    logging.info("Creating receiver with datatype_code=%s", datatype_code)
    receiver = rtmslib.RtmsReceiver(
        receiver_ip=args.address,
        interface_name=args.netif,
        port=args.port,
        data_type=datatype_code,
        rcv_cfg=cfg,
    )

    # Force socket initialization by doing a quick receive (socket is created lazily)
    try:
        logging.info("Forcing socket initialization...")
        receiver.receive_sample(0.001)
        logging.info("Socket initialized (no error)")
    except Exception as e:
        logging.info("Socket initialization result: %s (expected timeout)", e)

    # Signal that we are ready to receive (for test synchronization)
    print(f"READY: Receiver bound to port {args.port}", flush=True)
    logging.info("Receiver config: addr=%s, netif=%s, port=%s, datatype=%s, topic=%s, mtu=%s, ext=%s",
                 args.address, args.netif, args.port, args.data_type, args.topic_id, args.mtu_size, args.ext_info)

    received = 0
    failed = 0
    total_bytes = 0
    last_info = time.time()
    last_error = 0.0
    start_time = last_info
    prev_ts = None
    mean_interval = 0.0
    m2_interval = 0.0
    last_extinfo = None
    total_deadline = start_time + args.total_timeout if args.total_timeout > 0 else None
    logging.info("Starting receive loop, expecting %s samples, total_timeout=%s",
                 args.samples, args.total_timeout if args.total_timeout > 0 else "none")
    while args.samples == 0 or received < args.samples:
        # Check total timeout
        if total_deadline is not None and time.time() >= total_deadline:
            logging.info("Total timeout reached after %.1fs, exiting with %s samples received",
                         time.time() - start_time, received)
            break
        try:
            ok, ext_info, payload = receiver.receive_sample(args.timeout)
        except Exception as exc:  # keep the loop alive on transient receive errors
            failed += 1
            now = time.time()
            if last_error == 0.0 or (now - last_error) >= 2.0:
                logging.error("Receive failed: %s", exc)
                last_error = now
            continue
        if not ok:
            failed += 1
            now = time.time()
            if last_error == 0.0 or (now - last_error) >= 10.0:
                logging.warning("Timeout waiting for sample (failed=%s)", failed)
                last_error = now
            continue
        logging.info("Received sample %s: %s bytes, ts=%.6f", received, len(payload), ext_info.timestamp)
        received += 1
        total_bytes += len(payload)
        now = time.time()
        if prev_ts is not None:
            interval = now - prev_ts
            n = received - 1  # intervals count
            delta = interval - mean_interval
            mean_interval += delta / n
            m2_interval += delta * (interval - mean_interval)
        prev_ts = now
        logging.debug(
            "Sample %s: ts=%.6f, %sx%s, bytes=%s, datatype=%s",
            received,
            ext_info.timestamp,
            ext_info.width,
            ext_info.height,
            len(payload),
            ext_info.datatype,
        )
        last_extinfo = ext_info
        if now - last_info >= 10.0:
            elapsed = now - start_time
            rate = received / elapsed if elapsed > 0 else 0.0
            intervals = max(received - 1, 1)
            stddev = math.sqrt(m2_interval / intervals) if intervals > 1 else 0.0
            extinfo_str = ""
            if last_extinfo and cfg.ext_info_mode == rtmslib.ExtendedInfoMode.BasicImageInfo:
                extinfo_str = (
                    f" last_extinfo: ts={last_extinfo.timestamp:.6f}, "
                    f"dt={last_extinfo.datatype}, "
                    f"{last_extinfo.width}x{last_extinfo.height}, "
                    f"offset=({last_extinfo.offset_x},{last_extinfo.offset_y})"
                )
            logging.info(
                "Stats: received=%s failed=%s rate=%.2f/s stddev=%.6fs bytes=%s last_ts=%.6f%s",
                received,
                failed,
                rate,
                stddev,
                total_bytes,
                ext_info.timestamp,
                extinfo_str,
            )
            last_info = now

    # Final stats
    elapsed = time.time() - start_time
    rate = received / elapsed if elapsed > 0 else 0.0
    logging.info(
        "Stats: received=%s failed=%s rate=%.2f/s elapsed=%.3fs bytes=%s",
        received,
        failed,
        rate,
        elapsed,
        total_bytes,
    )

    # Return success if at least one sample was received
    if received > 0:
        return 0
    logging.error("No samples received")
    return 1


if __name__ == "__main__":
    sys.exit(main())
