"""Minimal, dependency-free reader for ROS 2 sqlite3 bags with sensor_msgs/PointCloud2.

Used for offline data exploration on Windows (no ROS installed).
Parses CDR (little-endian) encapsulated PointCloud2 messages.
"""
from __future__ import annotations

import sqlite3
import struct
import sys

import numpy as np

_DTYPES = {
    1: ("i1", 1), 2: ("u1", 1), 3: ("i2", 2), 4: ("u2", 2),
    5: ("i4", 4), 6: ("u4", 4), 7: ("f4", 4), 8: ("f8", 8),
}


class _Cdr:
    """Little-endian CDR reader with proper primitive alignment."""

    def __init__(self, buf: bytes):
        self.buf = buf
        self.origin = 4  # skip encapsulation header
        self.pos = 4

    def _align(self, size: int) -> None:
        rel = self.pos - self.origin
        pad = (-rel) % size
        self.pos += pad

    def u8(self) -> int:
        v = self.buf[self.pos]
        self.pos += 1
        return v

    def u32(self) -> int:
        self._align(4)
        v = struct.unpack_from("<I", self.buf, self.pos)[0]
        self.pos += 4
        return v

    def i32(self) -> int:
        self._align(4)
        v = struct.unpack_from("<i", self.buf, self.pos)[0]
        self.pos += 4
        return v

    def string(self) -> str:
        n = self.u32()
        s = self.buf[self.pos:self.pos + n - 1].decode("utf-8", "replace")
        self.pos += n
        return s

    def bytes(self, n: int) -> bytes:
        v = self.buf[self.pos:self.pos + n]
        self.pos += n
        return v


def parse_pointcloud2(blob: bytes) -> dict:
    c = _Cdr(blob)
    sec = c.i32()
    nsec = c.u32()
    frame_id = c.string()
    height = c.u32()
    width = c.u32()
    nfields = c.u32()
    fields = []
    for _ in range(nfields):
        name = c.string()
        offset = c.u32()
        datatype = c.u8()
        count = c.u32()
        fields.append((name, offset, datatype, count))
    is_bigendian = c.u8()
    point_step = c.u32()
    row_step = c.u32()
    n_bytes = c.u32()
    data = c.bytes(n_bytes)
    is_dense = c.u8()

    dt_fields = []
    used = 0
    for name, offset, datatype, count in fields:
        kind, size = _DTYPES[datatype]
        if offset > used:
            dt_fields.append((f"_pad{used}", f"V{offset - used}"))
        dt_fields.append((name, kind if count == 1 else f"{count}{kind}"))
        used = offset + size * count
    if point_step > used:
        dt_fields.append((f"_pad{used}", f"V{point_step - used}"))
    dtype = np.dtype([(n, t) for n, t in dt_fields])
    assert dtype.itemsize == point_step, (dtype.itemsize, point_step)
    arr = np.frombuffer(data, dtype=dtype, count=height * width)
    return dict(stamp=sec + nsec * 1e-9, frame_id=frame_id, height=height, width=width,
                fields=fields, point_step=point_step, row_step=row_step,
                is_dense=is_dense, is_bigendian=is_bigendian, points=arr)


def frames(db_path: str, limit: int | None = None, start: int = 0):
    con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
    q = "SELECT timestamp, data FROM messages ORDER BY timestamp"
    if limit is not None:
        q += f" LIMIT {limit} OFFSET {start}"
    for ts, blob in con.execute(q):
        yield ts, parse_pointcloud2(blob)
    con.close()


if __name__ == "__main__":
    path = sys.argv[1]
    idx = int(sys.argv[2]) if len(sys.argv) > 2 else 0
    for ts, m in frames(path, limit=1, start=idx):
        print("stamp", m["stamp"], "frame_id", m["frame_id"])
        print("height", m["height"], "width", m["width"], "point_step", m["point_step"],
              "dense", m["is_dense"], "bigendian", m["is_bigendian"])
        print("fields:")
        for f in m["fields"]:
            print("   ", f)
        p = m["points"]
        print("npoints", p.shape)
        for name in p.dtype.names:
            if name.startswith("_pad"):
                continue
            v = p[name]
            print(f"  {name:12s} dtype={v.dtype} min={np.min(v)} max={np.max(v)} mean={np.mean(v.astype(np.float64)):.4f}")