import json
from pathlib import Path

import cv2
import numpy as np
from nvidia import nvimgcodec

decoder = nvimgcodec.Decoder()
encoder = nvimgcodec.Encoder()


rec_folder = Path("/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/record2/")

robot_folder = rec_folder / "robot"
robot_folder = rec_folder / "robot"

d: dict[int, dict] = dict()
for file in robot_folder.glob("*"):
    for line in file.read_text().splitlines():
        s = line.split(",")
        v = dict(
            host_ts=int(s[0]),
            device_ts=int(float(s[1]) * 10**9),
            x=float(s[2]),
            y=float(s[3]),
            z=float(s[4]),
            rx=float(s[5]),
            ry=float(s[6]),
            rz=float(s[7]),
            fx=float(s[8]),
            fy=float(s[9]),
            fz=float(s[10]),
            frx=float(s[11]),
            fry=float(s[12]),
            frz=float(s[13]),
        )
        d[v["host_ts"]] = v
(rec_folder / "robot.json").write_text(json.dumps(d))

for file in Path(
    "/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/record2/camera_raw"
).glob("*.bin"):
    fb = Path(file).read_bytes()
    b = np.frombuffer(fb, dtype=np.uint8).reshape((1080, 1920, 3))
    b = cv2.cvtColor(b, cv2.COLOR_BGR2RGB)
    # cv2.imshow('',)
    # cv2.waitKey(0)
    # enc_params = nvimgcodec.EncodeParams(quality=5, chroma_subsampling=nvimgcodec.ChromaSubsampling.CSS_GRAY)
    # encoder.write('a.jpg', b, params=enc_params)
    encoder.write(
        (file.parent.parent / "camera" / file.with_suffix(".jpg").name).__str__(), b
    )