use std::path::{Path, PathBuf};

use crate::error::Result;
use crate::schema::{InferenceRequest, InferenceResponse};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelFormat {
    PyTorch,
    Onnx,
}

impl ModelFormat {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::PyTorch => "pytorch",
            Self::Onnx => "onnx",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelArtifact {
    pub name: String,
    pub format: ModelFormat,
    pub path: PathBuf,
}

impl ModelArtifact {
    pub fn pytorch(name: impl Into<String>, path: impl Into<PathBuf>) -> Self {
        Self {
            name: name.into(),
            format: ModelFormat::PyTorch,
            path: path.into(),
        }
    }

    pub fn onnx(name: impl Into<String>, path: impl Into<PathBuf>) -> Self {
        Self {
            name: name.into(),
            format: ModelFormat::Onnx,
            path: path.into(),
        }
    }

    pub fn path(&self) -> &Path {
        &self.path
    }
}

pub trait UserServingAdapter<WireRequest, WireResponse>: Send + Sync + 'static {
    fn model(&self) -> &ModelArtifact;

    fn preprocess(&self, request: WireRequest) -> Result<InferenceRequest>;

    fn postprocess(&self, response: InferenceResponse) -> Result<WireResponse>;
}
