//! Live capture from the physical rig: a Kinect v2 (depth + registered
//! RGB, via libfreenect2 through [`kinect2_sys`]) flanked by two webcams
//! ([`nokhwa`]), composed into raw [`rig::DepthFrame`]s.
//!
//! Each sensor runs on its own thread; frames are timestamp-aligned by
//! pairing each Kinect frame with the *latest* webcam frames (the plan
//! tolerates ~1 frame of skew). The Kinect's factory calibration
//! provides the depth camera's intrinsics (512x424, the registered
//! pixels' own convention); the config's center-camera pose — the
//! offline calibration's extrinsics — still applies.
//!
//! **Hardware note:** none of the device paths run in the dev container
//! — there is no USB. The composition, device-matching, and conversion
//! logic is unit-tested with synthetic buffers; the devices themselves
//! are exercised on real hardware. [`KinectRigSource::open`] fails fast
//! with a clear error when no Kinect is present.
//!
//! Platform notes (see the README's "Kinect v2 capture" section):
//!
//! * **Windows** — libfreenect2 works over the UsbDk or libusbK driver;
//!   kinect-view links a prebuilt libfreenect2 via `FREENECT2_DIR`.
//! * **Linux** — libfreenect2 over libusb, found via pkg-config or
//!   `/usr/local`; a udev rule grants device access.

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
    mpsc::{Receiver, SyncSender, sync_channel},
};
use std::thread::JoinHandle;
use std::time::Instant;

use anyhow::{Context, Result, bail};
use glam::UVec2;

use kinect2_sys::{
    Device as Freenect2Device, FrameFormat, Listener as Freenect2Listener, Registration,
};
use rig::{Camera, CameraIntrinsics, DepthFrame, DepthSource, Rig, SideRgb};

/// A composed frame from the capture threads.
struct PairedFrame {
    depth_mm: Vec<u16>,
    rgb: Vec<u8>,
    left: Option<Vec<u8>>,
    right: Option<Vec<u8>>,
    timestamp_secs: f64,
}

/// The live rig source: Kinect v2 + two webcams.
pub struct KinectRigSource {
    rig: Rig,
    receiver: Receiver<PairedFrame>,
    /// Set on drop; the worker threads check it and exit.
    shutdown: Arc<AtomicBool>,
    kinect_worker: Option<JoinHandle<()>>,
    /// Kept alive so their capture threads keep streaming; stopping
    /// on drop.
    webcams: Vec<Webcam>,
}

impl KinectRigSource {
    /// Open the rig: the Kinect plus the two webcams named by the rig
    /// config's camera ids (a numeric id is a device index; otherwise
    /// the id must match the camera's human-readable name).
    ///
    /// The rig's center camera is *replaced* by the Kinect's factory
    /// depth calibration (512x424 at the depth pixels' own intrinsics);
    /// its pose comes from the config. The returned
    /// [`DepthSource::rig`] is the effective rig the pipeline must
    /// use.
    pub fn open(rig: &Rig) -> Result<Self> {
        let shutdown = Arc::new(AtomicBool::new(false));

        // Kinect first, so a missing device fails before any webcam is
        // touched. The worker reports the effective depth camera back;
        // the webcam receivers are handed to it once opened.
        let (ready_tx, ready_rx) = std::sync::mpsc::channel::<Result<Camera>>();
        let (frames_tx, frames_rx) = sync_channel::<PairedFrame>(2);
        let (side_tx, side_rx) = sync_channel::<(Receiver<Vec<u8>>, Receiver<Vec<u8>>)>(1);
        let kinect_worker = {
            let shutdown = shutdown.clone();
            let center_config = rig.center.clone();
            std::thread::spawn(move || {
                kinect_loop(center_config, ready_tx, side_rx, frames_tx, &shutdown);
            })
        };
        // If the Kinect fails before the ready signal, this surfaces it.
        let center = ready_rx
            .recv()
            .map_err(|_| anyhow::anyhow!("the Kinect worker stopped before starting"))??;

        // The webcams: each streams decoded RGB into its channel from
        // nokhwa's own capture thread.
        let mut webcams = Vec::new();
        let mut side_channels = Vec::new();
        for config in [&rig.left, &rig.right] {
            let (tx, rx) = sync_channel(2);
            let webcam = Webcam::open(config, tx)
                .with_context(|| format!("opening webcam {:?}", config.id))?;
            webcams.push(webcam);
            side_channels.push(rx);
        }
        let mut it = side_channels.into_iter();
        let (left_rx, right_rx) = (
            it.next().expect("two side cameras"),
            it.next().expect("two side cameras"),
        );
        side_tx
            .send((left_rx, right_rx))
            .map_err(|_| anyhow::anyhow!("the Kinect worker stopped"))?;

        let rig = Rig {
            center,
            left: rig.left.clone(),
            right: rig.right.clone(),
        };
        Ok(Self {
            rig,
            receiver: frames_rx,
            shutdown,
            kinect_worker: Some(kinect_worker),
            webcams,
        })
    }
}

impl Drop for KinectRigSource {
    fn drop(&mut self) {
        self.shutdown.store(true, Ordering::Relaxed);
        if let Some(worker) = self.kinect_worker.take() {
            let _ = worker.join();
        }
        // The webcams stop their own capture threads on drop.
        self.webcams.clear();
    }
}

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

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

    fn next_depth_frame(&mut self) -> Result<DepthFrame> {
        let paired = self
            .receiver
            .recv()
            .map_err(|_| anyhow::anyhow!("the Kinect stopped producing frames"))?;
        Ok(DepthFrame {
            timestamp_secs: paired.timestamp_secs,
            depth_mm: paired.depth_mm,
            rgb: paired.rgb,
            side_rgb: Some(SideRgb {
                left: paired.left.unwrap_or_default(),
                left_resolution: self.rig.left.resolution,
                right: paired.right.unwrap_or_default(),
                right_resolution: self.rig.right.resolution,
            }),
        })
    }
}

/// The webcam receivers, one per side camera.
type SideChannels = (Receiver<Vec<u8>>, Receiver<Vec<u8>>);

/// The Kinect v2 worker: opens the device through libfreenect2, then
/// blocks on its synchronized color+depth listener, registers the pair,
/// and composes frames with the latest webcam images.
///
/// `ready` receives the effective depth camera (factory intrinsics,
/// config pose) once the device is up — or the open error.
fn kinect_loop(
    center_config: Camera,
    ready: std::sync::mpsc::Sender<Result<Camera>>,
    side_handoff: Receiver<SideChannels>,
    frames: SyncSender<PairedFrame>,
    shutdown: &Arc<AtomicBool>,
) {
    let result = (|| -> Result<()> {
        let device = Freenect2Device::open_default()
            .map_err(|e| anyhow::anyhow!("opening the Kinect failed — is it plugged in? {e}"))?;
        device.start().context("starting the Kinect")?;
        log::info!(
            "Kinect v2 open (serial {}), {} device(s) connected",
            device.serial(),
            Freenect2Device::count().unwrap_or(1)
        );

        // The depth camera: factory intrinsics at 512x424, pose from
        // the offline calibration.
        let ir = device.ir_params();
        let camera = Camera {
            id: center_config.id.clone(),
            pose: center_config.pose,
            intrinsics: CameraIntrinsics {
                fx: ir.fx(),
                fy: ir.fy(),
                cx: ir.cx(),
                cy: ir.cy(),
            },
            resolution: UVec2::new(
                kinect2_sys::DEPTH_WIDTH as u32,
                kinect2_sys::DEPTH_HEIGHT as u32,
            ),
        };
        ready
            .send(Ok(camera))
            .map_err(|_| anyhow::anyhow!("the app went away"))?;

        let mut listener =
            Freenect2Listener::attach(&device).context("attaching the frame listener")?;
        let mut registration =
            Registration::from_device(&device).context("creating the registration")?;

        // The webcam receivers, handed over once the app has opened
        // the side cameras.
        let (left_rx, right_rx) = side_handoff
            .recv()
            .map_err(|_| anyhow::anyhow!("the app went away"))?;

        let started = Instant::now();
        while !shutdown.load(Ordering::Relaxed) {
            let composed = listener.with_frames(1000, |color, depth| {
                // Registered pair: undistorted depth (float mm) and the
                // color image resampled onto the depth grid.
                registration
                    .apply(color, depth)
                    .ok()
                    .map(|(undistorted, registered)| {
                        // SAFETY: the views are alive for the closure body.
                        let depth_mm = convert_depth(unsafe { undistorted.floats() });
                        let rgb = convert_registered(unsafe { registered.bytes() }, color.format());
                        PairedFrame {
                            depth_mm,
                            rgb,
                            left: latest(&left_rx),
                            right: latest(&right_rx),
                            timestamp_secs: started.elapsed().as_secs_f64(),
                        }
                    })
            });
            match composed {
                Ok(None) | Ok(Some(None)) => continue, // timeout or a malformed pair
                Ok(Some(Some(frame))) => {
                    if frames.send(frame).is_err() {
                        break; // receiver dropped: shut down
                    }
                }
                Err(e) => {
                    log::error!("kinect worker: {e:#}");
                    break;
                }
            }
        }
        let _ = device.stop();
        Ok(())
    })();
    if let Err(e) = result {
        log::error!("kinect worker: {e:#}");
        let _ = ready.send(Err(e));
    }
}

/// libfreenect2's undistorted depth: float millimetres (non-positive,
/// NaN, infinity are invalid) → u16 millimetres with 0 for invalid.
fn convert_depth(floats: &[f32]) -> Vec<u16> {
    floats
        .iter()
        .map(|&d| {
            if d.is_finite() && (500.0..=6000.0).contains(&d) {
                d as u16
            } else {
                0
            }
        })
        .collect()
}

/// The registered color frame: 4 bytes per pixel at depth resolution,
/// in the color pipeline's channel order (BGRX or RGBX) — reorder to
/// the RGB our frames carry.
fn convert_registered(bytes: &[u8], color_format: FrameFormat) -> Vec<u8> {
    let mut rgb = vec![0_u8; bytes.len()];
    match color_format {
        FrameFormat::Rgbx => rgb.copy_from_slice(bytes),
        FrameFormat::Bgrx => {
            let (rgb_chunks, _) = rgb.as_chunks_mut::<4>();
            let (byte_chunks, _) = bytes.as_chunks::<4>();
            for (dst, px) in rgb_chunks.iter_mut().zip(byte_chunks) {
                let (dst, px): (&mut [u8; 4], &[u8; 4]) = (dst, px);
                let [r, g, b, a] = *px;
                *dst = [b, g, r, a];
            }
        }
        FrameFormat::Float | FrameFormat::Other(_) => {
            log::warn!("unexpected color frame format {color_format:?}; passing through");
            rgb.copy_from_slice(bytes);
        }
    }
    rgb
}

/// The most recent frame from a webcam channel, if one arrived.
fn latest(channel: &Receiver<Vec<u8>>) -> Option<Vec<u8>> {
    let mut newest = None;
    while let Ok(frame) = channel.try_recv() {
        newest = Some(frame);
    }
    newest
}

/// A webcam: a nokhwa [`CallbackCamera`] streaming decoded RGB into a
/// channel from its own capture thread (dropping it stops the stream).
struct Webcam {
    /// Held for its side effect: dropping it stops the capture thread.
    #[allow(dead_code)]
    camera: nokhwa::CallbackCamera,
}

impl Webcam {
    fn open(config: &Camera, frames: SyncSender<Vec<u8>>) -> Result<Self> {
        use nokhwa::pixel_format::RgbFormat;
        use nokhwa::utils::{
            ApiBackend, CameraFormat, FrameFormat, RequestedFormat, RequestedFormatType, Resolution,
        };

        // Match the device: numeric id → index, otherwise the
        // human-readable name.
        let devices = nokhwa::query(ApiBackend::Auto)
            .map_err(|e| anyhow::anyhow!("querying cameras: {e}"))?;
        let index = device_index(&devices, &config.id)
            .with_context(|| format!("finding webcam {:?}", config.id))?;

        let format = CameraFormat::new(
            Resolution::new(config.resolution.x, config.resolution.y),
            FrameFormat::MJPEG,
            30,
        );
        let mut camera = nokhwa::CallbackCamera::new(
            index,
            RequestedFormat::new::<RgbFormat>(RequestedFormatType::Closest(format)),
            move |buffer| {
                // Decode on the capture thread and hand the RGB over;
                // a full channel just drops the frame (we always want
                // the latest).
                if let Ok(image) = buffer.decode_image::<RgbFormat>() {
                    let _ = frames.try_send(image.into_raw());
                }
            },
        )
        .map_err(|e| anyhow::anyhow!("opening webcam {:?}: {e}", config.id))?;

        // The camera must deliver the configured resolution — the rig
        // intrinsics are calibrated for it.
        let actual = camera
            .camera_format()
            .map_err(|e| anyhow::anyhow!("webcam {:?} format: {e}", config.id))?;
        if actual.resolution().width_x != config.resolution.x
            || actual.resolution().height_y != config.resolution.y
        {
            bail!(
                "webcam {:?} delivers {}x{}, but the rig is calibrated for {}x{}",
                config.id,
                actual.resolution().width_x,
                actual.resolution().height_y,
                config.resolution.x,
                config.resolution.y
            );
        }
        camera
            .open_stream()
            .map_err(|e| anyhow::anyhow!("starting webcam {:?}: {e}", config.id))?;

        Ok(Self { camera })
    }
}

/// Resolve a configured camera id to a nokhwa device index: a numeric
/// id is the index itself; otherwise it must (uniquely) match a
/// device's human-readable name.
fn device_index(
    devices: &[nokhwa::utils::CameraInfo],
    id: &str,
) -> Result<nokhwa::utils::CameraIndex> {
    // Numeric id: use it directly.
    if let Ok(index) = id.parse::<u32>() {
        return Ok(nokhwa::utils::CameraIndex::Index(index));
    }
    let matches: Vec<_> = devices
        .iter()
        .filter(|d| d.human_name().to_lowercase().contains(&id.to_lowercase()))
        .collect();
    match matches.as_slice() {
        [device] => Ok(device.index().clone()),
        [] => bail!(
            "no camera name matches {id:?} (have: {:?})",
            devices.iter().map(|d| d.human_name()).collect::<Vec<_>>()
        ),
        _ => bail!(
            "camera id {id:?} is ambiguous: {:?}",
            matches.iter().map(|d| d.human_name()).collect::<Vec<_>>()
        ),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use glam::Vec3;
    use glam::uvec2;
    use rig::CameraPose;

    #[test]
    fn depth_conversion_drops_invalid_and_quantizes() {
        // libfreenect2 depth: float millimetres, invalid as
        // non-positive/NaN/infinity; out-of-range values are dropped.
        let floats = [
            1900.0,
            0.0,
            -2.0,
            f32::NAN,
            f32::INFINITY,
            2845.6,
            250.0,
            6500.0,
        ];
        let expected = [1900, 0, 0, 0, 0, 2845, 0, 0];
        assert_eq!(convert_depth(&floats), expected);
    }

    #[test]
    fn registered_colour_reordering() {
        // BGRX in → RGB out.
        let bgrx = [1_u8, 2, 3, 255, 5, 6, 7, 255];
        assert_eq!(
            convert_registered(&bgrx, FrameFormat::Bgrx),
            vec![3, 2, 1, 255, 7, 6, 5, 255]
        );
        // RGBX passes through.
        let rgbx = [1_u8, 2, 3, 255, 5, 6, 7, 255];
        assert_eq!(
            convert_registered(&rgbx, FrameFormat::Rgbx),
            vec![1, 2, 3, 255, 5, 6, 7, 255]
        );
    }

    #[test]
    fn latest_takes_only_the_newest_frame() {
        let (tx, rx) = sync_channel(4);
        tx.send(vec![1]).unwrap();
        tx.send(vec![2]).unwrap();
        tx.send(vec![3]).unwrap();
        assert_eq!(latest(&rx), Some(vec![3]));
        assert_eq!(latest(&rx), None);
    }

    #[test]
    fn device_matching_by_index_and_name() {
        let camera = |index: u32, name: &str| {
            nokhwa::utils::CameraInfo::new(name, "", "", nokhwa::utils::CameraIndex::Index(index))
        };
        let devices = vec![camera(0, "Logitech C920"), camera(1, "Some Other Cam")];

        // Numeric id → index.
        assert_eq!(
            device_index(&devices, "1").unwrap(),
            nokhwa::utils::CameraIndex::Index(1)
        );
        // Name match (case-insensitive substring).
        assert_eq!(
            device_index(&devices, "c920").unwrap(),
            nokhwa::utils::CameraIndex::Index(0)
        );
        // No match.
        assert!(device_index(&devices, "nonexistent").is_err());
    }

    #[test]
    fn opening_without_a_kinect_fails_clearly() {
        // In the container there is no USB, so opening must fail with
        // the "is it plugged in?" error rather than hang or panic.
        let rig = rig::Rig {
            center: Camera {
                id: "kinect".into(),
                pose: CameraPose::looking_at(Vec3::ZERO, Vec3::new(0.0, 0.0, 2.0)),
                intrinsics: CameraIntrinsics {
                    fx: 365.0,
                    fy: 365.0,
                    cx: 256.0,
                    cy: 212.0,
                },
                resolution: uvec2(512, 424),
            },
            left: Camera {
                id: "left".into(),
                pose: CameraPose::looking_at(Vec3::new(-0.4, 0.0, 0.0), Vec3::new(0.0, 0.0, 2.0)),
                intrinsics: CameraIntrinsics {
                    fx: 554.0,
                    fy: 554.0,
                    cx: 320.0,
                    cy: 240.0,
                },
                resolution: uvec2(640, 480),
            },
            right: Camera {
                id: "right".into(),
                pose: CameraPose::looking_at(Vec3::new(0.4, 0.0, 0.0), Vec3::new(0.0, 0.0, 2.0)),
                intrinsics: CameraIntrinsics {
                    fx: 554.0,
                    fy: 554.0,
                    cx: 320.0,
                    cy: 240.0,
                },
                resolution: uvec2(640, 480),
            },
        };
        if kinect2_sys::Device::count().unwrap_or(0) == 0 {
            let err = match KinectRigSource::open(&rig) {
                Ok(_) => panic!("opened a rig with no Kinect in the container?"),
                Err(e) => e,
            };
            assert!(
                err.to_string().to_lowercase().contains("kinect"),
                "unexpected error: {err:#}"
            );
        }
    }
}