//! Live capture from the physical rig: a Kinect v2 (depth + registered
//! RGB, via [`kinect_one`]) 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 this can run in the container — there is
//! no USB. The composition, device-matching, and conversion logic is
//! unit-tested with synthetic buffers; the device paths are exercised
//! only on real hardware. [`KinectRigSource::open`] fails fast with a
//! clear error when no Kinect is present.

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 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) =
            std::sync::mpsc::channel::<(Receiver<Vec<u8>>, Receiver<Vec<u8>>)>();
        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, then polls packets, decodes,
/// registers, 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 runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .context("building the Kinect runtime")?;
        runtime.block_on(async {
            let mut device = kinect_one::DeviceEnumerator::open_default(true)
                .await
                .map_err(|e| {
                    anyhow::anyhow!("opening the Kinect failed — is it plugged in? ({e})")
                })?;
            device
                .start()
                .await
                .map_err(|e| anyhow::anyhow!("starting the Kinect: {e}"))?;

            // The depth camera: factory intrinsics at 512x424, pose
            // from the offline calibration.
            let ir = device.get_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(
                    kinect_one::DEPTH_WIDTH as u32,
                    kinect_one::DEPTH_HEIGHT as u32,
                ),
            };
            ready
                .send(Ok(camera))
                .map_err(|_| anyhow::anyhow!("the app went away"))?;

            let mut registration = kinect_one::processor::Registration::new();
            registration.set_ir_params(device.get_ir_params());
            registration.set_color_params(device.get_color_params());
            let color_processor = kinect_one::processor::color::ZuneColorProcessor::new(
                kinect_one::processor::color::ColorSpace::RGB,
            )
            .map_err(|e| anyhow::anyhow!("color processor: {e}"))?;
            use kinect_one::processor::depth::DepthProcessorTrait;
            let mut depth_processor = kinect_one::processor::depth::CpuDepthProcessor::new()
                .map_err(|e| anyhow::anyhow!("depth processor: {e}"))?;
            depth_processor
                .set_p0_tables(device.get_p0_tables())
                .map_err(|e| anyhow::anyhow!("depth p0 tables: {e}"))?;
            depth_processor
                .set_ir_params(device.get_ir_params())
                .map_err(|e| anyhow::anyhow!("depth ir params: {e}"))?;

            // 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 mut packet_sync = kinect_one::PacketSync::new();
            let started = Instant::now();
            while !shutdown.load(Ordering::Relaxed) {
                if let Ok(Some(packet)) = device.poll_color_packet().await {
                    packet_sync.push_color_packet(packet);
                }
                if let Ok(Some(packet)) = device.poll_depth_packet().await {
                    packet_sync.push_depth_packet(packet);
                }
                let Some((color_packet, depth_packet)) = packet_sync.poll_packets() else {
                    continue;
                };

                use kinect_one::processor::ProcessTrait;
                let color_frame = color_packet
                    .process(&color_processor)
                    .await
                    .map_err(|e| anyhow::anyhow!("color decode: {e}"))?;
                let depth_frame = depth_packet
                    .process(&depth_processor)
                    .await
                    .map_err(|e| anyhow::anyhow!("depth decode: {e}"))?
                    .1;
                let (registered, undistorted) =
                    registration.undistort_depth_and_color(&color_frame, &depth_frame, true);

                let depth_mm = convert_depth(&undistorted.buffer);
                if frames
                    .send(PairedFrame {
                        depth_mm,
                        rgb: registered.buffer,
                        left: latest(&left_rx),
                        right: latest(&right_rx),
                        timestamp_secs: started.elapsed().as_secs_f64(),
                    })
                    .is_err()
                {
                    break; // receiver dropped: shut down
                }
            }
            Ok(())
        })
    })();
    if let Err(e) = result {
        log::error!("kinect worker: {e:#}");
        let _ = ready.send(Err(e));
    }
}

/// libfreenect2-style depth (f32 metres, invalid as 0/NaN) → u16
/// millimetres with 0 for invalid.
fn convert_depth(buffer: &[f32]) -> Vec<u16> {
    buffer
        .iter()
        .map(|&d| {
            if d.is_finite() && (0.5..=6.0).contains(&d) {
                (d * 1000.0).round() as u16
            } else {
                0
            }
        })
        .collect()
}

/// 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::{
            CameraFormat, FrameFormat, RequestedFormat, RequestedFormatType, Resolution,
        };

        // Match the device: numeric id → index, otherwise the
        // human-readable name.
        let devices = nokhwa::query(nokhwa::utils::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 nokhwa::utils::CameraIndex;
    use rig::CameraPose;

    #[test]
    fn depth_conversion_drops_invalid_and_quantizes() {
        // libfreenect2 depth: f32 metres, invalid as 0/NaN; out-of-range
        // values are dropped too.
        let out = convert_depth(&[1.9, 0.0, f32::NAN, 0.25, 6.5, 2.845]);
        assert_eq!(out, vec![1900, 0, 0, 0, 0, 2845]);
    }

    #[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, "", "", 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(), CameraIndex::Index(1));
        // Name match (case-insensitive substring).
        assert_eq!(
            device_index(&devices, "c920").unwrap(),
            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),
            },
        };
        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:#}"
        );
    }
}