def encode_basic_header(fmt: int, cs_id: int) -> bytes:
    """编码基本头"""
    fmt = fmt & 0x03
    if 2 <= cs_id <= 63:
        return bytes([(fmt << 6) | cs_id])
    elif 64 <= cs_id <= 319:
        return bytes([fmt << 6, cs_id - 64])
    elif 320 <= cs_id <= 65599:
        return bytes([(fmt << 6) | 1,
                       (cs_id - 64) & 0xFF,
                       ((cs_id - 64) >> 8) & 0xFF])
    else:
        raise ValueError(f"Invalid cs_id: {cs_id}")


def decode_basic_header(data: bytes, offset: int = 0):
    """解码基本头，返回 (fmt, cs_id, header_length)"""
    byte0 = data[offset]
    fmt = (byte0 >> 6) & 0x03
    cs_id = byte0 & 0x3F

    if cs_id == 0:
        # 2 字节格式
        cs_id = data[offset + 1] + 64
        return fmt, cs_id, 2
    elif cs_id == 1:
        # 3 字节格式
        cs_id = data[offset + 1] + (data[offset + 2] << 8) + 64
        return fmt, cs_id, 3
    else:
        # 1 字节格式
        return fmt, cs_id, 1

DEFAULT_CHUNK_SIZE = 128


def split_message_into_chunks(
    message: bytes,
    msg_type_id: int,
    msg_stream_id: int,
    timestamp: int,
    chunk_stream_id: int,
    chunk_size: int = DEFAULT_CHUNK_SIZE
) -> list:
    """将消息拆分为多个块"""
    chunks = []
    msg_length = len(message)
    offset = 0
    first_chunk = True

    while offset < msg_length:
        body_length = min(chunk_size, msg_length - offset)
        body = message[offset:offset + body_length]

        if first_chunk:
            # 第一个块：使用 fmt=0（完整头）
            header = encode_chunk_header(
                fmt=0,
                cs_id=chunk_stream_id,
                timestamp=timestamp,
                msg_length=msg_length,
                msg_type_id=msg_type_id,
                msg_stream_id=msg_stream_id
            )
            first_chunk = False
        else:
            # 后续块：使用 fmt=3（无头）
            header = encode_chunk_header(
                fmt=3,
                cs_id=chunk_stream_id
            )

        chunks.append(header + body)
        offset += body_length

    return chunks


def encode_chunk_header(fmt, cs_id, timestamp=0, msg_length=0,
                        msg_type_id=0, msg_stream_id=0) -> bytes:
    """编码块头"""
    parts = []

    # 基本头
    parts.append(encode_basic_header(fmt, cs_id))

    # 消息头（根据 fmt）
    if fmt == 0:
        ts = min(timestamp, 0xFFFFFF)
        parts.append(struct.pack('>I', ts)[1:])  # 3 bytes
        parts.append(struct.pack('>I', msg_length)[1:])  # 3 bytes
        parts.append(struct.pack('B', msg_type_id))
        parts.append(struct.pack('<I', msg_stream_id))  # little-endian!
        if timestamp >= 0xFFFFFF:
            parts.append(struct.pack('>I', timestamp))  # 扩展时间戳
    elif fmt == 1:
        ts = min(timestamp, 0xFFFFFF)
        parts.append(struct.pack('>I', ts)[1:])
        parts.append(struct.pack('>I', msg_length)[1:])
        parts.append(struct.pack('B', msg_type_id))
        if timestamp >= 0xFFFFFF:
            parts.append(struct.pack('>I', timestamp))
    elif fmt == 2:
        ts = min(timestamp, 0xFFFFFF)
        parts.append(struct.pack('>I', ts)[1:])
        if timestamp >= 0xFFFFFF:
            parts.append(struct.pack('>I', timestamp))
    # fmt == 3: 无消息头

    return b''.join(parts)

def create_set_chunk_size(chunk_size: int) -> bytes:
    """创建 Set Chunk Size 协议控制消息"""
    # 消息类型 ID = 1
    # 块流 ID = 2（协议控制消息）
    msg = struct.pack('>I', chunk_size & 0x7FFFFFFF)

    # 块头 (fmt=0, cs_id=2, timestamp=0)
    header = encode_basic_header(0, 2)
    header += struct.pack('>I', 0)[1:]  # timestamp = 0
    header += struct.pack('>I', 4)[1:]  # msg length = 4
    header += struct.pack('B', 1)       # msg type = 1 (Set Chunk Size)
    header += struct.pack('<I', 0)      # msg stream = 0

    return header + msg


def parse_set_chunk_size(data: bytes) -> int:
    """解析 Set Chunk Size"""
    return struct.unpack('>I', data[:4])[0] & 0x7FFFFFFF

def create_abort(cs_id: int) -> bytes:
    """创建 Abort 消息"""
    msg = struct.pack('>I', cs_id)
    header = encode_basic_header(0, 2)
    header += struct.pack('>I', 0)[1:]  # timestamp
    header += struct.pack('>I', 4)[1:]  # msg length
    header += struct.pack('B', 2)       # msg type = 2 (Abort)
    header += struct.pack('<I', 0)      # msg stream
    return header + msg

#!/usr/bin/env python3
"""
RTMP Chunk Stream Reassembler
重组块流中的消息
"""

import struct
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class ChunkState:
    """单个块流的状态"""
    cs_id: int
    # fmt=0 的完整信息
    timestamp: int = 0
    timestamp_delta: int = 0
    msg_length: int = 0
    msg_type_id: int = 0
    msg_stream_id: int = 0
    # 当前消息缓冲
    msg_buffer: bytes = b''
    has_extended_ts: bool = False
    prev_fmt: int = 3


class ChunkStreamReassembler:
    """RTMP 块流重组器"""

    DEFAULT_CHUNK_SIZE = 128

    def __init__(self):
        self.chunk_size = self.DEFAULT_CHUNK_SIZE
        self.streams: dict[int, ChunkState] = {}

    def set_chunk_size(self, size: int):
        """设置接收块大小"""
        self.chunk_size = size

    def get_state(self, cs_id: int) -> ChunkState:
        """获取或创建块流状态"""
        if cs_id not in self.streams:
            self.streams[cs_id] = ChunkState(cs_id=cs_id)
        return self.streams[cs_id]

    def parse_chunk(self, data: bytes, offset: int = 0) -> tuple:
        """
        解析一个块，返回 (message, new_offset)
        如果消息不完整返回 (None, offset)
        """
        # 解析基本头
        fmt, cs_id, basic_hdr_len = decode_basic_header(data, offset)
        pos = offset + basic_hdr_len
        state = self.get_state(cs_id)

        # 解析消息头
        if fmt == 0:
            ts = struct.unpack('>I', b'\x00' + data[pos:pos+3])[0]
            pos += 3
            msg_len = struct.unpack('>I', b'\x00' + data[pos:pos+3])[0]
            pos += 3
            msg_type = data[pos]; pos += 1
            msg_stream = struct.unpack('<I', data[pos:pos+4])[0]
            pos += 4

            if ts == 0xFFFFFF:
                ts = struct.unpack('>I', data[pos:pos+4])[0]
                pos += 4
                state.has_extended_ts = True

            state.timestamp = ts
            state.msg_length = msg_len
            state.msg_type_id = msg_type
            state.msg_stream_id = msg_stream
            state.msg_buffer = b''
            state.prev_fmt = 0

        elif fmt == 1:
            ts_delta = struct.unpack('>I', b'\x00' + data[pos:pos+3])[0]
            pos += 3
            msg_len = struct.unpack('>I', b'\x00' + data[pos:pos+3])[0]
            pos += 3
            msg_type = data[pos]; pos += 1

            if ts_delta == 0xFFFFFF:
                ts_delta = struct.unpack('>I', data[pos:pos+4])[0]
                pos += 4

            state.timestamp += ts_delta
            state.timestamp_delta = ts_delta
            state.msg_length = msg_len
            state.msg_type_id = msg_type
            state.msg_buffer = b''
            state.prev_fmt = 1

        elif fmt == 2:
            ts_delta = struct.unpack('>I', b'\x00' + data[pos:pos+3])[0]
            pos += 3

            if ts_delta == 0xFFFFFF:
                ts_delta = struct.unpack('>I', data[pos:pos+4])[0]
                pos += 4

            state.timestamp += ts_delta
            state.timestamp_delta = ts_delta
            state.msg_buffer = b''
            state.prev_fmt = 2

        elif fmt == 3:
            if state.has_extended_ts:
                pos += 4
            if state.prev_fmt in (0, 1):
                pass  # 已在之前更新
            elif state.prev_fmt == 2:
                state.timestamp += state.timestamp_delta
            # else: fmt=3 after fmt=3, no change

        # 读取块体
        remaining = min(self.chunk_size, state.msg_length - len(state.msg_buffer))
        if pos + remaining > len(data):
            return None, offset  # 数据不足

        state.msg_buffer += data[pos:pos + remaining]
        pos += remaining

        # 检查消息是否完整
        if len(state.msg_buffer) >= state.msg_length:
            message = {
                'timestamp': state.timestamp,
                'msg_type_id': state.msg_type_id,
                'msg_stream_id': state.msg_stream_id,
                'data': state.msg_buffer[:state.msg_length]
            }
            state.msg_buffer = b''
            return message, pos

        return None, pos  # 消息不完整