//! Recorded depth datasets: the on-disk format, writer, and loader.
//!
//! A recording is a directory:
//!
//! ```text
//! my-recording/
//! ├── recording.toml     # the rig's cameras (center + optional left /
//! │                     #   right), fps, frame count
//! ├── frame_0000.depth   # raw little-endian u16 depth, width*height
//! ├── frame_0000.rgb     # raw RGB bytes, width*height*3 (registered)
//! ├── frame_0000.left    # optional: left webcam RGB, w*h*3
//! ├── frame_0000.right   # optional: right webcam RGB, w*h*3
//! └── ...
//! ```
//!
//! Real rig 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. The full recorded rig (center + side cameras) is
//! available via [`RecordingSource::rig`] — the pipeline needs the side
//! cameras for the occlusion proxy.

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

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

use rig::{Camera, DepthFrame, DepthSource, Rig, SideRgb};

/// `recording.toml`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Manifest {
    /// The depth (center) camera.
    pub center: CameraManifest,
    /// The left webcam, if the recording has side RGB.
    pub left: Option<CameraManifest>,
    /// The right webcam, if the recording has side RGB.
    pub right: Option<CameraManifest>,
    pub fps: f32,
    pub frames: u32,
}

/// A 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: rig::CameraPose {
                position: manifest.position.into(),
                rotation: glam::Quat::from_array(manifest.rotation),
            },
            intrinsics: rig::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")
}

fn side_file(frame: u32, side: &str) -> String {
    format!("frame_{frame:04}.{side}")
}

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

impl Recording {
    /// Record `frames` frames from `source` into `dir`, replacing it.
    ///
    /// The rig (center + side cameras) comes from the source; side RGB
    /// is stored when the source provides it.
    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 rig = source.rig();
        let manifest = Manifest {
            center: CameraManifest::from(&rig.center),
            left: Some(CameraManifest::from(&rig.left)),
            right: Some(CameraManifest::from(&rig.right)),
            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)?;
            if let Some(sides) = &depth_frame.side_rgb {
                fs::write(dir.join(side_file(frame, "left")), &sides.left)?;
                fs::write(dir.join(side_file(frame, "right")), &sides.right)?;
            }
        }
        Ok(())
    }
}

/// Plays a recording back as a [`DepthSource`], looping forever.
#[derive(Debug)]
pub struct RecordingSource {
    dir: PathBuf,
    rig: Rig,
    fps: f32,
    frames: u32,
    has_sides: bool,
    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());
        }

        let has_left = manifest.left.is_some();
        let has_right = manifest.right.is_some();
        let rig = Rig {
            center: manifest.center.into(),
            left: manifest.left.map(Into::into).unwrap_or_else(|| Camera {
                id: "left".into(),
                pose: rig::CameraPose::identity(),
                intrinsics: rig::CameraIntrinsics {
                    fx: 1.0,
                    fy: 1.0,
                    cx: 0.0,
                    cy: 0.0,
                },
                resolution: glam::UVec2::ONE,
            }),
            right: manifest.right.map(Into::into).unwrap_or_else(|| Camera {
                id: "right".into(),
                pose: rig::CameraPose::identity(),
                intrinsics: rig::CameraIntrinsics {
                    fx: 1.0,
                    fy: 1.0,
                    cx: 0.0,
                    cy: 0.0,
                },
                resolution: glam::UVec2::ONE,
            }),
        };
        // Side RGB is present only when both cameras and both files for
        // frame 0 exist.
        let has_sides = has_left
            && has_right
            && dir.join(side_file(0, "left")).exists()
            && dir.join(side_file(0, "right")).exists();

        Ok(Self {
            rig,
            fps: manifest.fps,
            frames: manifest.frames,
            has_sides,
            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
    }

    /// The recorded rig: center + side cameras.
    pub fn rig(&self) -> &Rig {
        &self.rig
    }
}

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

    fn rig(&self) -> Rig {
        self.rig.clone()
    }

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

        let (width, height) = (self.rig.center.resolution.x, self.rig.center.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();

        let side_rgb = if self.has_sides {
            let read_side = |side: &str, camera: &Camera| -> Result<Vec<u8>> {
                let path = self.dir.join(side_file(frame, side));
                let bytes =
                    fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
                let expected = (camera.resolution.x * camera.resolution.y) as usize * 3;
                if bytes.len() != expected {
                    bail!(
                        "{}: expected {expected} rgb bytes, got {}",
                        path.display(),
                        bytes.len()
                    );
                }
                Ok(bytes)
            };
            Some(SideRgb {
                left: read_side("left", &self.rig.left)?,
                left_resolution: self.rig.left.resolution,
                right: read_side("right", &self.rig.right)?,
                right_resolution: self.rig.right.resolution,
            })
        } else {
            None
        };

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

#[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, 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_and_rig() {
        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);
        // The recorded rig round-trips: center + side cameras.
        let recorded = player.rig();
        assert_eq!(recorded.center.pose, source.rig().center.pose);
        assert_eq!(recorded.left.pose, source.rig().left.pose);
        assert_eq!(recorded.right.intrinsics, source.rig().right.intrinsics);

        // The recorded frames match what the source produced, in order,
        // including the side RGB.
        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.side_rgb, original.side_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"), "[[fake]]").expect("write");
        let err = RecordingSource::load(&dir).expect_err("bad manifest");
        assert!(err.to_string().contains("parsing"));
        fs::remove_dir_all(&dir).ok();
    }
}