#!/usr/bin/env python3
"""
AMF0 Encoder/Decoder
完整的 AMF0 编解码实现
"""

import struct
from datetime import datetime, timezone


class AMF0:
    """AMF0 类型标记常量"""
    NUMBER = 0x00
    BOOLEAN = 0x01
    STRING = 0x02
    OBJECT = 0x03
    NULL = 0x05
    UNDEFINED = 0x06
    ECMA_ARRAY = 0x08
    OBJECT_END = 0x09  # Object 结束标记 (0x00 0x00 0x09)
    STRICT_ARRAY = 0x0A
    DATE = 0x0B
    LONG_STRING = 0x0C
    XML_DOCUMENT = 0x0F
    TYPED_OBJECT = 0x10
    AMF3_SWITCH = 0x11


class AMF0Encoder:
    """AMF0 编码器"""

    def encode(self, *values) -> bytes:
        """编码多个 AMF0 值"""
        result = b''
        for value in values:
            result += self._encode_value(value)
        return result

    def _encode_value(self, value) -> bytes:
        """编码单个值"""
        if value is None:
            return bytes([AMF0.NULL])
        elif isinstance(value, bool):
            return self._encode_boolean(value)
        elif isinstance(value, (int, float)):
            return self._encode_number(float(value))
        elif isinstance(value, str):
            return self._encode_string(value)
        elif isinstance(value, dict):
            return self._encode_object(value)
        elif isinstance(value, list):
            return self._encode_strict_array(value)
        elif isinstance(value, datetime):
            return self._encode_date(value)
        else:
            return bytes([AMF0.UNDEFINED])

    def _encode_number(self, value: float) -> bytes:
        return bytes([AMF0.NUMBER]) + struct.pack('>d', value)

    def _encode_boolean(self, value: bool) -> bytes:
        return bytes([AMF0.BOOLEAN, 0x01 if value else 0x00])

    def _encode_string(self, value: str) -> bytes:
        encoded = value.encode('utf-8')
        if len(encoded) <= 0xFFFF:
            return bytes([AMF0.STRING]) + struct.pack('>H', len(encoded)) + encoded
        else:
            return bytes([AMF0.LONG_STRING]) + struct.pack('>I', len(encoded)) + encoded

    def _encode_object(self, value: dict) -> bytes:
        result = bytes([AMF0.OBJECT])
        for key, val in value.items():
            key_bytes = key.encode('utf-8')
            result += struct.pack('>H', len(key_bytes)) + key_bytes
            result += self._encode_value(val)
        result += b'\x00\x00' + bytes([AMF0.OBJECT_END])
        return result

    def _encode_strict_array(self, value: list) -> bytes:
        result = bytes([AMF0.STRICT_ARRAY]) + struct.pack('>I', len(value))
        for item in value:
            result += self._encode_value(item)
        return result

    def _encode_date(self, value: datetime) -> bytes:
        timestamp_ms = int(value.timestamp() * 1000)
        timezone_offset = 0  # UTC
        return (bytes([AMF0.DATE]) +
                struct.pack('>d', float(timestamp_ms)) +
                struct.pack('>h', timezone_offset))


class AMF0Decoder:
    """AMF0 解码器"""

    def __init__(self, data: bytes, offset: int = 0):
        self.data = data
        self.offset = offset

    def decode(self) -> list:
        """解码所有值，返回值列表"""
        results = []
        while self.offset < len(self.data):
            try:
                value = self._decode_value()
                results.append(value)
            except (struct.error, IndexError):
                break
        return results

    def _decode_value(self):
        """解码单个值"""
        if self.offset >= len(self.data):
            return None

        marker = self.data[self.offset]
        self.offset += 1

        if marker == AMF0.NUMBER:
            return self._decode_number()
        elif marker == AMF0.BOOLEAN:
            return self._decode_boolean()
        elif marker == AMF0.STRING:
            return self._decode_string()
        elif marker == AMF0.OBJECT:
            return self._decode_object()
        elif marker == AMF0.NULL:
            return None
        elif marker == AMF0.UNDEFINED:
            return None
        elif marker == AMF0.ECMA_ARRAY:
            return self._decode_ecma_array()
        elif marker == AMF0.STRICT_ARRAY:
            return self._decode_strict_array()
        elif marker == AMF0.DATE:
            return self._decode_date()
        elif marker == AMF0.LONG_STRING:
            return self._decode_long_string()
        elif marker == AMF0.AMF3_SWITCH:
            return self._decode_value()  # 递归解码后续值
        else:
            return f'[Unknown marker: 0x{marker:02x}]'

    def _decode_number(self) -> float:
        value = struct.unpack('>d', self.data[self.offset:self.offset+8])[0]
        self.offset += 8
        return value

    def _decode_boolean(self) -> bool:
        value = self.data[self.offset] != 0
        self.offset += 1
        return value

    def _decode_string(self) -> str:
        length = struct.unpack('>H', self.data[self.offset:self.offset+2])[0]
        self.offset += 2
        value = self.data[self.offset:self.offset+length].decode('utf-8')
        self.offset += length
        return value

    def _decode_long_string(self) -> str:
        length = struct.unpack('>I', self.data[self.offset:self.offset+4])[0]
        self.offset += 4
        value = self.data[self.offset:self.offset+length].decode('utf-8')
        self.offset += length
        return value

    def _decode_object(self) -> dict:
        result = {}
        while self.offset < len(self.data):
            # 检查结束标记
            if (self.data[self.offset] == 0 and
                self.data[self.offset+1] == 0 and
                self.data[self.offset+2] == AMF0.OBJECT_END):
                self.offset += 3
                break

            key = self._decode_string_key()
            value = self._decode_value()
            result[key] = value
        return result

    def _decode_string_key(self) -> str:
        length = struct.unpack('>H', self.data[self.offset:self.offset+2])[0]
        self.offset += 2
        value = self.data[self.offset:self.offset+length].decode('utf-8')
        self.offset += length
        return value

    def _decode_ecma_array(self) -> dict:
        count = struct.unpack('>I', self.data[self.offset:self.offset+4])[0]
        self.offset += 4
        # ECMA Array 使用与 Object 相同的键值对格式
        return self._decode_object()

    def _decode_strict_array(self) -> list:
        count = struct.unpack('>I', self.data[self.offset:self.offset+4])[0]
        self.offset += 4
        result = []
        for _ in range(count):
            result.append(self._decode_value())
        return result

    def _decode_date(self) -> datetime:
        timestamp_ms = struct.unpack('>d', self.data[self.offset:self.offset+8])[0]
        self.offset += 8
        timezone_offset = struct.unpack('>h', self.data[self.offset:self.offset+2])[0]
        self.offset += 2
        return datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc)

def amf3_read_integer(data: bytes, offset: int) -> tuple:
    """读取 AMF3 可变长度整数，返回 (value, new_offset)"""
    result = 0
    b = data[offset]
    offset += 1

    if b < 0x80:
        return b, offset

    result = (b & 0x7F) << 7
    b = data[offset]; offset += 1
    if b < 0x80:
        return result | b, offset

    result = (result | (b & 0x7F)) << 7
    b = data[offset]; offset += 1
    if b < 0x80:
        return result | b, offset

    result = (result | (b & 0x7F)) << 8
    b = data[offset]; offset += 1
    return result | b, offset


def amf3_write_integer(value: int) -> bytes:
    """写入 AMF3 可变长度整数"""
    if value < 0x80:
        return bytes([value])
    elif value < 0x4000:
        return bytes([0x80 | (value >> 7), value & 0x7F])
    elif value < 0x200000:
        return bytes([0x80 | (value >> 14), 0x80 | ((value >> 7) & 0x7F), value & 0x7F])
    else:
        return bytes([0x80 | (value >> 22), 0x80 | ((value >> 15) & 0x7F),
                       0x80 | ((value >> 8) & 0x7F), value & 0xFF])

def create_connect_command(
    app: str,
    tc_url: str,
    flash_ver: str = "FMLE/3.0",
    fpad: bool = False,
    capabilities: int = 239,
    audio_codecs: int = 3575,
    video_codecs: int = 252,
    video_function: int = 1,
    object_encoding: int = 0,
    transaction_id: float = 1.0
) -> list:
    """构造 connect 命令的 AMF 值列表"""
    return [
        "connect",                # Command Name
        transaction_id,           # Transaction ID
        {                         # Command Object
            "app": app,
            "flashVer": flash_ver,
            "fpad": fpad,
            "capabilities": capabilities,
            "audioCodecs": audio_codecs,
            "videoCodecs": video_codecs,
            "videoFunction": video_function,
            "objectEncoding": object_encoding,
            "tcUrl": tc_url,
        }
    ]

def create_connect_result(transaction_id: float) -> list:
    """构造 connect 成功响应"""
    return [
        "_result",                # Command Name
        transaction_id,           # Transaction ID (对应 connect 的 ID)
        {                         # Properties
            "fmsVer": "FMS/3,0,1,123",
            "capabilities": 31,
        },
        {                         # Information
            "level": "status",
            "code": "NetConnection.Connect.Success",
            "description": "Connection succeeded.",
            "objectEncoding": 0,
        }
    ]


def create_connect_error(transaction_id: float, code: str, desc: str) -> list:
    """构造 connect 错误响应"""
    return [
        "_error",
        transaction_id,
        None,
        {
            "level": "error",
            "code": code,
            "description": desc,
        }
    ]

def create_create_stream(transaction_id: float) -> list:
    """构造 createStream 命令"""
    return [
        "createStream",
        transaction_id,
        None,  # 无 Command Object
    ]

def create_create_stream_result(transaction_id: float, stream_id: float = 1.0) -> list:
    """构造 createStream 响应"""
    return [
        "_result",
        transaction_id,
        None,
        stream_id,  # 新建的流 ID
    ]

def create_play(
    stream_name: str,
    stream_id: int = 1,
    start: float = -2,
    duration: float = -1,
    reset: bool = True
) -> list:
    """构造 play 命令"""
    return [
        "play",
        0,  # Transaction ID (通常为 0)
        None,
        stream_name,  # 流名称
        start,        # 开始时间 (-2=直播, -1=从头)
        duration,     # 持续时间 (-1=到结束)
        reset,        # 是否重置
    ]

def create_publish(
    stream_name: str,
    publish_type: str = "live"
) -> list:
    """构造 publish 命令"""
    return [
        "publish",
        0,
        None,
        stream_name,    # 流名称
        publish_type,   # "live", "record", "append"
    ]

def create_rpc_call(method: str, transaction_id: float, *args) -> list:
    """构造 RPC 调用命令"""
    return ["call", transaction_id, None, method] + list(args)

# 示例：调用服务端的 getServerTime
encoder = AMF0Encoder()
rpc_msg = create_rpc_call("getServerTime", 3.0)
encoded = encoder.encode(*rpc_msg)

def create_on_status(
    level: str,
    code: str,
    description: str,
    info: dict = None
) -> list:
    """构造 onStatus 事件"""
    return [
        "onStatus",
        0,
        None,
        {
            "level": level,
            "code": code,
            "description": description,
            **(info or {}),
        }
    ]


# 常见事件代码：
# NetStream.Play.Start    - 开始播放
# NetStream.Play.Stop     - 停止播放
# NetStream.Play.Reset    - 播放重置
# NetStream.Publish.Start - 开始推流
# NetStream.Unpublish.Success - 停止推流
# NetStream.Play.StreamNotFound - 流未找到

def test_amf0_codec():
    """AMF0 编解码测试"""
    encoder = AMF0Encoder()
    decoder = AMF0Decoder

    # 测试 connect 命令
    values = create_connect_command(
        app="live",
        tc_url="rtmp://localhost:1935/live"
    )
    encoded = encoder.encode(*values)
    print(f"Encoded {len(values)} values: {len(encoded)} bytes")
    print(f"Hex: {encoded.hex()}")

    # 解码验证
    d = decoder(encoded)
    decoded = d.decode()
    print(f"Decoded: {decoded}")

    # 验证
    assert decoded[0] == "connect"
    assert decoded[1] == 1.0
    assert decoded[2]["app"] == "live"
    print("✅ AMF0 codec test passed!")


if __name__ == '__main__':
    test_amf0_codec()