//! Recorded depth datasets: the on-disk format, writer, and loader.
//!
//! A recording is a directory:
//!
//! ```text
//! my-recording/
//! ├── recording.toml     # the depth camera (id, pose, intrinsics,
//! │                     #   resolution), fps, frame count
//! ├── frame_0000.depth   # raw little-endian u16 depth, width*height
//! ├── frame_0000.rgb    # raw RGB bytes, width*height*3
//! ├── frame_0001.depth
//! └── ...
//! ```
//!
//! Real Kinect recordings can be converted into this shape; for dev and
//! CI, [`Recording::write`] produces recordings from the synthetic
//! source. [`RecordingSource`] plays them back as a [`DepthSource`],
//! looping forever.

use std::fs;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};

use rig::{Camera, CameraIntrinsics, CameraPose, DepthFrame, DepthSource};

/// `recording.toml`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Manifest {
    pub camera: CameraManifest,
    pub fps: f32,
    pub frames: u32,
}

/// The depth camera a recording was made with.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CameraManifest {
    pub id: String,
    /// Camera position, metres.
    pub position: [f32; 3],
    /// Camera-to-world rotation, `x,y,z,w` quaternion.
    pub rotation: [f32; 4],
    /// `(width, height)` in pixels.
    pub resolution: [u32; 2],
    /// `fx, fy, cx, cy` in pixels.
    pub intrinsics: [f32; 4],
}

impl From<&Camera> for CameraManifest {
    fn from(camera: &Camera) -> Self {
        Self {
            id: camera.id.clone(),
            position: camera.pose.position.to_array(),
            rotation: camera.pose.rotation.to_array(),
            resolution: [camera.resolution.x, camera.resolution.y],
            intrinsics: [
                camera.intrinsics.fx,
                camera.intrinsics.fy,
                camera.intrinsics.cx,
                camera.intrinsics.cy,
            ],
        }
    }
}

impl From<CameraManifest> for Camera {
    fn from(manifest: CameraManifest) -> Self {
        Self {
            id: manifest.id,
            pose: CameraPose {
                position: manifest.position.into(),
                rotation: glam::Quat::from_array(manifest.rotation),
            },
            intrinsics: CameraIntrinsics {
                fx: manifest.intrinsics[0],
                fy: manifest.intrinsics[1],
                cx: manifest.intrinsics[2],
                cy: manifest.intrinsics[3],
            },
            resolution: glam::UVec2::new(manifest.resolution[0], manifest.resolution[1]),
        }
    }
}

fn depth_file(frame: u32) -> String {
    format!("frame_{frame:04}.depth")
}

fn rgb_file(frame: u32) -> String {
    format!("frame_{frame:04}.rgb")
}

/// Writes recordings: [`Recording::write`].
pub struct Recording;

impl Recording {
    /// Record `frames` frames from `source` into `dir`, replacing it.
    ///
    /// The source must yield exactly the frame stream worth keeping;
    /// its first `frames` frames become the recording.
    pub fn write(
        dir: impl AsRef<Path>,
        source: &mut dyn DepthSource,
        fps: f32,
        frames: u32,
    ) -> Result<()> {
        let dir = dir.as_ref();
        fs::create_dir_all(dir).with_context(|| format!("creating {}", dir.display()))?;
        let manifest = Manifest {
            camera: CameraManifest::from(source.camera()),
            fps,
            frames,
        };
        fs::write(
            dir.join("recording.toml"),
            toml::to_string_pretty(&manifest)?,
        )
        .context("writing recording.toml")?;

        for frame in 0..frames {
            let depth_frame = source
                .next_depth_frame()
                .with_context(|| format!("recording frame {frame}"))?;
            let depth_le: Vec<u8> = depth_frame
                .depth_mm
                .iter()
                .flat_map(|d| d.to_le_bytes())
                .collect();
            fs::write(dir.join(depth_file(frame)), depth_le)?;
            fs::write(dir.join(rgb_file(frame)), &depth_frame.rgb)?;
        }
        Ok(())
    }
}

/// Plays a recording back as a [`DepthSource`], looping forever.
#[derive(Debug)]
pub struct RecordingSource {
    dir: PathBuf,
    camera: Camera,
    fps: f32,
    frames: u32,
    frame_index: u32,
}

impl RecordingSource {
    /// Load the recording in `dir`.
    pub fn load(dir: impl AsRef<Path>) -> Result<Self> {
        let dir = dir.as_ref().to_path_buf();
        let manifest_path = dir.join("recording.toml");
        let manifest: Manifest = toml::from_str(
            &fs::read_to_string(&manifest_path)
                .with_context(|| format!("reading {}", manifest_path.display()))?,
        )
        .with_context(|| format!("parsing {}", manifest_path.display()))?;

        if manifest.frames == 0 {
            bail!("recording in {} has no frames", dir.display());
        }
        if manifest.fps <= 0.0 {
            bail!("recording in {} has non-positive fps", dir.display());
        }

        Ok(Self {
            camera: manifest.camera.into(),
            fps: manifest.fps,
            frames: manifest.frames,
            dir,
            frame_index: 0,
        })
    }

    /// The recording's frame count.
    pub fn frames(&self) -> u32 {
        self.frames
    }

    /// The recording's playback rate.
    pub fn fps(&self) -> f32 {
        self.fps
    }
}

impl DepthSource for RecordingSource {
    fn camera(&self) -> &Camera {
        &self.camera
    }

    fn next_depth_frame(&mut self) -> Result<DepthFrame> {
        let frame = self.frame_index % self.frames;
        self.frame_index += 1;

        let (width, height) = (self.camera.resolution.x, self.camera.resolution.y);
        let depth_path = self.dir.join(depth_file(frame));
        let rgb_path = self.dir.join(rgb_file(frame));

        let depth_bytes =
            fs::read(&depth_path).with_context(|| format!("reading {}", depth_path.display()))?;
        let rgb_bytes =
            fs::read(&rgb_path).with_context(|| format!("reading {}", rgb_path.display()))?;

        let expected_depth = (width * height) as usize * 2;
        if depth_bytes.len() != expected_depth {
            bail!(
                "{}: expected {expected_depth} depth bytes, got {}",
                depth_path.display(),
                depth_bytes.len()
            );
        }
        let expected_rgb = (width * height) as usize * 3;
        if rgb_bytes.len() != expected_rgb {
            bail!(
                "{}: expected {expected_rgb} rgb bytes, got {}",
                rgb_path.display(),
                rgb_bytes.len()
            );
        }

        let depth_mm: Vec<u16> = depth_bytes
            .as_chunks::<2>()
            .0
            .iter()
            .map(|b| u16::from_le_bytes(*b))
            .collect();

        Ok(DepthFrame {
            timestamp_secs: frame as f64 / self.fps as f64,
            depth_mm,
            rgb: rgb_bytes,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::synth::{self, NoiseParams};
    use glam::uvec2;

    /// A tiny clean source for exact comparisons.
    fn tiny_clean_source() -> synth::SyntheticSource {
        let rig = synth::dev_rig_with_resolutions(uvec2(32, 24), uvec2(64, 48));
        synth::dev_source_with(rig.center, synth::SUBJECT, 30.0).with_noise(NoiseParams::OFF)
    }

    fn temp_dir(name: &str) -> PathBuf {
        let dir =
            std::env::temp_dir().join(format!("kinect-view-rec-{name}-{}", std::process::id()));
        let _ = fs::remove_dir_all(&dir);
        dir
    }

    #[test]
    fn round_trip_preserves_frames() {
        let dir = temp_dir("roundtrip");
        let mut source = tiny_clean_source();
        Recording::write(&dir, &mut source, 30.0, 3).expect("write recording");

        let mut player = RecordingSource::load(&dir).expect("load recording");
        assert_eq!(player.frames(), 3);
        assert_eq!(player.fps(), 30.0);
        assert_eq!(player.camera(), source.camera());

        // The recorded frames match what the source produced, in order.
        let mut source = tiny_clean_source();
        for _ in 0..3 {
            let recorded = player.next_depth_frame().expect("frame");
            let original = source.next_depth_frame().expect("frame");
            assert_eq!(recorded.depth_mm, original.depth_mm);
            assert_eq!(recorded.rgb, original.rgb);
            assert_eq!(recorded.timestamp_secs, original.timestamp_secs);
        }

        // And playback loops.
        let wrapped = player.next_depth_frame().expect("frame");
        assert_eq!(wrapped.timestamp_secs, 0.0);

        fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn missing_files_are_reported() {
        let err = RecordingSource::load("/nonexistent/recording").expect_err("missing dir");
        assert!(err.to_string().contains("reading"));

        // A manifest with frames but no frame files.
        let dir = temp_dir("missing-frames");
        fs::create_dir_all(&dir).expect("dir");
        fs::write(
            dir.join("recording.toml"),
            r#"
[[fake]]"#,
        )
        .expect("write");
        let err = RecordingSource::load(&dir).expect_err("bad manifest");
        assert!(err.to_string().contains("parsing"));
        fs::remove_dir_all(&dir).ok();
    }
}