r"""Encoding of Asterix CAT240 messages.

Use :func:`encode` to compile new Asterix CAT240 messages and feed it with the output of
either :func:`make_summary` or :func:`make_video_message` for creating a type `001` or
`002` message, respectively. The output is a byte sequence of a single CAT240 message
that can directly be concatenated with others:

Example:
    .. code-block:: python

        payloads = [
            pycatzao.make_summary(summary="FoObaR"),
            pycatzao.make_video_message(
                np.array([255, 0, 255], dtype=np.uint8),  # signal at range = [9m, 12m, 15m]
                msg_index=1,
                header=pycatzao.make_video_header(
                    start_az=45.0,
                    end_az=45.0,
                    cell_offset=9,  # (center of) first cell is at 9m
                    cell_width=3,   # cell width is 3m
                )
            ),
        ]

        tod = ...  # get ToD in seoncds
        blocks = [pycatzao.encode(msg, sac=7, sic=42, tod=tod) for msg in payloads]
        print(b"".join(blocks).hex())  # outputs: f00012d108072a01...
"""  # noqa: E501

import zlib

import numpy as np


def make_summary(summary):
    """Create Video Summary.

    Creates a type `001` message (Video Summary message) to be passed as the payload to
    :func:`encode`.

    Args:
        summary (str):
            (ASCII) string to deliver stream meta data and/or labels.

    Returns:
        typing.Any:
            Payload for :func:`encode`.

    Raises:
        ValueError:
            The ASCII representation of the summary has to be not longer than 0xFFFF
            characters.
    """
    chars = summary.encode("ascii")
    n = len(chars)
    if n > 0xFFFF:  # pragma: no cover
        raise ValueError(
            f"Summary message is too long. Not more than {0xFFFF} ASCII symbols are"
            " allowed. See specification of I240/030 for details."
        )

    return {
        "uap": 0x1000,
        "data": n.to_bytes(1, byteorder="big") + chars,
    }


def make_video_header(*, start_az, end_az, cell_offset, cell_width):
    """Create Video Header.

    Creates a video header (I240/040 or I240/041) for :func:`make_video_message`.

    Args:
        start_az (float):
            Start azimuth in degrees.

        end_az (float):
            End azimuth in degrees.

        cell_offset (float):
            Offset of first cell in meters.

        cell_width (float):
            Cell width in meters. If this value is larger than 500m, the header type
            is I240/040 (Video Header Nano) and I240/041 (Video Header Femto) otherwise.

    Returns:
        typing.Any:
            The header for func:`make_video_message`.
    """
    start_az %= 360
    end_az %= 360

    az_scale = 2**16 / 360
    start_az = round(start_az * az_scale).to_bytes(2, byteorder="big")
    end_az = round(end_az * az_scale).to_bytes(2, byteorder="big")

    item, t_scale = ("I240/040", 1e-9) if cell_width > 500 else ("I240/041", 1e-15)
    cell_dur = round(2 * cell_width / (299_792_458 * t_scale)).to_bytes(
        4, byteorder="big"
    )
    start_rg = round(cell_offset / cell_width).to_bytes(4, byteorder="big")

    return {
        "uap": 0x0800 if item == "I240/040" else 0x0400,
        "data": start_az + end_az + start_rg + cell_dur,
    }


def make_video_message(amp, *, msg_index, header, compress=True):
    """Create Video Message.

    Creates a video message (I240/050, I240/051 or I240/052) for :func:`encode` from a
    video header.

    Args:
        amp (np.ndarray):
            Amplitude array given as an `np.ndarray` with `dtype` set to either
            `np.uint8`, `np.uint16`, or `np.uint32`.

        msg_index (int):
            Message Sequence Identifier (video record cyclic counter).

        header:
            Header generated by :func:`make_video_header`.

        compress (bool):
            Compress the amplitude array before compiling the message. Note that this is
            typically cheap yet very effective and should be enabled unless you know
            better.

    Returns:
        typing.Any:
            Payload for :func:`encode`.
    """
    if not isinstance(amp, np.ndarray):  # pragma: no cover
        raise ValueError("Amplitude array has to be of type `np.ndarray`.")

    if amp.dtype not in [np.uint8, np.uint16, np.uint32]:  # pragma: no cover
        raise ValueError(
            "Data type of amplitude array is not supported. Data type has to be either"
            " 'np.uint8', 'np.uint16', or 'np.uint32'."
        )

    nb_cells = len(amp)
    amp = amp.tobytes()

    res = ((len(amp) * 8) // nb_cells).bit_length().to_bytes(1, byteorder="big")

    if compress:
        amp = zlib.compress(amp)

    nb_vb = len(amp)

    items = {4: "I240/050", 64: "I240/051", 256: "I240/052"}
    keys = list(items.keys())
    i = np.argmin([(1 + (nb_vb - 1) // n) * (n + 1) for n in keys])
    n = keys[i]
    item = items[n]

    amp += b"\x00" * ((1 + (nb_vb - 1) // n) * n - nb_vb)
    assert len(amp) % n == 0

    n = (len(amp) // n).to_bytes(1, byteorder="big")
    nb_vb = nb_vb.to_bytes(2, byteorder="big")
    nb_cells = nb_cells.to_bytes(3, byteorder="big")

    uap = {"I240/050": 0x02C0, "I240/051": 0x02A0, "I240/052": 0x0290}[item]
    uap = uap | header["uap"]

    data = msg_index.to_bytes(4, byteorder="big")
    data += header["data"]
    data += (b"\x80" if compress else b"\x00") + res + nb_vb + nb_cells + n + amp

    return {
        "uap": uap,
        "data": data,
    }


def encode(payload, *, sac, sic, tod=-1):
    """Compiles data into a CAT240 message.

    Takes a payload generated by either :func:`make_summary` or
    :func:`make_video_message` and wraps it as an Asterix CAT240 message.

    Args:
        payload: The payload to encode.
        sac (int):
            System Area Code (SAC)
        sic (int):
            System Identification Code (SIC)
        tod (float):
            Time of Day (ToD) in seconds. If negative, no ToD is included into the
            compiled message.

    Returns:
        bytes:
            A single Asterix CAT240 message.
    """
    sac = sac.to_bytes(1, byteorder="big")
    sic = sic.to_bytes(1, byteorder="big")
    msg_type = b"\x01" if payload["uap"] & 0x1000 else b"\x02"

    if tod <= 0:
        tod = b""
    else:
        tod = round(tod * 128)
        if tod > 0xFFFFFF:  # pragma: no cover
            raise ValueError(
                "Time of Day (given in seconds) is too large. See specification of"
                " I240/140 for details."
            )

        tod = tod.to_bytes(3, byteorder="big")

    uap = 0xC100 if tod == b"" else 0xC108
    uap |= payload["uap"]
    uap = uap.to_bytes(2, byteorder="big")

    data = uap + sac + sic + msg_type + payload["data"] + tod

    return b"\xf0" + (len(data) + 3).to_bytes(2, byteorder="big") + data
