from base64 import b64decode, b64encode
from struct import unpack as s_unpack
from io import BytesIO
from collections import defaultdict
import math, gzip
import msgpack
read_int = lambda stream, nbytes, signed=False: int.from_bytes(stream.read(nbytes), 'little',signed=signed)
read_float = lambda stream: s_unpack('<f', stream.read(4))[0]
class SSEDataLengthOutOfRangeException(Exception):
    needed : int
    current : int
    def __init__(self, needed, current):
        self.needed = needed
        self.current = current
        super().__init__(f'Needed {needed} bytes, but only {current} bytes available.')    
# Sekai_Streaming_StreamingCommon__CheckHeader
def decode_buffer_base64(buffer : bytes) -> tuple[int, bytes]:
    '''Decodes the 'RTVL' SSE Message payload into a header signature and data.

    Args:
        buffer (bytes): Encoded buffer

    Raises:
        SSEDataLengthOutOfRangeException: If the buffer length does not match the expected length.

    Returns:
        (int, bytes): Header Signature and decoded data payload
    '''
    # They really want you to believe it's Base64...
    # [4 bytes: RTVL][6 bytes : length in hex][Base64 T/F][Split T/F][3 bytes: Signature][data, Base64]
    # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 15 bytes
    # Sekai_Streaming_StreamingCommon__cctor
    SYM_HEADER = b'RTVL'
    SYM_TRUE = b'T'
    SYM_FALSE = b'F'
    stream = BytesIO(buffer)
    assert stream.read(4) == SYM_HEADER, 'bad header'
    encoded_length = stream.read(6)
    encoded_length = int(encoded_length, 16)
    is_base64_encoded = stream.read(1) == SYM_TRUE
    is_split = stream.read(1) == SYM_TRUE
    header_signature = int(stream.read(3).decode(), 10)
    data = stream.read()
    current_length = len(data) + 15
    if current_length != encoded_length:
        raise SSEDataLengthOutOfRangeException(encoded_length, current_length)    
    if is_base64_encoded:
        data = b64decode(data)
    return header_signature, data
# Sekai_Streaming_SubscribeDecoder__Deserialize
def decode_buffer_payload(buffer : bytes) -> tuple[int, bytes]:
    '''Decodes the decoded data payload (in binary form) into a decoder signature and payload data.

    Args:
        buffer (bytes): Decoded buffer

    Returns:
        (int, bytes): Decoder signature and decoded payload data. The signature should match the Header Signature.        
    '''
    stream = BytesIO(buffer)
    decoder_signature = read_int(stream, 1)
    unk1 = read_int(stream, 4)
    if not unk1:
        unk2 = read_int(stream, 4)
        if unk2:
            payload = stream.read()
            payload = gzip.decompress(payload)
            return decoder_signature, payload        
    return decoder_signature, stream.read()
# Sekai_Streaming_StreamingData__Deserialize
def decode_streaming_data(version : tuple, decoder_signature, buffer, strict=True) -> dict:
    '''Decodes the streaming data payload into a dictionary.
    
    Args:
        version (tuple): RLA version
        decoder_signature (int): Decoder signature
        buffer (bytes): Decoded payload data
        strict (bool, optional): If False, incomplete packets will be returned as is. Defaults to True.

    Raises:
        Exception: If the packet is incomplete and strict is True.

    Returns:
        dict: Parsed streaming data payload
    '''
    stream = BytesIO(buffer)
    n_mask_offset = read_int(stream, 4)
    n_init_pos = stream.tell()
    stream.seek(n_mask_offset)
    n_mask_length = read_int(stream, 2)
    bitmask = stream.read(n_mask_length)
    # assert not stream.read() # EOF
    stream.seek(n_init_pos)
    # CP_Serialize_SerializableValueSet    
    mask_i = -1
    def get_next_mask():
        nonlocal mask_i
        mask_i += 1
        return bitmask[mask_i // 8] & (1 << (mask_i % 8)) != 0
    get_next_pred = lambda: get_next_mask() | (get_next_mask() << 1)
    get_next_byte = lambda: [lambda: 0, lambda: 1, lambda: -1, lambda: read_int(stream, 1)][get_next_pred()]() # ReadByte
    get_next_ushort = lambda: [lambda: 0.0, lambda: 0.0, lambda: read_int(stream, 1), lambda: read_int(stream, 2)][get_next_pred()]() # ReadUShort
    get_next_short =  lambda: [lambda: 0.0, lambda: 0.0, lambda: read_int(stream, 1, True), lambda: read_int(stream, 2)][get_next_pred()]() # ReadShort
    get_next_int = lambda: [lambda: 0, lambda: read_int(stream,1), lambda: read_int(stream,1,True), lambda: read_int(stream, 4)][get_next_pred()]() # ReadInt
    get_next_long = lambda: [lambda: 0, lambda: read_int(stream,1), lambda: read_int(stream,2), lambda: read_int(stream, 8)][get_next_pred()]() # ReadLong
    get_next_float = lambda: [lambda: 0.0, lambda: 0.0, lambda: 3.4028e38, lambda: read_float(stream)][get_next_pred()]() # ReadSingle
    get_next_vector3 = lambda: (get_next_float(), get_next_float(), get_next_float()) # ReadVector3
    get_next_ushort_vector3 = lambda: (get_next_ushort() * 0.01, get_next_ushort() * 0.01, get_next_ushort() * 0.01) # ReadUShortVector3
    get_next_tiny_int = lambda type: {'+': lambda: read_int(stream, 2), '*': lambda: read_int(stream, 1, True), ')': lambda: read_int(stream, 1)}[type]() # ReadTinyInt
    get_next_string = lambda: stream.read(get_next_tiny_int(chr(read_int(stream, 1)))).decode() if get_next_pred() else None # ReadString
    # HACK: Workaround for incomplete packets.
    def gen_exception_handler(generator):
        try:
            for item in generator:
                yield item
        except Exception as e: 
            # XXX: This shouldn't happen unless the packet is corrupted.
            if strict: raise e
            else:
                return
    gen_get_next_array = lambda reader: (reader() for _ in range(get_next_int()))  # ReadArray<T>
    get_next_array = lambda reader: list(gen_exception_handler(gen_get_next_array(reader)))
    # Sekai_Streaming_StreamingData__Deserialize
    assert get_next_byte() == decoder_signature, 'bad signature'
    compress_type = get_next_int()
    sequence_no = get_next_int()
    target_time = get_next_long()
    # StreamingData$$ReadValue implementations
    deg_to_rad = lambda value: tuple(deg * (math.pi / 180) for deg in value)
    get_next_deg_as_rad = lambda: deg_to_rad(get_next_ushort_vector3()) # Eulers [0, 2pi]
    get_next_pose_data = lambda: {
        'bodyPosition': get_next_vector3(),
        'bodyRotation': get_next_deg_as_rad(), # Eulers
        **({
            'musicItemPropPosition': get_next_vector3(),
            'musicItemPropRotation': get_next_deg_as_rad()
        } if version >= (1,1) else {}),                    
        'boneDatas': get_next_array(get_next_deg_as_rad), # Eulers [0, 2pi]
        'shapeDatas': get_next_array(get_next_float), # [0,100]
        **({
            'propBoneDatas': get_next_array(get_next_deg_as_rad) # Eulers [0, 2pi]
        } if version >= (1,1) else {}),
        'heightOffset': get_next_short() * 0.01,
        'isActive': get_next_mask(),
        'useActiveFx': get_next_mask(),
        **({
            'isEyeLookAt': get_next_mask()
        } if version >= (1,4) else {})
    }
    match decoder_signature:
        case 0:
            # Sekai_Streaming_MotionData
            timeStamps = get_next_array(get_next_long) if get_next_pred() == 2 else None
            poses = get_next_array(get_next_pose_data)
            return {'type': 'MotionData', 'timeStamps': timeStamps, 'poses': poses}
        case 1:
            # Sekai_Streaming_MotionCaptureData
            read_character_capture_data = lambda: {
                'id': get_next_int(),
                'timestamp': get_next_long(),
                'pose': get_next_pose_data()
            }
            data = get_next_array(read_character_capture_data)
            return {'type': 'MotionCaptureData', 'data': data}
        case 2:
            # Sekai_Streaming_SoundData
            channels = get_next_int()
            sample_rate = get_next_int()
            data_length = get_next_int()
            if compress_type != 1:
                data = stream.read(data_length) # Raw sampling data
                return {'type': 'SoundData', 'channels': channels, 'sampleRate': sample_rate, 'encoding': 'raw', 'data': b64encode(data).decode()}
            else:
                data = stream.read(data_length) # HCA w/o metadata
                return {'type': 'SoundData', 'channels': channels, 'sampleRate': sample_rate, 'encoding': 'hca', 'data': b64encode(data).decode()}              
        case 3:
            # Sekai_Streaming_StatusData
            read_stage_status = lambda: {
                'liveState': get_next_byte(),
                'lightIntensity': get_next_float(),
                'gayaVolume': get_next_float(),
                'cheerVolume': get_next_float(),
                'characterSpotlightIndex': get_next_int(),
                'characterSpotlightIntensity': get_next_float(),
                'stageSetlistIndex': get_next_int(),
                'musicSetlistIndex': get_next_int(),
                'musicTargetTime': get_next_long(),
                'musicStartTime': get_next_float(),
                'seId': get_next_int(),
                'seStartTime': get_next_long(),
                'timeStamp': get_next_long(),
                'unk0': get_next_byte(),
                'unk1': get_next_byte(),
                'playTimelineId': get_next_int(),                    
                ** ({
                    'screenFadeColorR': get_next_byte(),
                    'screenFadeColorG': get_next_byte(),
                    'screenFadeColorB': get_next_byte(),
                    'screenFade': get_next_float(),
                    'characterFormationRotate': get_next_int(),
                    'stageCenterPosition': get_next_vector3(),
                    'playerAvatarStartPosition': get_next_vector3()
                } if version >= (1,4) else {})
            }
            read_character_status = lambda: {
                'costumeIndex': get_next_int(),
                ** ({'useFx': get_next_mask()} if version >= (1,2) else {}),
                'timeStamp': get_next_long(),
            }
            stage_status_list = get_next_array(read_stage_status)
            stage_status_length = get_next_int()                
            charcter_status_list = [get_next_array(read_character_status) for _ in range(stage_status_length)]
            return {'type': 'StatusData', 'stageStatus': stage_status_list, 'characterStatus': charcter_status_list}
        case 4:
            # Sekai_Streaming_VirtualLiveMessageData
            message_id = get_next_int()
            user_id = get_next_string()
            data_length = get_next_int()
            data = stream.read(data_length)
            # CP_BinarySerializer__Deserialize_OtherRoomMessageData
            # CP_BinarySerializer__Deserialize_OtherRoomActionData
            # NOTE: The msgpack payload does not contain the message type
            data = msgpack.unpackb(data)
            return {'type': 'VirtualLiveMessageData','messageId': message_id, 'userId': user_id, 'data': data}
        case 5:
            # Sekai_Streaming_ComplementInfoData
            info_type = get_next_int()
            info_data = get_next_string()
            return {'type': 'ComplementInfoData', 'infoType': info_type, 'infoData': info_data}            
    return {'type': 'Unknown', 'data': buffer}
result = defaultdict(dict)

def read_rla(src : BytesIO, version=(1,0), strict=True) -> dict:
    '''Parses the Sekai RLA file format used in 'streaming_live/archive' assets.

    Args:
        src (BytesIO): Source RLA file stream
        version (tuple, optional): RLA version, found in respective RLH (JSON) header files. i.e. one of (1,0), (1,1), (1,2), (1,3), (1,4) Defaults to (1,0). 
        strict (bool, optional): If False, incomplete packets will be returned as is. Defaults to True.

    Returns:
        dict: Parsed RLA data. The dictionary is sorted by the frame ticks.
    '''    
    result = defaultdict(dict)
    def read_frames():
        ticks = read_int(src, 8)
        if ticks:
            buffer_length = read_int(src, 4)
            buffer = src.read(buffer_length)
            header_signature, data = decode_buffer_base64(buffer)
            decoder_signature, data = decode_buffer_payload(data)
            assert header_signature == decoder_signature, 'mismatching signature (header/decoder). packet may be corrupt'
            payload = decode_streaming_data(version, decoder_signature, data, strict)
            result[ticks].setdefault(payload['type'], list()).append(payload)
            return True
        return False
    while read_frames(): pass        
    result = dict(sorted(result.items(), key=lambda x: x[0]))
    return result

if __name__ == '__main__':
    from timeit import timeit
    fp = open(r'C:\Users\mos9527\Desktop\sekai_streaming\sekai_30_00000060.rla.bytes','rb')
    buffer = fp.read()
    fp = BytesIO(buffer)
    print('took %.2f seconds' % timeit("read_rla(fp, (1,0))",globals=globals()))
