import io
import os
import inspect
import logging
import random
import time
import base64
import struct
from functools import partial
from typing import Optional, Union

import av
import s3path
import bz2
import lz4.frame
import lz4.block
import zstandard
import numpy as np
import DracoPy as draco
from PIL import Image as ImagePIL
from PIL import ImageDraw
from mcap.data_stream import ReadDataStream
from mcap.records import Chunk
from pythonjsonlogger import jsonlogger


class ContextLoggerAdapter(logging.LoggerAdapter):
    def process(self, msg, kwargs):
        kwargs["extra"] = self.extra
        return msg, kwargs

    def log(self, level, msg, *args, **kwargs):
        """
        Override the default 'log' method to inject custom contextual information into the LogRecord.
        """
        if self.isEnabledFor(level):
            frame = inspect.currentframe()
            # the frame is two levels up
            stack_info = inspect.getouterframes(frame)[3][0]
            if stack_info:
                # Extracting filename, line number, and function name from the frame
                filename = stack_info.f_code.co_filename
                lineno = stack_info.f_lineno
                func_name = stack_info.f_code.co_name
                record = self.logger.makeRecord(
                    self.logger.name,
                    level,
                    filename,
                    lineno,
                    msg,
                    args,
                    None,
                    func_name,
                    extra=self.extra,
                )
                self.logger.handle(record)


def get_logger(
    name,
    level: Optional[str] = None,
    log_to_file: bool = False,
    json_logging: bool = False,
    correlation_id: str | None = None,
):
    log_level = level or os.environ.get("LOG_LEVEL", "INFO")
    if isinstance(log_level, str):
        log_level = log_level.upper()
    logger = logging.getLogger(name)
    logger.setLevel(log_level)
    # logger.propagate = False
    if logger.hasHandlers():
        return ContextLoggerAdapter(logger, {"correlation_id": correlation_id})
    correlation_id_str = "%(correlation_id)s:" if correlation_id else ""
    if json_logging:
        formatter = jsonlogger.JsonFormatter(
            f"%(asctime)s %(levelname)s {correlation_id_str} %(name)s %(filename)s %(funcName)s %(lineno)s %(message)s"
        )
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)
        logger.addHandler(handler)
    else:
        formatter = logging.Formatter(
            f"%(asctime)s  (%(levelname)s - %(name)s): {correlation_id_str} %(message)s"
        )
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)
        logger.addHandler(handler)

        if log_to_file:
            handler = logging.FileHandler("lqs.log")
            handler.setFormatter(formatter)
            logger.addHandler(handler)
    return ContextLoggerAdapter(logger, {"correlation_id": correlation_id})


def attempt_with_retries(
    func,
    fargs=None,
    fkwargs=None,
    exceptions=Exception,
    tries=3,
    delay=1.0,
    max_delay=None,
    backoff=2,
    jitter=0,
    message=None,
    logging_func=None,
):
    args = fargs if fargs else list()
    kwargs = fkwargs if fkwargs else dict()
    f = partial(func, *args, **kwargs)
    _tries, _delay = tries, delay
    attempt = 1
    while _tries:
        try:
            return f()
        except exceptions as e:
            _tries -= 1
            if not _tries:
                raise

            if logging_func is not None:
                if message is None:
                    warning_message = f"Failed on attempt {attempt}: {e}; retrying in {_delay} seconds..."
                else:
                    warning_message = f'Failed "{message}" on attempt {attempt}: {e}; retrying in {_delay} seconds...'
                logging_func(warning_message)

            attempt += 1
            time.sleep(_delay)
            _delay *= backoff
            if isinstance(jitter, tuple):
                _delay += random.uniform(*jitter)
            else:
                _delay += jitter
            if max_delay is not None:
                _delay = min(_delay, max_delay)


def get_relative_object_path(object_key: str, source: str):
    if source is None:
        return object_key
    if not object_key.endswith("/"):
        object_key = object_key.rsplit("/", 1)[0] + "/"
    return str(s3path.PureS3Path(object_key + source))


def decompress_chunk_bytes(chunk_bytes, chunk_compression, chunk_length=None):
    if chunk_compression == "bz2":
        decompressed_bytes = bz2.decompress(chunk_bytes)
    elif chunk_compression == "lz4":
        decompressed_bytes = lz4.frame.decompress(chunk_bytes)
    elif chunk_compression == "zstd_mcap":
        chunk = ModifiedChunk.read(ReadDataStream(io.BytesIO(chunk_bytes)))
        return zstandard.decompress(chunk.data, chunk.uncompressed_size)
    elif chunk_compression == "zstd":
        decompressed_bytes = zstandard.decompress(chunk_bytes)
    elif chunk_compression == "lz4_block":
        if chunk_length is None:
            raise ValueError("Must specify chunk_length when decompressing lz4_block")
        decompressed_bytes = lz4.block.decompress(chunk_bytes, chunk_length)
    elif chunk_compression in [None, "none", ""]:
        decompressed_bytes = chunk_bytes
    elif chunk_compression == "mp4":
        decompressed_bytes = chunk_bytes
    else:
        raise ValueError(f"Unknown chunk compression type: {chunk_compression}")
    return decompressed_bytes


# Image

image_type_names = [
    "sensor_msgs/Image",
    "sensor_msgs/CompressedImage",
    "stereo_msgs/DisparityImage",
    "sensor_msgs/msg/Image",
    "sensor_msgs/msg/CompressedImage",
    "stereo_msgs/msg/DisparityImage",
    "ark::image::Image",
    "logqs.image",
    "logqs.inference.depth-estimation",
    "logqs.inference.image-to-image",
    "lqs_cutter/navcon_image",
]


def image_to_base64(
    image: ImagePIL.Image, format: str = "WEBP", format_params: dict = {}
) -> str:
    buffered_img = io.BytesIO()
    image.save(buffered_img, format=format, **format_params)
    return base64.b64encode(buffered_img.getvalue()).decode("utf-8")


def get_complete_segmentation_image(output, meta):
    image_size = (
        meta["image"]["width"],
        meta["image"]["height"],
    )
    img = ImagePIL.new("RGB", image_size, (0, 0, 0))
    for segmentation in output:
        mask_data = base64.b64decode(segmentation["mask"])
        mask = ImagePIL.open(io.BytesIO(mask_data)).convert("L")

        label = segmentation.get("label")
        if label is None:
            color = (
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
            )
        else:
            color = int(hash(label) % 0xFFFFFF)
            color = (
                color & 0xFF,
                (color >> 8) & 0xFF,
                (color >> 16) & 0xFF,
            )
        draw = ImageDraw.Draw(img)
        draw.bitmap((0, 0), mask, fill=color)
    return img


def convert_avc1_to_annexb(data, filter_parameter_sets=True):
    # Wrap data in a memoryview and cache its length.
    mv = memoryview(data)
    mv_len = len(mv)
    offset = 0
    annexb = bytearray()
    start_code = b"\x00\x00\x00\x01"

    while offset + 4 <= mv_len:
        nal_length = struct.unpack_from(">I", mv, offset)[0]
        offset += 4
        # Validate NAL length.
        if nal_length <= 0 or offset + nal_length > mv_len:
            break
        # Avoid parameter set duplication.
        nal_unit = mv[offset : offset + nal_length].tobytes()
        offset += nal_length
        nal_type = nal_unit[0] & 0x1F
        if filter_parameter_sets and nal_type in (7, 8):
            continue
        annexb.extend(start_code)
        annexb.extend(nal_unit)
    return bytes(annexb)


def convert_hvc1_to_annexb(data):
    mv = memoryview(data)
    total_len = len(mv)
    offset = 0
    result = bytearray()
    start_code = b"\x00\x00\x00\x01"

    while offset + 4 <= total_len:
        nal_length = struct.unpack_from(">I", mv, offset)[0]
        offset += 4
        if nal_length <= 0 or offset + nal_length > total_len:
            break
        result.extend(start_code)
        result.extend(mv[offset : offset + nal_length])
        offset += nal_length
    return bytes(result)


def generate_image_from_frames(extradata, frames, frame_data, encoding):
    # 1. Initialize Codec
    codec = av.codec.CodecContext.create(encoding, "r")
    codec.extradata = extradata

    # Precompute start_offset and wrap frame_data in a memoryview.
    start_offset = frames[0][2]
    fd_mv = memoryview(frame_data)
    fd_len = len(fd_mv)

    # Configure converter
    if encoding == "hevc":

        def converter(d):
            return convert_hvc1_to_annexb(d)

    else:
        # Explicitly disable filtering to keep in-band headers if they exist
        def converter(d):
            return convert_avc1_to_annexb(d, filter_parameter_sets=False)

    # 2. Prepare Packet Stream
    # We must prepare valid packets first. We cannot send 0-byte packets.
    packets_to_send = []

    for i, frame in enumerate(frames):
        # Compute slice boundaries
        offset_in_frame = frame[2] - start_offset
        size = frame[1]

        # Check bounds
        if fd_len < offset_in_frame + size:
            break

        # Slice out the NAL unit data
        nal_data = fd_mv[offset_in_frame : offset_in_frame + size]
        annexb_data = converter(nal_data)

        # Handle Payload Construction
        if i == 0:
            # CRITICAL FIX 1: HEADER INJECTION
            # Prepend global extradata (SPS/PPS) to the very first packet.
            # This forces the fresh decoder to accept this frame as a valid start point,
            # even if it is technically a dependent slice in the file.
            packet_payload = extradata + annexb_data
        else:
            packet_payload = annexb_data

        # CRITICAL FIX 2: FILTER EMPTY PACKETS
        # Sending a 0-byte packet triggers EOF in FFmpeg. We must skip them.
        if len(packet_payload) > 0:
            packets_to_send.append(av.Packet(packet_payload))

    # 3. Decode Loop
    # We accumulate frames. We do NOT flush inside this loop.
    decoded_images = []

    for packet in packets_to_send:
        frames = codec.decode(packet)
        decoded_images.extend(frames)

    # 4. Final Flush
    # CRITICAL FIX 3: DEFERRED FLUSH
    # Only now, after sending ALL packets, do we flush buffers.
    frames = codec.decode(None)
    decoded_images.extend(frames)

    # 5. Get the Result
    assert (
        len(decoded_images) > 0
    ), "Decoding failed: No frames produced even after flushing."

    # The image we want is the LAST one (final state after all deltas)
    final_image = decoded_images[-1].to_image()
    return final_image


def get_record_image_lqs(record_data, **kwargs) -> ImagePIL.Image:
    if "frames" in kwargs:
        # if we've passed in frame information, then this is a video frame
        frames = kwargs["frames"]
        encoding = kwargs["encoding"]

        extradata = kwargs.get("extradata", None)
        if isinstance(extradata, str):
            extradata = base64.b64decode(extradata)

        frame_data = record_data
        if isinstance(frame_data, str):
            frame_data = base64.b64decode(frame_data)

        img = generate_image_from_frames(extradata, frames, frame_data, encoding)
    elif isinstance(record_data, str):
        # if the record data is a string, it is a base64 encoded image
        try:
            image_data = base64.b64decode(record_data)
        except Exception as e:
            raise Exception(f"Failed to decode record data: {e}")
        try:
            img = ImagePIL.open(io.BytesIO(image_data))
        except Exception as e:
            raise Exception(f"Failed to open image: {e}")
    else:
        meta = record_data.get("meta")
        if meta:
            pipeline = meta.get("pipeline")
            if pipeline:
                task = pipeline.get("task")
                if task is None:
                    raise Exception(
                        "Expected to find 'task' in record data's 'meta.pipeline' field."
                    )
                if task == "depth-estimation":
                    image_data = base64.b64decode(record_data["output"]["depth"])
                    img = ImagePIL.open(io.BytesIO(image_data))
                elif task == "image-segmentation":
                    img = get_complete_segmentation_image(
                        output=record_data["output"], meta=meta
                    )
                elif task == "image-to-image":
                    image_data = base64.b64decode(record_data["output"])
                    img = ImagePIL.open(io.BytesIO(image_data))
                else:
                    raise Exception(
                        f"Task '{task}' not supported for image generation."
                    )
            else:
                raise Exception(
                    "Expected to find 'pipeline' in record data's 'meta' field."
                )
        else:
            raise Exception("Expected to find 'meta' in record data.")
    return img


def normalize_disparity_image(
    image: ImagePIL.Image, min_clip=74, max_clip=2592, **kwargs
):
    pixels = np.asarray(image)
    pixels = np.clip(pixels, min_clip, max_clip)
    pixels = ((pixels - min_clip) / (max_clip - min_clip)) * 255.0
    return ImagePIL.fromarray(pixels).convert("L")


def image_to_rle_bmp(image) -> io.BytesIO:
    if image.mode != "P":
        image = image.convert("P")

    width, height = image.size
    pixels = np.array(image)

    # 1. Compress Data (RLE8)
    compressed_data = bytearray()

    # BMP RLE stores data bottom-up
    for row in pixels[::-1]:
        i = 0
        while i < width:
            count = 1
            # Count repeating pixels (max 255 per chunk)
            while i + count < width and row[i + count] == row[i] and count < 255:
                count += 1

            # RLE Pair: [Count, Color Index]
            compressed_data.append(count)
            compressed_data.append(row[i])
            i += count

        # End of Line marker (00 00)
        compressed_data += b"\x00\x00"

    # End of Bitmap marker (00 01)
    compressed_data += b"\x00\x01"

    # 2. Build Headers
    # Generate Palette Data from the image
    palette_data = bytearray()
    raw_palette = image.getpalette()

    # Ensure palette is full 768 bytes (256 colors * 3 RGB)
    if len(raw_palette) < 768:
        raw_palette += [0] * (768 - len(raw_palette))

    for i in range(0, 768, 3):
        # BMP palette is BGR, PIL is RGB
        palette_data.append(raw_palette[i + 2])  # B
        palette_data.append(raw_palette[i + 1])  # G
        palette_data.append(raw_palette[i])  # R
        palette_data.append(0)  # Reserved (Alpha)

    # File Header (14 bytes)
    # 14 (FileHeader) + 40 (InfoHeader) + 1024 (Palette)
    offset = 14 + 40 + 1024
    file_size = offset + len(compressed_data)
    file_header = struct.pack("<2sIHHI", b"BM", file_size, 0, 0, offset)

    # Info Header (40 bytes)
    # Compression 1 = BI_RLE8, BitCount = 8
    info_header = struct.pack(
        "<IIIHHIIIIII",
        40,
        width,
        height,
        1,
        8,
        1,
        len(compressed_data),
        2835,
        2835,
        256,
        0,
    )

    # 3. Write to Bytes
    image_bytes = bytearray()
    image_bytes += file_header
    image_bytes += info_header
    image_bytes += palette_data
    image_bytes += compressed_data

    return io.BytesIO(image_bytes)


def get_record_image_ros(record_data, renormalize=False, **kwargs) -> ImagePIL.Image:
    img_modes = {
        "16UC1": "I;16",
        "mono8": "L",
        "mono16": "I;16",
        "32FC1": "F",
        "8UC1": "L",
        "8UC3": "RGB",
        "rgb8": "RGB",
        "bgr8": "RGB",
        "rgba8": "RGBA",
        "bgra8": "RGBA",
        "bayer_rggb": "L",
        "bayer_rggb8": "L",
        "bayer_gbrg": "L",
        "bayer_gbrg8": "L",
        "bayer_grbg": "L",
        "bayer_grbg8": "L",
        "bayer_bggr": "L",
        "bayer_bggr8": "L",
        "yuv422": "YCbCr",
        "yuv411": "YCbCr",
    }

    record_image_format = record_data.get("format", None)
    if record_image_format is None:
        encoding = record_data.get("encoding", None)
        if not encoding:
            return None

        mode = img_modes[encoding]
        img = ImagePIL.frombuffer(
            mode,
            (record_data["width"], record_data["height"]),
            bytes(record_data["data"]),
            "raw",
            mode,
            0,
            1,
        )
        if encoding == "bgr8":
            b, g, r = img.split()
            img = ImagePIL.merge("RGB", (r, g, b))
        elif encoding == "bgra8":
            b, g, r, a = img.split()
            img = ImagePIL.merge("RGBA", (r, g, b, a))
        elif encoding in ["mono16", "16UC1", "32FC1"] and renormalize:
            if kwargs.get("normalize_disparity"):
                img = normalize_disparity_image(img, **kwargs)
            pixels = np.asarray(img)
            pixel_max = np.max(pixels)
            pixel_range = pixel_max - np.min(pixels)
            if pixel_range == 0:
                pixels = np.zeros_like(pixels)
            else:
                pixels = ((pixels - np.min(pixels)) / pixel_range) * 255.0
            img = ImagePIL.fromarray(pixels)
            img = img.convert("L")
    else:
        if record_data["format"] == "h264":
            import av

            codec = av.CodecContext.create("h264", "r")
            packets = av.packet.Packet(bytes(record_data["data"]))
            img = codec.decode(packets)[0].to_image()
        else:
            img = ImagePIL.open(io.BytesIO(bytes(record_data["data"])))
    return img


def get_record_image_ark(record_data, renormalize=False, **kwargs) -> ImagePIL.Image:
    data = bytes(record_data["data"])
    image_format = record_data["data_format"]
    image_format_mapping = {
        42: "BayerRg10",
        44: "BayerRg12",
        43: "BayerRg16",
        41: "BayerRg8",
        31: "Bgr",
        33: "Bgra",
        51: "DepthZ16",
        21: "Grey",
        81: "H264",
        2: "Jpeg",
        1: "MotionJpeg",
        71: "Nv12",
        32: "Rgb",
        34: "Rgba",
        0: "Unset",
        27: "Uv8",
        26: "Yuv420",
        25: "Yuyv",
    }
    if image_format in image_format_mapping:
        image_format = image_format_mapping[image_format]
    else:
        raise NotImplementedError(f"Message type {image_format} not supported.")
    image_format = image_format.lower()

    if image_format == "depthz16":
        ark_img = np.frombuffer(data, dtype=np.float16)

        if renormalize:
            pixel_range = np.max(ark_img) - np.min(ark_img)
            if pixel_range != 0:
                ark_img = ((ark_img - np.min(ark_img)) / pixel_range) * 65535.0

        ark_img = ark_img.astype(np.uint16)
        data = ark_img
        image_mode = "I;16"
        img = ImagePIL.frombuffer(
            image_mode,
            (record_data["width"], record_data["height"]),
            data,
            "raw",
            image_mode,
            0,
            1,
        )
        pixels = np.asarray(img)
        pixel_range = np.max(pixels) - np.min(pixels)
        if pixel_range == 0:
            pixels = np.zeros_like(pixels)
        else:
            pixels = ((pixels - np.min(pixels)) / pixel_range) * 255.0
        img = ImagePIL.fromarray(pixels)
        img = img.convert("L")
    elif image_format == "grey":
        image_mode = "L"
        img = ImagePIL.frombuffer(
            image_mode,
            (record_data["width"], record_data["height"]),
            data,
            "raw",
            image_mode,
            0,
            1,
        )
    elif image_format == "rgb" or image_format == "bgr":
        image_mode = "RGB"
        img = ImagePIL.frombuffer(
            image_mode,
            (record_data["width"], record_data["height"]),
            data,
            "raw",
            image_mode,
            0,
            1,
        )
        if image_format == "bgr":
            b, g, r = img.split()
            img = ImagePIL.merge("RGB", (r, g, b))
    elif image_format == "rgba" or image_format == "bgra":
        image_mode = "RGBA"
        img = ImagePIL.frombuffer(
            image_mode,
            (record_data["width"], record_data["height"]),
            data,
            "raw",
            image_mode,
            0,
            1,
        )
        if image_format == "bgra":
            b, g, r, a = img.split()
            img = ImagePIL.merge("RGBA", (r, g, b, a))
    elif image_format == "jpeg":
        img = ImagePIL.open(io.BytesIO(bytes(record_data["data"])))
    else:
        if image_format == "h264":
            import av

            codec = av.CodecContext.create("h264", "r")
            packets = av.packet.Packet(record_data.data)
            img = codec.decode(packets)[0].to_image()
            # this may not work
        elif image_format.lower() == "nv12":
            width = record_data["width"]
            height = record_data["height"]
            data = record_data["data"]
            Y_size = width * height
            UV_size = (
                Y_size // 2
            )  # There are half as many UV samples as Y samples in NV12

            # Separate Y, U, and V values
            Y = np.array(data[:Y_size], dtype=np.uint8).reshape((height, width))
            UV = np.array(data[Y_size : Y_size + UV_size], dtype=np.uint8).reshape(
                (height // 2, width)
            )

            # Deinterleave U and V components
            U = UV[:, ::2]
            V = UV[:, 1::2]

            # Up-sample U and V components to match Y's dimension
            U_upsampled = np.repeat(np.repeat(U, 2, axis=0), 2, axis=1)
            V_upsampled = np.repeat(np.repeat(V, 2, axis=0), 2, axis=1)

            # Stack the YUV components together to form the YUV image
            YUV = np.stack((Y, U_upsampled, V_upsampled), axis=-1).astype(np.uint8)

            # Convert YUV to RGB
            rgb_image = ImagePIL.fromarray(YUV, "YCbCr").convert("RGB")

            # Save image as png
            img = rgb_image
        else:
            raise NotImplementedError(
                f"Message type {image_format} not supported for image conversion."
            )
    return img


def get_record_image_cutter(record_data, renormalize=False, **kwargs) -> ImagePIL.Image:
    image_format = record_data["type"]
    images_data = record_data["data"][0]
    if image_format != "JPG":
        raise NotImplementedError(
            f"Message type {image_format} not supported for image conversion."
        )
    img = ImagePIL.open(io.BytesIO(bytes(images_data)))
    return img


def get_record_image(
    record_data,
    max_size: Optional[int] = None,
    format: str = "WEBP",
    format_params: dict = {},
    renormalize: bool = False,
    resample=ImagePIL.Resampling.NEAREST,
    reset_position: bool = True,
    return_bytes: bool = False,
    **kwargs,
) -> Union[ImagePIL.Image, io.BytesIO, None]:
    if isinstance(record_data, bytes):
        img = get_record_image_lqs(
            record_data=record_data, renormalize=renormalize, **kwargs
        )
    else:
        _format = record_data.get("format", None)
        _encoding = record_data.get("encoding", None)
        _data_format = record_data.get("data_format", None)
        _image = record_data.get("image", None)
        if record_data.get("type") and record_data.get("cameras"):
            img = get_record_image_cutter(
                record_data=record_data, renormalize=renormalize, **kwargs
            )
        elif _encoding is not None:
            # ROS sensor_msgs/Image
            img = get_record_image_ros(
                record_data=record_data, renormalize=renormalize, **kwargs
            )
        elif _format is not None:
            # ROS sensor_msgs/CompressedImage
            img = get_record_image_ros(
                record_data=record_data, renormalize=renormalize, **kwargs
            )
        elif _data_format is not None:
            # ARK ark::image::Image
            img = get_record_image_ark(
                record_data=record_data, renormalize=renormalize, **kwargs
            )
        elif _image is not None:
            # ROS stereo_msgs/DisparityImage
            img = get_record_image_ros(
                record_data=_image, renormalize=renormalize, **kwargs
            )
        else:
            # we default to LQS image formats
            img = get_record_image_lqs(
                record_data=record_data, renormalize=renormalize, **kwargs
            )

    if max_size is not None:
        if img.width > max_size or img.height > max_size:
            img.thumbnail((max_size, max_size), resample=resample)

    if kwargs.get("use_rle_bmp", False):
        result = image_to_rle_bmp(img)
        return result

    if return_bytes:
        buffered_img = io.BytesIO()
        img.save(buffered_img, format=format, **format_params)
        if reset_position:
            buffered_img.seek(0)
        return buffered_img
    else:
        return img


# Point Cloud

point_cloud_type_names = [
    "sensor_msgs/PointCloud2",
    "sensor_msgs/msg/PointCloud2",
]


def jet_colormap(values, normalize=True):
    if normalize:
        values = (values - np.min(values)) / (np.max(values) - np.min(values))
    four_values = 4 * values
    red = np.minimum(four_values - 1.5, -four_values + 4.5)
    green = np.minimum(four_values - 0.5, -four_values + 3.5)
    blue = np.minimum(four_values + 0.5, -four_values + 2.5)
    colors = np.stack(
        [
            255 * np.clip(red, 0, 1),
            255 * np.clip(green, 0, 1),
            255 * np.clip(blue, 0, 1),
        ],
        axis=-1,
    )
    return colors.astype(np.uint8)


def voxel_downsample(positions, colors, max_points, voxel_size=1.0):
    # Calculate voxel indices for each point
    voxel_indices = np.floor(positions / voxel_size).astype(int)

    # Create a structured array to associate positions with their voxel
    dtype = [("x", int), ("y", int), ("z", int)]
    structured_voxel_indices = np.array(
        list(zip(voxel_indices[:, 0], voxel_indices[:, 1], voxel_indices[:, 2])),
        dtype=dtype,
    )

    # Find unique voxels and the indices of points in those voxels
    unique_voxels, inverse_indices, voxel_counts = np.unique(
        structured_voxel_indices, return_inverse=True, return_counts=True, axis=0
    )

    # Calculate sampling probabilities inversely proportional to voxel density
    sampling_probabilities = 1.0 / voxel_counts[inverse_indices]
    sampling_probabilities /= sampling_probabilities.sum()

    # Sample points based on calculated probabilities
    selected_indices = np.random.choice(
        np.arange(len(positions)),
        size=min(max_points, len(positions)),
        replace=False,
        p=sampling_probabilities,
    )

    return (
        positions[selected_indices],
        colors[selected_indices] if colors is not None else None,
    )


def process_pointcloud_record_data(
    record_data,
    max_points=None,
    random_sample=True,
    voxel_size=10.0,
    use_colormap=False,
    max_radius=None,
    force_colors=True,
    forced_color=[255, 0, 0],
    **kwargs,
):
    fields = record_data["fields"]
    valid_fields = ["x", "y", "z", "rgb", "bgr", "intensity"]
    fields = [f for f in record_data["fields"] if f["name"] in valid_fields]
    if len(fields) == 0:
        if len(record_data["data"]) > 0:
            raise ValueError(
                "No fields found in point cloud data with non-zero data length."
            )
        return [], []

    assert [f["name"] for f in fields[:3]] == [
        "x",
        "y",
        "z",
    ], f"Expected fields to be ['x', 'y', 'z'], but got {fields[:3]}"

    point_cloud_data = bytes(record_data["data"])
    dtype_map = {
        1: np.int8,
        2: np.uint8,
        3: np.int16,
        4: np.uint16,
        5: np.int32,
        6: np.uint32,
        7: np.float32,
        8: np.float64,
    }
    dtype_list = []
    current_offset = 0
    for field in fields:
        field_name = field["name"]
        field_type = dtype_map[field["datatype"]]
        field_offset = field["offset"]
        # Check if there is padding needed
        if field_offset > current_offset:
            padding_size = field_offset - current_offset
            dtype_list.append(("padding" + str(current_offset), np.uint8, padding_size))
        dtype_list.append((field_name, field_type))
        current_offset = field_offset + np.dtype(field_type).itemsize

    if current_offset < record_data["point_step"]:
        padding_size = record_data["point_step"] - current_offset
        dtype_list.append(("padding" + str(current_offset), np.uint8, padding_size))

    dtype = np.dtype(dtype_list)
    all_data = np.frombuffer(point_cloud_data, dtype=dtype)

    # Access positions directly via field names, assuming the first three fields are x, y, z
    positions = np.vstack([all_data["x"], all_data["y"], all_data["z"]]).T

    colors = None
    if len(fields) > 3 and fields[3]["name"] == "rgb":
        # Extract and process RGB colors directly from the structured array
        rgb_ints = all_data["rgb"].view(
            np.uint32
        )  # Assuming 'rgb' is stored as floats and can be viewed as uint32
        red = (rgb_ints >> 16) & 255
        green = (rgb_ints >> 8) & 255
        blue = rgb_ints & 255
        colors = np.vstack([red, green, blue]).T.astype(np.uint8)
    if len(fields) > 3 and fields[3]["name"] == "bgr":
        # Extract and process BGR colors directly from the structured array
        bgr_ints = all_data["bgr"].view(
            np.uint32
        )  # Assuming 'bgr' is stored as floats and can be viewed as uint32

        # NOTE: this seems to be the correct way to extract BGR colors,
        # but in testing, it's producing incorrect colors, so swapping to RGB extraction above
        # blue = (bgr_ints >> 16) & 255
        # green = (bgr_ints >> 8) & 255
        # red = bgr_ints & 255

        red = (bgr_ints >> 16) & 255
        green = (bgr_ints >> 8) & 255
        blue = bgr_ints & 255
        colors = np.vstack([red, green, blue]).T.astype(np.uint8)
    elif len(fields) > 3 and fields[3]["name"] == "intensity":
        # Process intensity data if needed
        colors = all_data["intensity"].astype(
            np.uint8
        )  # Assuming 'intensity' directly gives the desired data
        if use_colormap:
            colors = jet_colormap(colors)
        else:
            colors = np.stack([colors] * 3, axis=-1)

    if max_radius is not None:
        indices = np.linalg.norm(positions, axis=1) < max_radius
        positions = positions[indices]
        if colors is not None:
            colors = colors[indices]

    if max_points and len(positions) > max_points:
        if random_sample:
            indices = np.random.choice(len(positions), max_points, replace=False)
            positions = positions[indices]
            if colors is not None:
                colors = colors[indices]
        else:
            positions, colors = voxel_downsample(
                positions, colors, max_points, voxel_size=voxel_size
            )

    if colors is None and force_colors:
        if use_colormap:
            colors = jet_colormap(np.linalg.norm(positions, axis=1), normalize=True)
        else:
            colors = np.full((len(positions), 3), forced_color, dtype=np.uint8)

    return positions, colors


def encode_to_draco(
    positions,
    quant_bits=11,
    compression=7,
    remove_dups=True,
    colors=None,
    max_points=None,
):
    if len(positions) == 0:
        raise ValueError("Empty positions array for encoding to Draco.")

    draco_binary = draco.encode(
        positions,
        faces=None,
        quantization_bits=quant_bits,
        compression_level=0,
        quantization_range=-1,
        quantization_origin=None,
        create_metadata=False,
        preserve_order=False,
        colors=None if remove_dups else colors,
        tex_coord=None,
    )
    if remove_dups:
        _draco_points = draco.decode(draco_binary).points
        _, unique_indices = np.unique(_draco_points, axis=0, return_index=True)
        if max_points is not None and len(unique_indices) > max_points:
            unique_indices = np.random.choice(unique_indices, max_points, replace=False)

        _draco_points = _draco_points[unique_indices]
        if colors is not None:
            colors = colors[unique_indices]

        draco_binary = draco.encode(
            _draco_points,
            faces=None,
            quantization_bits=quant_bits,
            compression_level=compression,
            quantization_range=-1,
            quantization_origin=None,
            create_metadata=False,
            preserve_order=False,
            colors=colors,
            tex_coord=None,
        )
    return draco_binary


def decode_draco(draco_binary):
    return draco.decode(draco_binary).points, draco.decode(draco_binary).colors


# MCAP


class ModifiedChunk(Chunk):
    _chunk_start_offset: int
    _chunk_length: int

    def __init__(self, **kwargs):
        self._chunk_start_offset = kwargs.pop("_chunk_start_offset")
        self._chunk_length = kwargs.pop("_chunk_length")
        super().__init__(**kwargs)

    @staticmethod
    def read(stream: ReadDataStream):
        _chunk_start_offset = stream.count
        message_start_time = stream.read8()
        message_end_time = stream.read8()
        uncompressed_size = stream.read8()
        uncompressed_crc = stream.read4()
        compression_length = stream.read4()
        compression = str(stream.read(compression_length), "utf-8")
        data_length = stream.read8()
        if compression != "zstd":
            _chunk_start_offset = stream.count
            _chunk_length = data_length
        else:
            # we need to modify the chunk_length
            _chunk_length = (stream.count - _chunk_start_offset) + data_length
        data = stream.read(data_length)
        return ModifiedChunk(
            compression=compression,
            data=data,
            message_end_time=message_end_time,
            message_start_time=message_start_time,
            uncompressed_crc=uncompressed_crc,
            uncompressed_size=uncompressed_size,
            _chunk_start_offset=_chunk_start_offset,
            _chunk_length=_chunk_length,
        )
