from .lib import b4ToInt, b2ToInt

__all__ = ["RTP"]


class Header:
    """
    Reference:
        1. https://datatracker.ietf.org/doc/html/rfc7798
        2. https://blog.csdn.net/u010140427/article/details/127773028

    Structure:
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |0                   1                   2                   3  |  10进制分列
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1|  10进制分列
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |      Byte     |      Byte     |      Byte     |      Byte     |  字节分列，以下为真实header结构
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |V=2|P|X|  CC=4 |M|    PT=7     |           SN=16               |
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |                           timestamp                           |
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |                             SSRC                              |
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |                             CSRC                              |
        |                             ....                              |
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|

    Notes:
1、V：RTP协议的版本号，占2位，当前协议版本号为2
2、P：填充标志，占1位，如果P=1，则在该报文的尾部填充一个或多个额外的八位组，它们不是有效载荷的一部分。
3、X：扩展标志，占1位，如果X=1，则在RTP报头后跟有一个扩展报头
4、CC：CSRC计数器，占4位，指示CSRC 标识符的个数
5、M: 标记，占1位，不同的有效载荷有不同的含义，对于视频，标记一帧的结束；对于音频，标记会话的开始。
6、PT: 有效荷载类型，占7位，用于说明RTP报文中有效载荷的类型，如GSM音频、JPEM图像等，
    在流媒体中大部分是用来区分音频流和视频流的，这样便于客户端进行解析。
7、序列号SN：占16位，用于标识发送者所发送的RTP报文的序列号，每发送一个报文，序列号增1。
    这个字段当下层的承载协议用UDP的时候，网络状况不好的时候可以用来检查丢包。同时出现网络抖动的情况可以用来对数据进行重新排序，
    序列号的初始值是随机的，同时音频包和视频包的sequence是分别记数的。
8、时间戳(timestamp)：占32位，必须使用90 kHz时钟频率。时戳反映了该RTP报文的第一个八位组的采样时刻。
    接收者使用时戳来计算延迟和延迟抖动，并进行同步控制。
9、同步信源(SSRC)标识符：占32位，用于标识同步信源。该标识符是随机选择的，参加同一视频会议的两个同步信源不能有相同的SSRC。
10、特约信源(CSRC)标识符：每个CSRC标识符占32位，可以有0～15个。每个CSRC标识了包含在该RTP报文有效载荷中的所有特约信源。
!!! 注：基本的RTP说明并不定义任何头扩展本身，如果遇到X=1，需要特殊处理。

    Demo:
        假设header数据为：
            80 60 00 01  ==> V_P_X_CC M_PT SN
            00 00 00 00  ==> timestamp
            39 6e d3 46  ==> SSRC
            ...
        对应二进制：
            |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
            |0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1|
            |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|  第一行数据
            |       80      |       60      |       00      |       01      |  16进制
            |1 0 0 0 0 0 0 0|0 1 1 0 0 0 0 0|0 0 0 0 0 0 0 0|0 0 0 0 0 0 0 1|  2进制
            |V=2|P|X|  CC=4 |M|    PT=7     |           SN=16               |  对应位
            |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|  第二行数据
            |       00      |       00      |       00      |       00      |  16进制
            |                           timestamp                           |
            |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|  第三行数据
            |       39      |       6e      |       d3      |       46      |  16进制
            |                             SSRC                              |
            |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
            ...
    """
    __SIZE = 12

    def __init__(self, header: "bytes"):
        self.V = header[0] >> 6
        self.P = (header[0] & 0x20) > 0
        self.X = (header[0] & 0x10) > 0
        self.CC = header[0] & 0x0F
        self.M = (header[1] & 0x80) > 0
        self.PT = header[1] & 0x7F
        self.SN = b2ToInt(header[2:4])
        self.timestamp = b4ToInt(header[4:8])
        self.SSRC = b4ToInt(header[8:12])

        self.size = self.__SIZE + self.CC * 4
        self.CSRC = header[self.__SIZE:self.size]

    def __str__(self):
        return (
            f"<RTP-Header: V:{self.V}, P:{self.P}, X:{self.X}, CC:{self.CC}, M:{self.M}, PT:{self.PT}, "
            f"SN:{self.SN}, Time:{self.timestamp}, SSRC:{self.SSRC}, CSRC:{self.CSRC}, size:{self.size}>"
        )


class Extension:
    """
    Structure：
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |0                   1                   2                   3  |
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1|
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |      defined by profile       |           length              |  长度为本行之后的内容长度，可为0
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
        |                        header extension                       |
        |                             ....                              |
        |-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
    """
    __SIZE = 4

    def __init__(self):
        self.dbp, self.len, self.size, self.body = b"", 0, 0, b""

    def parse(self, ext: "bytes"):
        self.dbp = b2ToInt(ext[0:2])
        self.len = b2ToInt(ext[2:4])
        self.size = self.__SIZE + self.len
        self.body = ext[self.__SIZE:self.size]

    def __str__(self):
        return f"<RTP-Extension: dbp:{self.dbp}, len:{self.len}, body:{self.body}, size:{self.size}>"


class RTP:
    def __init__(self, packet: "bytes"):
        self.header = Header(packet)
        self.extension = Extension()
        if self.header.X:
            self.extension.parse(packet[self.header.size:])
        if self.header.P:
            self.payload = packet[self.header.size + self.extension.size:packet[-1]]
        else:
            self.payload = packet[self.header.size + self.extension.size:]

    def __str__(self):
        return f"<RTP-Decoder:\n\t{self.header}\n\t{self.extension}\n\tpayload length: {len(self.payload)}\n>"