#!/usr/bin/env python3
"""
AudioSpecificConfig 解析器
用于解析 AAC Sequence Header
"""

import struct
from dataclasses import dataclass


@dataclass
class AudioSpecificConfig:
    """AAC AudioSpecificConfig"""
    audio_object_type: int      # 音频对象类型
    sampling_frequency_index: int  # 采样率索引
    sampling_frequency: int     # 实际采样率 (Hz)
    channel_configuration: int  # 声道配置
    frame_length_flag: int = 0  # 帧长度标志
    depends_on_core_coder: int = 0
    extension_flag: int = 0

    # AAC-LC 特有
    sbr_present_flag: int = 0
    ps_present_flag: int = 0


# 采样率映射表
SAMPLING_FREQ_TABLE = {
    0: 96000, 1: 88200, 2: 64000, 3: 48000,
    4: 44100, 5: 32000, 6: 24000, 7: 22050,
    8: 16000, 9: 12000, 10: 11025, 11: 8000,
    12: 7350
}

# 音频对象类型名称
AUDIO_OBJECT_TYPE_NAMES = {
    0: "Null", 1: "AAC Main", 2: "AAC-LC", 3: "AAC-SSR",
    4: "AAC-LTP", 5: "SBR", 6: "AAC Scalable",
    7: "TwinVQ", 8: "CELP", 9: "HVXC",
    12: "TTSI", 13: "Main Synthesis", 14: "Wavetable Synthesis",
    39: "ER AAC LD", 42: "ER AAC ELD"
}

# 声道配置名称
CHANNEL_CONFIG_NAMES = {
    0: "Defined in AOT", 1: "Center (Mono)", 2: "L+R (Stereo)",
    3: "Center+L+R", 4: "Center+L+R+Rear", 5: "5 Surround",
    6: "5.1 Surround", 7: "7.1 Surround"
}


def parse_audio_specific_config(data: bytes) -> AudioSpecificConfig:
    """
    解析 AudioSpecificConfig
    
    参数:
        data: AAC Sequence Header 的数据部分（不含 Audio Header 和 AAC Packet Type）
    
    返回:
        AudioSpecificConfig 对象
    """
    if len(data) < 2:
        raise ValueError("AudioSpecificConfig 至少需要 2 字节")

    # 位读取器
    bits = BitReader(data)

    # Audio Object Type (5 bits)
    audio_object_type = bits.read_bits(5)
    if audio_object_type == 31:
        # 扩展类型：再读 6 位 + 32
        audio_object_type = 32 + bits.read_bits(6)

    # Sampling Frequency Index (4 bits)
    sampling_freq_index = bits.read_bits(4)
    if sampling_freq_index == 15:
        # 自定义采样率：读 24 位
        sampling_freq = bits.read_bits(24)
    else:
        sampling_freq = SAMPLING_FREQ_TABLE.get(sampling_freq_index, 0)

    # Channel Configuration (4 bits)
    channel_config = bits.read_bits(4)

    config = AudioSpecificConfig(
        audio_object_type=audio_object_type,
        sampling_frequency_index=sampling_freq_index,
        sampling_frequency=sampling_freq,
        channel_configuration=channel_config,
    )

    # GASpecificConfig (如果 AOT 是 1-7 等)
    if audio_object_type in (1, 2, 3, 4, 6, 7):
        config.frame_length_flag = bits.read_bits(1)
        config.depends_on_core_coder = bits.read_bits(1)
        config.extension_flag = bits.read_bits(1)

    return config


def parse_aac_sequence_header(body: bytes) -> AudioSpecificConfig:
    """
    解析 AAC Sequence Header（完整的 Audio Tag Body）
    
    参数:
        body: Audio Tag Body（包含 1 byte Audio Header + 1 byte AAC Packet Type + ASC）
    """
    if len(body) < 4:
        raise ValueError("AAC Sequence Header 至少需要 4 字节")

    audio_header = body[0]  # 0xAF
    aac_packet_type = body[1]  # 0 = Sequence Header

    if aac_packet_type != 0:
        raise ValueError(f"不是 AAC Sequence Header，type={aac_packet_type}")

    asc_data = body[2:]
    return parse_audio_specific_config(asc_data)


class BitReader:
    """位读取器"""

    def __init__(self, data: bytes):
        self.data = data
        self.byte_pos = 0
        self.bit_pos = 0

    def read_bits(self, n: int) -> int:
        """读取 n 位"""
        result = 0
        for _ in range(n):
            if self.byte_pos >= len(self.data):
                break
            bit = (self.data[self.byte_pos] >> (7 - self.bit_pos)) & 1
            result = (result << 1) | bit
            self.bit_pos += 1
            if self.bit_pos >= 8:
                self.bit_pos = 0
                self.byte_pos += 1
        return result


def create_audio_specific_config(
    object_type: int = 2,
    sample_rate: int = 44100,
    channels: int = 2
) -> bytes:
    """
    创建 AudioSpecificConfig
    
    参数:
        object_type: 音频对象类型 (2=AAC-LC)
        sample_rate: 采样率
        channels: 声道数
    """
    # 查找采样率索引
    freq_index = None
    for idx, freq in SAMPLING_FREQ_TABLE.items():
        if freq == sample_rate:
            freq_index = idx
            break
    if freq_index is None:
        raise ValueError(f"不支持的采样率: {sample_rate}")

    # 编码 AudioSpecificConfig
    bits = 0
    bits |= (object_type & 0x1F) << (32 - 5)
    bits |= (freq_index & 0x0F) << (32 - 5 - 4)
    bits |= (channels & 0x0F) << (32 - 5 - 4 - 4)

    result = bits.to_bytes(2, 'big')
    return result


# 测试
def test_audio_specific_config():
    """测试 AudioSpecificConfig 解析"""
    # AAC-LC, 44100Hz, Stereo
    # Object Type=2, Freq Index=4, Channel=2
    # Binary: 00010 0100 0010 0000
    #         00010010 00010000 = 0x12 0x10
    data = bytes([0x12, 0x10])

    config = parse_audio_specific_config(data)
    print(f"Audio Object Type: {config.audio_object_type} ({AUDIO_OBJECT_TYPE_NAMES.get(config.audio_object_type, 'Unknown')})")
    print(f"Sampling Frequency: {config.sampling_frequency} Hz (index={config.sampling_frequency_index})")
    print(f"Channel Configuration: {config.channel_configuration} ({CHANNEL_CONFIG_NAMES.get(config.channel_configuration, 'Unknown')})")

    # 验证
    assert config.audio_object_type == 2  # AAC-LC
    assert config.sampling_frequency == 44100
    assert config.channel_configuration == 2  # Stereo

    # 编码验证
    encoded = create_audio_specific_config(2, 44100, 2)
    assert encoded == data

    print("✅ AudioSpecificConfig 测试通过")


if __name__ == '__main__':
    test_audio_specific_config()

def create_adts_header(
    object_type: int,
    sample_rate: int,
    channels: int,
    data_length: int
) -> bytes:
    """
    创建 ADTS Header
    
    参数:
        object_type: AAC Profile (0=Main, 1=LC, 2=SSR)
        sample_rate: 采样率
        channels: 声道数
        data_length: AAC 原始数据长度
    """
    # 采样率索引
    freq_table = {
        96000: 0, 88200: 1, 64000: 2, 48000: 3,
        44100: 4, 32000: 5, 24000: 6, 22050: 7,
        16000: 8, 12000: 9, 11025: 10, 8000: 11,
        7350: 12
    }
    freq_index = freq_table.get(sample_rate, 4)
    profile = object_type - 1  # AAC-LC=2, profile=1
    frame_length = 7 + data_length  # 7 bytes ADTS header + data

    adts = bytearray(7)
    adts[0] = 0xFF                        # syncword high
    adts[1] = 0xF0                        # syncword low + ID=0 + layer=00 + protection=1
    adts[1] |= (profile & 0x03) << 2     # profile
    adts[1] |= (freq_index >> 2) & 0x01  # sampling_freq_index bit 3
    adts[2] = (freq_index & 0x03) << 6   # sampling_freq_index bits 2-1
    adts[2] |= (channels & 0x07) << 2    # channel_config
    adts[3] = 0                           # 其他标志
    adts[4] = (frame_length >> 5) & 0xFF # frame_length 高 8 位
    adts[5] = (frame_length & 0x1F) << 3 # frame_length 低 5 位
    adts[6] = 0xFC                        # buffer_fullness + num_frames

    return bytes(adts)


def add_adts_to_aac_raw(aac_raw: bytes, object_type: int = 2,
                         sample_rate: int = 44100, channels: int = 2) -> bytes:
    """为 AAC Raw 数据添加 ADTS 头"""
    adts = create_adts_header(object_type, sample_rate, channels, len(aac_raw))
    return adts + aac_raw

# AAC 编码（推荐）
ffmpeg -i input.mp4 -c:a aac -b:a 128k -ar 44100 -ac 2 -f flv rtmp://...

# AAC HE v1（低码率优化）
ffmpeg -i input.mp4 -c:a libfdk_aac -profile:a aac_he -b:a 64k -f flv rtmp://...

# AAC HE v2（更低码率）
ffmpeg -i input.mp4 -c:a libfdk_aac -profile:a aac_he_v2 -b:a 32k -f flv rtmp://...

# MP3 编码
ffmpeg -i input.mp4 -c:a libmp3lame -b:a 128k -ar 44100 -f flv rtmp://...

# Opus 编码（Enhanced RTMP，需服务端支持）
ffmpeg -i input.mp4 -c:a libopus -b:a 64k -f flv rtmp://...

class AVMuxer:
    """音视频复用器 — 同步音频和视频"""

    def __init__(self):
        self.video_timestamp = 0  # ms
        self.audio_timestamp = 0  # ms
        self.audio_frame_duration = 23  # ms (44100Hz, 1024 samples/frame)
        self.video_frame_duration = 33  # ms (30fps)

    def get_next_video_timestamp(self) -> int:
        ts = self.video_timestamp
        self.video_timestamp += self.video_frame_duration
        return ts

    def get_next_audio_timestamp(self) -> int:
        ts = self.audio_timestamp
        self.audio_timestamp += self.audio_frame_duration
        return ts

#!/usr/bin/env python3
"""
RTMP Audio Decoder
完整的音频消息解析器
"""

import struct
from dataclasses import dataclass


@dataclass
class AudioFrame:
    """音频帧信息"""
    codec: str
    codec_id: int
    sample_rate: int
    sample_size: int
    channels: int
    is_stereo: bool
    data: bytes
    timestamp: int = 0
    # AAC 特有
    aac_packet_type: int = -1
    aac_sequence_header: object = None


SOUND_FORMAT_NAMES = {
    0: 'PCM', 1: 'ADPCM', 2: 'MP3',
    3: 'PCM_LE', 4: 'Nellymoser 16kHz', 5: 'Nellymoser 8kHz',
    6: 'Nellymoser', 7: 'G.711 A-law', 8: 'G.711 mu-law',
    10: 'AAC', 11: 'Speex', 14: 'MP3 8kHz', 15: 'Device'
}

SAMPLE_RATE_MAP = {
    0: 5512, 1: 11025, 2: 22050, 3: 44100
}


def decode_audio_message(body: bytes, timestamp: int = 0) -> AudioFrame:
    """
    解析 Audio Message Body
    
    参数:
        body: Type 8 消息的 Body
        timestamp: 消息时间戳
    
    返回:
        AudioFrame 对象
    """
    if len(body) < 1:
        raise ValueError("Audio message body too short")

    header = body[0]
    sound_format = (header >> 4) & 0x0F
    sound_rate = (header >> 2) & 0x03
    sound_size = (header >> 1) & 0x01
    sound_type = header & 0x01

    codec_name = SOUND_FORMAT_NAMES.get(sound_format, f'Unknown({sound_format})')
    sample_rate = SAMPLE_RATE_MAP.get(sound_rate, 44100)
    sample_size = 8 if sound_size == 0 else 16
    channels = 1 if sound_type == 0 else 2

    frame = AudioFrame(
        codec=codec_name,
        codec_id=sound_format,
        sample_rate=sample_rate,
        sample_size=sample_size,
        channels=channels,
        is_stereo=(sound_type == 1),
        data=body[1:],
        timestamp=timestamp,
    )

    # AAC 特殊处理
    if sound_format == 10 and len(body) >= 2:
        frame.aac_packet_type = body[1]
        if body[1] == 0:
            # AAC Sequence Header
            try:
                config = parse_audio_specific_config(body[2:])
                frame.aac_sequence_header = config
                frame.sample_rate = config.sampling_frequency
                frame.channels = config.channel_configuration
            except Exception:
                pass
        frame.data = body[2:]

    return frame


# 测试
def test_audio_decoder():
    """测试音频解码"""
    # AAC Sequence Header: 0xAF 0x00 0x12 0x10
    body = bytes([0xAF, 0x00, 0x12, 0x10])
    frame = decode_audio_message(body, timestamp=1000)

    print(f"Codec: {frame.codec}")
    print(f"Sample Rate: {frame.sample_rate} Hz")
    print(f"Channels: {frame.channels}")
    print(f"Is Stereo: {frame.is_stereo}")
    print(f"AAC Packet Type: {'Sequence Header' if frame.aac_packet_type == 0 else 'Raw'}")

    assert frame.codec == 'AAC'
    assert frame.sample_rate == 44100
    assert frame.channels == 2
    print("✅ Audio decoder 测试通过")


if __name__ == '__main__':
    test_audio_decoder()