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

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

Examples:
  rtmslibSender --src-ip 0.0.0.0 --dest-ip 239.10.0.1 --dest-port 50000 \\
    --bytes-per-pixel 2 --width 2048 --height 2048 --count 10 --ext-info basic
"""

import argparse
import sys
import time
import logging
from pathlib import Path

import ifw.rtmslib as rtmslib
import numpy as np

DATA_TYPES = {
    "Int8": (1, 64),       # map to BYTE
    "Int16": (2, 256),
    "UInt16": (2, 512),
    "Int32": (4, 1024),
    "UInt32": (4, 2048),
    "Float": (4, 65536),
    "Double": (8, 131072),
}


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="RTMS sender (Python bindings)")
    p.add_argument("-s", "--src-ip", required=True, help="Source IP to bind.")
    p.add_argument("-d", "--dest-ip", required=True, help="Destination IP (unicast or multicast).")
    p.add_argument("-p", "--dest-port", type=int, required=True, help="Destination UDP port.")
    p.add_argument("-m", "--mtu-size", type=int, default=1500, help="MTU size (1500 or 9000).")
    p.add_argument("-t", "--topic-id", type=int, default=10, help="RTMS topic ID.")
    p.add_argument("-c", "--component-id", type=int, default=1, help="RTMS component ID.")
    p.add_argument("--data-type", choices=DATA_TYPES.keys(), required=False,
                   help="Pixel data type (sets datatype field and bytes-per-pixel). Optional if FITS payload provided.")
    p.add_argument("--width", type=int, help="Image width (pixels). If a FITS payload is given, this is inferred.")
    p.add_argument("--height", type=int, help="Image height (pixels). If a FITS payload is given, this is inferred.")
    p.add_argument("-n", "--iterations", type=int, default=1,
                   help="Number of samples to send (0=loop forever).")
    p.add_argument("-q", "--frequency", type=float, default=-1.0,
                   help="Sending frequency in Hz (-1 = as fast as possible).")
    p.add_argument("-x", "--ext-info", choices=["none", "timestamp", "basic"], default="basic",
                   help="Extended info mode to emit.")
    p.add_argument("-i", "--payload", type=Path,
                   help="Optional payload file. If FITS, width/height/datatype are inferred and cubes play plane by plane. Otherwise sends raw bytes; if absent, sends a ramp pattern.")
    p.add_argument("--packet-send-delay-microsec", type=int, default=0,
                   help="Optional delay (microseconds) between UDP packets.")
    p.add_argument("--ignore-conn-refused", action="store_true",
                   help="Ignore ECONNREFUSED errors from UDP send (log but do not abort).")
    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:
    if name == "none":
        return getattr(rtmslib.ExtendedInfoMode, "None")
    if name == "timestamp":
        return rtmslib.ExtendedInfoMode.TimestampOnly
    return rtmslib.ExtendedInfoMode.BasicImageInfo


def make_payload(width: int, height: int, bpp: int, path: Path | None) -> bytes:
    size = width * height * bpp
    if path:
        data = path.read_bytes()
        if len(data) != size:
            raise SystemExit(f"Payload size mismatch: expected {size} bytes, got {len(data)}")
        return data
    # Ramp pattern
    if bpp == 1:
        return bytes((i % 256 for i in range(size)))
    if bpp == 2:
        import array
        arr = array.array("H", (i % 65535 for i in range(width * height)))
        return arr.tobytes()
    return bytes(size)


def load_fits_payloads(path: Path):
    try:
        from astropy.io import fits
    except ImportError as exc:
        raise SystemExit("FITS payload requested but astropy is not installed") from exc

    with fits.open(path) as hdul:
        data = hdul[0].data
    if data is None:
        raise SystemExit(f"No data found in FITS file: {path}")
    if data.ndim == 2:
        frames = [data]
    elif data.ndim == 3:
        frames = list(data)
    else:
        raise SystemExit(f"Unsupported FITS dimensionality {data.ndim}: expected 2D or 3D cube")

    dtype_map = {
        np.int8: ("Int8", 64),
        np.int16: ("Int16", 256),
        np.uint16: ("UInt16", 512),
        np.int32: ("Int32", 1024),
        np.uint32: ("UInt32", 2048),
        np.float32: ("Float", 65536),
        np.float64: ("Double", 131072),
    }
    base_dtype = frames[0].dtype.type
    if base_dtype not in dtype_map:
        raise SystemExit(f"Unsupported FITS dtype {frames[0].dtype}")
    data_type_name, data_type_code = dtype_map[base_dtype]
    bpp = frames[0].dtype.itemsize
    height, width = frames[0].shape[-2], frames[0].shape[-1]
    payloads = [np.ascontiguousarray(f).tobytes() for f in frames]
    return width, height, bpp, data_type_name, data_type_code, payloads


def main() -> None:
    args = parse_args()
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s %(levelname)s: %(message)s",
    )
    payloads = []
    if args.payload and args.payload.suffix.lower() in {".fits", ".fit", ".fts"}:
        width, height, bpp, dtype_name, datatype_code, payloads = load_fits_payloads(args.payload)
        logging.info("Loaded FITS payload(s): %s frames, %sx%s, dtype=%s", len(payloads), width, height, dtype_name)
    else:
        if args.data_type is None or args.width is None or args.height is None:
            raise SystemExit("data-type, width, and height are required unless a FITS payload is provided.")
        bpp, datatype_code = DATA_TYPES[args.data_type]
        payloads = [make_payload(args.width, args.height, bpp, args.payload)]
        width, height = args.width, args.height

    cfg = rtmslib.SenderConfig()
    cfg.ext_info_mode = extinfo_mode(args.ext_info)
    cfg.mtu_size = args.mtu_size
    cfg.topic_id = args.topic_id
    cfg.component_id = args.component_id
    cfg.packet_send_delay = args.packet_send_delay_microsec
    cfg.ignore_connection_refused = args.ignore_conn_refused

    sender = rtmslib.RtmsSender(
        src_ip=args.src_ip,
        dest_ip=args.dest_ip,
        dest_port=args.dest_port,
        bytes_per_pixel=bpp,
        raw_image_size=len(payloads[0]),
        sndr_cfg=cfg,
    )

    count = args.iterations
    interval_sec = 0.0 if args.frequency < 0 else (1.0 / args.frequency if args.frequency > 0 else 0.0)
    sent = 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
    sample_ext = rtmslib.SampleExtInfo()
    ext_mode = cfg.ext_info_mode
    if ext_mode == rtmslib.ExtendedInfoMode.BasicImageInfo:
        sample_ext.timestamp = time.time()
        sample_ext.datatype = datatype_code
        sample_ext.width = width
        sample_ext.height = height
        sample_ext.offset_x = 0
        sample_ext.offset_y = 0
    elif ext_mode == rtmslib.ExtendedInfoMode.TimestampOnly:
        sample_ext.timestamp = time.time()
        # Leave datatype/geometry at defaults so only timestamp is packed.
    else:
        # None: keep all defaults so no ext-info is emitted.
        sample_ext.timestamp = 0.0

    total_frames = len(payloads)
    if count == 1 and total_frames > 1:
        count = total_frames

    logging.info("Starting send loop: %s samples, freq=%.1f Hz, dest=%s:%s",
                 count, args.frequency, args.dest_ip, args.dest_port)

    while count == 0 or sent < count:
        if ext_mode in (
            rtmslib.ExtendedInfoMode.BasicImageInfo,
            rtmslib.ExtendedInfoMode.TimestampOnly,
        ):
            sample_ext.timestamp = time.time()
        payload = payloads[sent % total_frames]
        rc = sender.send_sample(sample_ext, payload)
        if rc != 0:
            failed += 1
            now = time.time()
            if last_error == 0.0 or (now - last_error) >= 10.0:
                logging.error("Send failed (rc=%s)", rc)
                last_error = now
        else:
            logging.info("Sent sample %s (%s bytes) to %s:%s",
                         sent, len(payload), args.dest_ip, args.dest_port)
            total_bytes += len(payload)
            now = time.time()
            if prev_ts is not None:
                interval = now - prev_ts
                n = sent  # intervals count equals sent so far
                delta = interval - mean_interval
                mean_interval += delta / n
                m2_interval += delta * (interval - mean_interval)
            prev_ts = now
        sent += 1
        now = time.time()
        if now - last_info >= 10.0:
            elapsed = now - start_time
            rate = sent / elapsed if elapsed > 0 else 0.0
            intervals = max(sent - 1, 1)
            stddev = (m2_interval / intervals) ** 0.5 if intervals > 1 else 0.0
            logging.info(
                "Stats: sent=%s failed=%s rate=%.2f/s stddev=%.6fs bytes=%s",
                sent,
                failed,
                rate,
                stddev,
                total_bytes,
            )
            last_info = now
        if interval_sec > 0:
            time.sleep(interval_sec)

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


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