use std::collections::HashMap;

use crate::error::{NarsilError, Result};

pub mod proto {
    tonic::include_proto!("narsil.inference.v1");
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DataType {
    F32,
    F64,
    I64,
    I32,
    U8,
    Bool,
}

impl DataType {
    pub fn byte_width(self) -> usize {
        match self {
            Self::F32 | Self::I32 => 4,
            Self::F64 | Self::I64 => 8,
            Self::U8 | Self::Bool => 1,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Tensor {
    pub name: String,
    pub dtype: DataType,
    pub shape: Vec<usize>,
    pub data: Vec<u8>,
}

impl Tensor {
    pub fn new(
        name: impl Into<String>,
        dtype: DataType,
        shape: Vec<usize>,
        data: Vec<u8>,
    ) -> Result<Self> {
        let tensor = Self {
            name: name.into(),
            dtype,
            shape,
            data,
        };
        tensor.validate()?;
        Ok(tensor)
    }

    pub fn from_f32(name: impl Into<String>, shape: Vec<usize>, values: Vec<f32>) -> Result<Self> {
        let mut data = Vec::with_capacity(values.len() * 4);
        for value in values {
            data.extend_from_slice(&value.to_le_bytes());
        }
        Self::new(name, DataType::F32, shape, data)
    }

    pub fn f32_values(&self) -> Result<Vec<f32>> {
        if self.dtype != DataType::F32 {
            return Err(NarsilError::Schema(format!(
                "tensor '{}' must be F32, got {:?}",
                self.name, self.dtype
            )));
        }
        if !self.data.len().is_multiple_of(4) {
            return Err(NarsilError::Schema(format!(
                "tensor '{}' has invalid f32 byte length {}",
                self.name,
                self.data.len()
            )));
        }

        Ok(self
            .data
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("chunk has length 4")))
            .collect())
    }

    pub fn validate(&self) -> Result<()> {
        if self.shape.is_empty() {
            return Err(NarsilError::Schema(format!(
                "tensor '{}' must have at least one dimension",
                self.name
            )));
        }

        let elements = self.shape.iter().try_fold(1usize, |acc, dim| {
            if *dim == 0 {
                return Err(NarsilError::Schema(format!(
                    "tensor '{}' has zero-sized dimension",
                    self.name
                )));
            }
            acc.checked_mul(*dim).ok_or_else(|| {
                NarsilError::Schema(format!("tensor '{}' shape overflows usize", self.name))
            })
        })?;
        let expected_bytes = elements
            .checked_mul(self.dtype.byte_width())
            .ok_or_else(|| {
                NarsilError::Schema(format!("tensor '{}' byte size overflows usize", self.name))
            })?;

        if self.data.len() != expected_bytes {
            return Err(NarsilError::Schema(format!(
                "tensor '{}' has {} bytes, expected {} for {:?} {:?}",
                self.name,
                self.data.len(),
                expected_bytes,
                self.dtype,
                self.shape
            )));
        }

        Ok(())
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct InferenceRequest {
    pub model_name: String,
    pub inputs: Vec<Tensor>,
    pub parameters: HashMap<String, String>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct InferenceResponse {
    pub outputs: Vec<Tensor>,
    pub metadata: HashMap<String, String>,
}

impl TryFrom<proto::Tensor> for Tensor {
    type Error = NarsilError;

    fn try_from(value: proto::Tensor) -> Result<Self> {
        let dtype = match proto::DataType::try_from(value.dtype) {
            Ok(proto::DataType::F32) => DataType::F32,
            Ok(proto::DataType::F64) => DataType::F64,
            Ok(proto::DataType::I64) => DataType::I64,
            Ok(proto::DataType::I32) => DataType::I32,
            Ok(proto::DataType::U8) => DataType::U8,
            Ok(proto::DataType::Bool) => DataType::Bool,
            Ok(proto::DataType::Unspecified) | Err(_) => {
                return Err(NarsilError::Schema(format!(
                    "tensor '{}' has unspecified dtype",
                    value.name
                )));
            }
        };

        let shape = value
            .shape
            .into_iter()
            .map(|dim| {
                usize::try_from(dim).map_err(|_| {
                    NarsilError::Schema(format!(
                        "tensor '{}' dimension {dim} exceeds usize",
                        value.name
                    ))
                })
            })
            .collect::<Result<Vec<_>>>()?;

        Self::new(value.name, dtype, shape, value.data)
    }
}

impl From<Tensor> for proto::Tensor {
    fn from(value: Tensor) -> Self {
        Self {
            name: value.name,
            dtype: proto_dtype(value.dtype) as i32,
            shape: value.shape.into_iter().map(|dim| dim as u64).collect(),
            data: value.data,
        }
    }
}

impl TryFrom<proto::PredictRequest> for InferenceRequest {
    type Error = NarsilError;

    fn try_from(value: proto::PredictRequest) -> Result<Self> {
        if value.inputs.is_empty() {
            return Err(NarsilError::InvalidRequest(
                "predict request must contain at least one input tensor".to_string(),
            ));
        }

        Ok(Self {
            model_name: value.model_name,
            inputs: value
                .inputs
                .into_iter()
                .map(Tensor::try_from)
                .collect::<Result<Vec<_>>>()?,
            parameters: value.parameters,
        })
    }
}

impl From<InferenceResponse> for proto::PredictResponse {
    fn from(value: InferenceResponse) -> Self {
        Self {
            outputs: value.outputs.into_iter().map(proto::Tensor::from).collect(),
            metadata: value.metadata,
        }
    }
}

fn proto_dtype(dtype: DataType) -> proto::DataType {
    match dtype {
        DataType::F32 => proto::DataType::F32,
        DataType::F64 => proto::DataType::F64,
        DataType::I64 => proto::DataType::I64,
        DataType::I32 => proto::DataType::I32,
        DataType::U8 => proto::DataType::U8,
        DataType::Bool => proto::DataType::Bool,
    }
}
