#![allow(dead_code)]
#![allow(non_camel_case_types)]
#![allow(non_upper_case_globals)]
#![allow(non_snake_case)]
#![allow(improper_ctypes)]
#![allow(deref_nullptr)]

//! This library brings Google's Mediapipe to Rust.

// LINKING
#[link(name = "mediagraph")]
extern "C" {}

use opencv::prelude::*;
use std::ffi::CString;

mod bindings;

use bindings::*;

/// The C++ mediagraph graph type.
pub type DetectorType = mediagraph_DetectorType;

/// The C++ mediagraph landmark type.
pub type Landmark = mediagraph_Landmark;

pub const FACE_GRAPH_TYPE: DetectorType = mediagraph_DetectorType_FACE;
pub const HANDS_GRAPH_TYPE: DetectorType = mediagraph_DetectorType_HANDS;
pub const POSE_GRAPH_TYPE: DetectorType = mediagraph_DetectorType_POSE;

impl Default for Landmark {
    fn default() -> Self {
        Self {
            x: 0.0,
            y: 0.0,
            z: 0.0,
            visibility: 0.0,
            presence: 0.0,
        }
    }
}

/// Represents a detected pose, as 33 landmarks.
/// Landmark names are in [pose::PoseLandmark].
pub struct Pose {
    pub data: [Landmark; 33],
}

impl Default for Pose {
    fn default() -> Self {
        Self {
            data: [Landmark::default(); 33],
        }
    }
}

/// Represents a detected hand, as 21 landmarks.
/// Landmark names are in [hands::HandLandmark]
#[derive(Default)]
pub struct Hand {
    pub data: [Landmark; 21],
}

/// Represents a detected face mesh, as 478 landmarks.
pub struct FaceMesh {
    pub data: [Landmark; 478],
}

impl Default for FaceMesh {
    fn default() -> Self {
        Self {
            data: [Landmark::default(); 478],
        }
    }
}

/// Detector calculator which interacts with the C++ library.
pub struct Detector {
    graph: *mut mediagraph_Detector,
    num_landmarks: u32,
}

impl Detector {
    /// Creates a new Mediagraph with the given config.
    pub fn new(graph_type: DetectorType, graph_config: &str, output_node: &str) -> Self {
        let graph_config = CString::new(graph_config).expect("CString::new failed");
        let output_node = CString::new(output_node).expect("CString::new failed");

        let graph: *mut mediagraph_Detector = unsafe {
            mediagraph_Detector::Create(graph_type, graph_config.as_ptr(), output_node.as_ptr())
        };

        let num_landmarks = match graph_type {
            FACE_GRAPH_TYPE => 478,
            HANDS_GRAPH_TYPE => 42,
            POSE_GRAPH_TYPE => 33,
            _ => 0,
        };

        Self {
            graph,
            num_landmarks,
        }
    }

    /// Processes the input frame, returns a slice of landmarks if any are detected.
    pub fn process(&mut self, input: &Mat) -> &[Landmark] {
        let mut data = input.clone();
        let raw_landmarks = unsafe {
            mediagraph_Detector_Process(
                self.graph as *mut std::ffi::c_void,
                data.data_mut(),
                data.cols(),
                data.rows(),
            )
        };
        if raw_landmarks.is_null() {
            return &[];
        }
        let landmarks =
            unsafe { std::slice::from_raw_parts(raw_landmarks, self.num_landmarks as usize) };
        landmarks
    }
}

impl Drop for Detector {
    fn drop(&mut self) {
        unsafe {
            mediagraph_Detector_Detector_destructor(self.graph);
        }
    }
}

pub mod pose {
    //! Pose detection utilities.
    use super::*;

    /// Pose landmark indices.
    pub enum PoseLandmark {
        NOSE = 0,
        LEFT_EYE_INNER = 1,
        LEFT_EYE = 2,
        LEFT_EYE_OUTER = 3,
        RIGHT_EYE_INNER = 4,
        RIGHT_EYE = 5,
        RIGHT_EYE_OUTER = 6,
        LEFT_EAR = 7,
        RIGHT_EAR = 8,
        MOUTH_LEFT = 9,
        MOUTH_RIGHT = 10,
        LEFT_SHOULDER = 11,
        RIGHT_SHOULDER = 12,
        LEFT_ELBOW = 13,
        RIGHT_ELBOW = 14,
        LEFT_WRIST = 15,
        RIGHT_WRIST = 16,
        LEFT_PINKY = 17,
        RIGHT_PINKY = 18,
        LEFT_INDEX = 19,
        RIGHT_INDEX = 20,
        LEFT_THUMB = 21,
        RIGHT_THUMB = 22,
        LEFT_HIP = 23,
        RIGHT_HIP = 24,
        LEFT_KNEE = 25,
        RIGHT_KNEE = 26,
        LEFT_ANKLE = 27,
        RIGHT_ANKLE = 28,
        LEFT_HEEL = 29,
        RIGHT_HEEL = 30,
        LEFT_FOOT_INDEX = 31,
        RIGHT_FOOT_INDEX = 32,
    }

    pub struct PoseDetector {
        pub mode: bool,         // false
        pub smooth: bool,       // true,
        pub detection_con: f32, // 0.5
        pub track_con: f32,     // 0.5
        graph: Detector,
    }

    impl PoseDetector {
        pub fn new(mode: bool, smooth: bool, detection_con: f32, track_con: f32) -> Self {
            let graph = Detector::new(
                POSE_GRAPH_TYPE,
                include_str!("pose_tracking_cpu.txt"),
                "pose_landmarks",
            );

            Self {
                mode,
                smooth,
                detection_con,
                track_con,
                graph,
            }
        }

        /// Processes the input frame, returns a pose if detected.
        pub fn process(&mut self, input: &Mat) -> Option<Pose> {
            let landmarks = self.graph.process(input);

            if landmarks.is_empty() {
                return None;
            }

            let mut pose = Pose::default();
            pose.data.copy_from_slice(landmarks);
            Some(pose)
        }
    }

    impl Default for PoseDetector {
        fn default() -> Self {
            Self::new(false, true, 0.5, 0.5)
        }
    }
}

pub mod face_mesh {
    //! Face detection utilities.
    use super::*;

    pub struct FaceMeshDetector {
        pub static_mode: bool,      // false
        pub max_faces: usize,       // 2
        pub min_detection_con: f32, // 0.5
        pub min_track_con: f32,     // 0.5
        graph: Detector,
    }

    impl FaceMeshDetector {
        pub fn new(
            static_mode: bool,
            max_faces: usize,
            min_detection_con: f32,
            min_track_con: f32,
        ) -> Self {
            let graph = Detector::new(
                FACE_GRAPH_TYPE,
                include_str!("face_mesh_desktop_live.txt"),
                "multi_face_landmarks",
            );

            Self {
                static_mode,
                max_faces,
                min_detection_con,
                min_track_con,
                graph,
            }
        }

        /// Processes the input frame, returns a face mesh if detected.
        pub fn process(&mut self, input: &Mat) -> Option<FaceMesh> {
            let landmarks = self.graph.process(input);

            if landmarks.is_empty() {
                return None;
            }

            let mut face_mesh = FaceMesh::default();
            face_mesh.data.copy_from_slice(landmarks);
            Some(face_mesh)
        }
    }

    impl Default for FaceMeshDetector {
        fn default() -> Self {
            Self::new(false, 2, 0.5, 0.5)
        }
    }
}

pub mod hands {
    //! Hand detection utilities.
    use super::*;

    /// Hand landmark indices.
    pub enum HandLandmark {
        WRIST = 0,
        THUMB_CMC = 1,
        THUMB_MCP = 2,
        THUMB_IP = 3,
        THUMB_TIP = 4,
        INDEX_FINGER_MCP = 5,
        INDEX_FINGER_PIP = 6,
        INDEX_FINGER_DIP = 7,
        INDEX_FINGER_TIP = 8,
        MIDDLE_FINGER_MCP = 9,
        MIDDLE_FINGER_PIP = 10,
        MIDDLE_FINGER_DIP = 11,
        MIDDLE_FINGER_TIP = 12,
        RING_FINGER_MCP = 13,
        RING_FINGER_PIP = 14,
        RING_FINGER_DIP = 15,
        RING_FINGER_TIP = 16,
        PINKY_MCP = 17,
        PINKY_PIP = 18,
        PINKY_DIP = 19,
        PINKY_TIP = 20,
    }

    pub struct HandDetector {
        pub mode: bool,
        pub max_hands: usize,
        pub detection_con: f32, // 0.5
        pub min_track_con: f32, // 0.5
        graph: Detector,
    }

    impl HandDetector {
        pub fn new(mode: bool, max_hands: usize, detection_con: f32, min_track_con: f32) -> Self {
            let graph = Detector::new(
                HANDS_GRAPH_TYPE,
                include_str!("hand_tracking_desktop_live.txt"),
                "hand_landmarks",
            );

            Self {
                mode,
                max_hands,
                detection_con,
                min_track_con,
                graph,
            }
        }

        /// Processes the input frame, returns a tuple of hands if detected.
        pub fn process(&mut self, input: &Mat) -> Option<[Hand; 2]> {
            let landmarks = self.graph.process(input);

            if landmarks.is_empty() {
                return None;
            }

            let mut lh = Hand::default();
            let mut rh = Hand::default();
            lh.data.copy_from_slice(&landmarks[0..21]);
            rh.data.copy_from_slice(&landmarks[21..42]);

            Some([lh, rh])
        }
    }

    impl Default for HandDetector {
        fn default() -> Self {
            Self::new(false, 2, 0.5, 0.5)
        }
    }
}