import pickle
with open('/home/thebears/source/infer/dump','rb') as ff:
    dump = pickle.load(ff)



# %%

import av
import numpy as np

#def rtp_to_matrix(rtp_packets):
if True:
    """Convert list of RTP packets to image matrix"""
    
    # Create codec context (adjust codec based on your stream)
    codec = av.CodecContext.create('h264', 'r')  # or 'h265', 'mjpeg', etc.
    
    frames = []
    for _,rtp_packet in ts_data.items():
        # Extract payload (skip RTP header - usually 12 bytes)
        payload = rtp_packet
        # Create packet
        print(len(payload))
        packet = av.Packet(payload)

        try:
            # Decode
            for frame in codec.decode(packet):
                # Convert to numpy array
                img = frame.to_ndarray(format='bgr24')
                frames.append(img)
        except Exception as e:
            print(f"Decode error: {e}")
            continue
    
    # Flush decoder
    for frame in codec.decode(None):
        img = frame.to_ndarray(format='bgr24')
        frames.append(img)
    


# Usage
# rtp_packets = [packet1, packet2, ...]  # Your raw RTP packets
# matrices = rtp_to_matrix(rtp_packets)


# %%

import subprocess
import numpy as np
import cv2
from io import BytesIO

class RTPDepayloader:
    def __init__(self, payload_type):
        """
        payload_type: 96=H264, 26=JPEG, etc.
        """
        self.payload_type = payload_type
        self.buffer = bytearray()
        
    def parse_rtp_header(self, packet):
        """Parse RTP header and return payload"""
        if len(packet) < 12:
            return None
            
        # RTP Header structure (12 bytes minimum)
        # Byte 0: V(2), P(1), X(1), CC(4)
        # Byte 1: M(1), PT(7)
        # Bytes 2-3: Sequence number
        # Bytes 4-7: Timestamp
        # Bytes 8-11: SSRC
        
        version = (packet[0] >> 6) & 0x03
        padding = (packet[0] >> 5) & 0x01
        extension = (packet[0] >> 4) & 0x01
        csrc_count = packet[0] & 0x0F
        
        marker = (packet[1] >> 7) & 0x01
        payload_type = packet[1] & 0x7F
        
        sequence = int.from_bytes(packet[2:4], 'big')
        timestamp = int.from_bytes(packet[4:8], 'big')
        
        # Calculate header size
        header_size = 12 + (csrc_count * 4)
        
        if extension:
            if len(packet) < header_size + 4:
                return None
            ext_len = int.from_bytes(packet[header_size+2:header_size+4], 'big') * 4
            header_size += 4 + ext_len
        
        # Return payload
        payload = packet[header_size:]
        
        return {
            'payload': payload,
            'marker': marker,
            'sequence': sequence,
            'timestamp': timestamp,
            'payload_type': payload_type
        }
    
    def depayload_h264(self, rtp_packets):
        """Depayload H.264 RTP packets to NAL units"""
        nal_units = bytearray()
        
        for packet in rtp_packets:
            parsed = self.parse_rtp_header(packet)
            if not parsed:
                continue
                
            payload = parsed['payload']
            if len(payload) == 0:
                continue
            
            # H.264 RTP payload format (RFC 6184)
            nal_unit_type = payload[0] & 0x1F
            
            if nal_unit_type >= 1 and nal_unit_type <= 23:
                # Single NAL unit
                nal_units.extend(b'\x00\x00\x00\x01')  # Start code
                nal_units.extend(payload)
            elif nal_unit_type == 28:
                # FU-A (Fragmentation Unit)
                fu_header = payload[1]
                start = (fu_header >> 7) & 0x01
                end = (fu_header >> 6) & 0x01
                nal_type = fu_header & 0x1F
                
                if start:
                    # Reconstruct NAL header
                    nal_header = (payload[0] & 0xE0) | nal_type
                    nal_units.extend(b'\x00\x00\x00\x01')
                    nal_units.append(nal_header)
                
                nal_units.extend(payload[2:])
        
        return bytes(nal_units)

def decode_h264_to_matrix(h264_data):
    """Decode H.264 bitstream to image matrix using FFmpeg"""
    
    # Use FFmpeg to decode
    process = subprocess.Popen([
        'ffmpeg',
        '-f', 'h264',
        '-i', 'pipe:0',
        '-f', 'rawvideo',
        '-pix_fmt', 'bgr24',
        'pipe:1'
    ], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    stdout, stderr = process.communicate(input=h264_data)
    
    # You need to know dimensions or parse from SPS
    # For now, assume 1920x1080
    width, height = 1920, 1080
    
    if len(stdout) >= width * height * 3:
        frame = np.frombuffer(stdout[:width*height*3], dtype=np.uint8)
        frame = frame.reshape((height, width, 3))
        return frame
    
    return None

# Usage
rtp_packets = [bytes(x) for x in dump]
depayloader = RTPDepayloader(payload_type=96)
# rtp_packets = [...]  # Your raw RTP packets
h264_stream = depayloader.depayload_h264(rtp_packets)
matrix = decode_h264_to_matrix(h264_stream)


# %%


bts = b''
for i in range(10000):
    bts+=process.stdout.read(1024)
    print(i)
# %%
with open('dumpb','wb') as ff:
    ff.write(bts)
# %%
import re
pattern = b'BM'

fcf = [m.start() for m in re.finditer(re.escape(pattern), bts)]


# %%

with open('dumpb','rb') as ff:
    byte_stream = ff.read()[10:]


import struct
import numpy as np
import cv2




# %%
import struct


def read_bmp_frames(stream):
    """Read individual BMP frames from a stream."""
    while True:
        # Read BMP signature (2 bytes)
        signature = stream.read(2)
        if len(signature) < 2:
            break
        
        if signature != b'BM':
            print(f"Warning: Expected 'BM', got {signature}")
            break
        
        # Read file size (4 bytes, little-endian)
        size_bytes = stream.read(4)
        if len(size_bytes) < 4:
            break
        
        file_size = struct.unpack('<I', size_bytes)[0]
        
        # Read the rest of the frame
        remaining_size = file_size - 6  # Already read 6 bytes
        frame_data = stream.read(remaining_size)
        
        if len(frame_data) < remaining_size:
            break
        
        # Complete frame
        complete_frame = signature + size_bytes + frame_data
        yield complete_frame

# Usage
import subprocess as sp

cmd = "ffmpeg -rtsp_transport tcp -i rtsp://admin:marybear@192.168.1.153:554/h264Preview_01_sub -f image2pipe -vcodec bmp -an pipe:1"

# process = sp.Popen(
#         cmd.split(" "),
#         stdout=sp.PIPE,
#         stderr=sp.PIPE,
#     bufsize =512*896*3*3
#     )




import asyncio
import subprocess

async def read_stream():
    """Read subprocess output asynchronously."""
    
    process = await asyncio.create_subprocess_exec(*cmd.split(" "),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    frame_count = 0
    try:
        while True:
            # Read up to 65KB at a time
            chunk = await process.stdout.read(65536)
            print('chunk')
            if not chunk:
                break
            
            frame_count += len(chunk)
            print(f"Read {len(chunk)} bytes (total: {frame_count})")
    
    except asyncio.CancelledError:
        process.kill()
        await process.wait()
    
    return frame_count

# Run
frame_count = asyncio.run(read_stream())
# %%





# %%

import numpy as np
fr = next(gen)
import cv2

# %%

import av
import numpy as np

def rtp_to_matrix(rtp_packets):
    """Convert list of RTP packets to image matrix"""
    
    # Create codec context (adjust codec based on your stream)
    codec = av.CodecContext.create('h264', 'r')  # or 'h265', 'mjpeg', etc.
    
    frames = []
    for rtp_packet in rtp_packets:
        # Extract payload (skip RTP header - usually 12 bytes)
        payload = rtp_packet[12:]
        
        # Create packet
        packet = av.Packet(payload)
        
        try:
            # Decode
            for frame in codec.decode(packet):
                # Convert to numpy array
                img = frame.to_ndarray(format='bgr24')
                frames.append(img)
        except Exception as e:
            print(f"Decode error: {e}")
            continue
    
    # Flush decoder
    for frame in codec.decode(None):
        img = frame.to_ndarray(format='bgr24')
        frames.append(img)
    
    return frames

# Usage
# rtp_packets = [packet1, packet2, ...]  # Your raw RTP packets
# matrices = rtp_to_matrix(rtp_packets)