//! This module contains variable profile of ONNX model related.
//!
use std;
use std::ffi::CString;
use std::mem;

use libc::int32_t;

use dtype::Dtype;
use error::{cvt_r, Error};
use ffi;
use model_data::ModelData;

pub struct VariableProfile {
    pub dtype: Dtype,
    pub dims: Vec<i32>,
}

/// Variable Profile table.
///
/// An instance of this struct is generated by
/// `VariableProfileTableBuilder::build_variable_profile_table`
pub struct VariableProfileTable {
    handle: ffi::menoh_variable_profile_table_handle,
}

/// Builder of Variable Profile Table.
pub struct VariableProfileTableBuilder {
    handle: ffi::menoh_variable_profile_table_builder_handle,
}

type Result<T> = std::result::Result<T, Error>;

impl VariableProfileTable {
    /// Get Variable profile detail by using variable name.
    pub fn get_variable_profile(&self, name: &str) -> Result<VariableProfile> {
        let dtype = variable_profile_table_get_dtype(self.handle, name)?;
        let dims_size = variable_profile_table_get_dims_size(self.handle, name)?;
        let mut dims = Vec::new();
        for index in 0..dims_size {
            dims.push(variable_profile_table_get_dims_at(
                self.handle,
                name,
                index,
            )?);
        }

        Ok(VariableProfile { dtype, dims })
    }

    #[doc(hidden)]
    pub unsafe fn get_handle(&self) -> ffi::menoh_variable_profile_table_handle {
        self.handle
    }
}

impl VariableProfileTableBuilder {
    pub fn new() -> Result<Self> {
        let mut handle: ffi::menoh_variable_profile_table_builder_handle =
            unsafe { mem::uninitialized() };
        cvt_r(|| unsafe {
            ffi::menoh_make_variable_profile_table_builder(
                &mut handle as *mut ffi::menoh_variable_profile_table_builder_handle,
            )
        })?;
        Ok(VariableProfileTableBuilder { handle })
    }

    /// Add input profile.
    /// dims length must be 2 or 4.
    pub fn add_input_profile(&mut self, name: &str, dtype: Dtype, dims: &[i32]) -> Result<()> {
        let name = CString::new(name).map_err(|_| Error::VariableNotFound)?;
        match dims.len() {
            2 => {
                cvt_r(|| unsafe {
                    ffi::menoh_variable_profile_table_builder_add_input_profile_dims_2(
                        self.handle,
                        name.as_ptr(),
                        dtype.value(),
                        dims[0],
                        dims[1],
                    )
                })?;
                return Ok(());
            }
            4 => {
                cvt_r(|| unsafe {
                    ffi::menoh_variable_profile_table_builder_add_input_profile_dims_4(
                        self.handle,
                        name.as_ptr(),
                        dtype.value(),
                        dims[0],
                        dims[1],
                        dims[2],
                        dims[3],
                    )
                })?;
                return Ok(());
            }
            _ => Err(Error::DimensionMismatch),
        }
    }

    /// Add output profile.
    pub fn add_output_profile(&mut self, name: &str, dtype: Dtype) -> Result<()> {
        let name = CString::new(name).map_err(|_| Error::VariableNotFound)?;
        cvt_r(|| unsafe {
            ffi::menoh_variable_profile_table_builder_add_output_profile(
                self.handle,
                name.as_ptr(),
                dtype.value(),
            )
        })?;
        Ok(())
    }

    /// Build variable profile table.
    pub fn build_variable_profile_table(
        &self,
        model_data: &ModelData,
    ) -> Result<VariableProfileTable> {
        let mut handle: ffi::menoh_variable_profile_table_handle = unsafe { mem::uninitialized() };
        cvt_r(|| unsafe {
            ffi::menoh_build_variable_profile_table(
                self.handle,
                model_data.get_handle(),
                &mut handle as *mut ffi::menoh_variable_profile_table_handle,
            )
        })?;
        Ok(VariableProfileTable { handle })
    }
}

impl Drop for VariableProfileTable {
    fn drop(&mut self) {
        unsafe { ffi::menoh_delete_variable_profile_table(self.handle) }
    }
}

impl Drop for VariableProfileTableBuilder {
    fn drop(&mut self) {
        unsafe {
            ffi::menoh_delete_variable_profile_table_builder(self.handle);
        }
    }
}

fn variable_profile_table_get_dtype(
    handle: ffi::menoh_variable_profile_table_handle,
    name: &str,
) -> Result<Dtype> {
    let name = CString::new(name).map_err(|_| Error::VariableNotFound)?;
    let mut dtype: ffi::menoh_dtype = ffi::menoh_dtype::default();
    cvt_r(|| unsafe {
        ffi::menoh_variable_profile_table_get_dtype(
            handle,
            name.as_ptr(),
            &mut dtype as *mut ffi::menoh_dtype,
        )
    })?;
    let dtype = Dtype::from(dtype);
    Ok(dtype)
}

fn variable_profile_table_get_dims_size(
    handle: ffi::menoh_variable_profile_table_handle,
    name: &str,
) -> Result<i32> {
    let name = CString::new(name).map_err(|_| Error::VariableNotFound)?;
    let mut dims_size: int32_t = int32_t::default();
    cvt_r(|| unsafe {
        ffi::menoh_variable_profile_table_get_dims_size(
            handle,
            name.as_ptr(),
            &mut dims_size as *mut int32_t,
        )
    })?;
    Ok(dims_size)
}

fn variable_profile_table_get_dims_at(
    handle: ffi::menoh_variable_profile_table_handle,
    name: &str,
    index: i32,
) -> Result<i32> {
    let name = CString::new(name).map_err(|_| Error::VariableNotFound)?;
    let mut dst_dim: int32_t = int32_t::default();
    cvt_r(|| unsafe {
        ffi::menoh_variable_profile_table_get_dims_at(
            handle,
            name.as_ptr(),
            index,
            &mut dst_dim as *mut int32_t,
        )
    })?;
    Ok(dst_dim)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn vpt_builder_new() {
        assert!(VariableProfileTableBuilder::new().is_ok());
    }
}