furnace-server 0.8.0

use burn::tensor::{Device, Tensor, TensorData};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::Instant;
use tracing::{debug, info};

use crate::backend::BackendKind;
use crate::burn_model::{BurnModelContainer, create_sample_model};
use crate::error::{ModelError, Result};
use crate::models::{BuiltInModel, resolve_weights};
use crate::types::TensorSpec;

/// Hard upper bound on the number of items in one inference batch
pub const MAX_BATCH_SIZE: usize = 1000;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    pub name: String,
    pub version: String,
    pub input_spec: TensorSpec,
    pub output_spec: TensorSpec,
    pub model_type: String,
    pub backend: String,
    pub created_at: String,
    pub model_size_bytes: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelStats {
    pub inference_count: u64,
    pub total_inference_time_ms: f64,
    pub last_inference_time: Option<DateTime<Utc>>,
    pub average_inference_time_ms: f64,
    pub error_count: u64,
    pub success_count: u64,
    pub min_inference_time_ms: Option<f64>,
    pub max_inference_time_ms: Option<f64>,
    /// Forward passes run. Below `success_count` when requests were batched together
    pub batch_count: u64,
    /// Items processed by those forward passes
    pub batched_items: u64,
}

impl ModelStats {
    fn new() -> Self {
        Self {
            inference_count: 0,
            total_inference_time_ms: 0.0,
            last_inference_time: None,
            average_inference_time_ms: 0.0,
            error_count: 0,
            success_count: 0,
            min_inference_time_ms: None,
            max_inference_time_ms: None,
            batch_count: 0,
            batched_items: 0,
        }
    }

    fn record_success(&mut self, elapsed_ms: f64) {
        self.inference_count += 1;
        self.success_count += 1;
        self.total_inference_time_ms += elapsed_ms;
        self.average_inference_time_ms = self.total_inference_time_ms / self.success_count as f64;
        self.min_inference_time_ms = Some(
            self.min_inference_time_ms
                .map_or(elapsed_ms, |m| m.min(elapsed_ms)),
        );
        self.max_inference_time_ms = Some(
            self.max_inference_time_ms
                .map_or(elapsed_ms, |m| m.max(elapsed_ms)),
        );
        self.last_inference_time = Some(Utc::now());
    }

    fn record_error(&mut self) {
        self.inference_count += 1;
        self.error_count += 1;
    }

    fn record_batch(&mut self, items: usize) {
        self.batch_count += 1;
        self.batched_items += items as u64;
    }
}

/// One request's input for [`Model::predict_many`]: `batch_size` items back to back
#[derive(Debug)]
pub struct FlatInput {
    pub data: Vec<f32>,
    pub batch_size: usize,
}

/// A model that can run inference on flattened `[batch, features]` input.
pub trait BurnModel: Send + Sync + std::fmt::Debug {
    fn predict(&self, input: Tensor<2>) -> Result<Tensor<2>>;
    /// Device the model's weights live on; inputs are created here.
    fn device(&self) -> &Device;
    /// Per-item input shape (without the batch dimension)
    fn get_input_shape(&self) -> &[usize];
    /// Per-item output shape (without the batch dimension)
    fn get_output_shape(&self) -> &[usize];
    fn get_name(&self) -> &str;
    fn backend(&self) -> BackendKind;
}

/// File-based MLP model (`.bpk` weights + `.json` metadata)
#[derive(Debug)]
pub struct MlpModel {
    container: BurnModelContainer,
    input_shape: Vec<usize>,
    output_shape: Vec<usize>,
}

impl MlpModel {
    pub fn new(container: BurnModelContainer) -> Self {
        let input_shape = container.input_shape();
        let output_shape = container.output_shape();
        Self {
            container,
            input_shape,
            output_shape,
        }
    }
}

impl BurnModel for MlpModel {
    fn predict(&self, input: Tensor<2>) -> Result<Tensor<2>> {
        Ok(self.container.predict(input))
    }

    fn device(&self) -> &Device {
        self.container.device()
    }

    fn get_input_shape(&self) -> &[usize] {
        &self.input_shape
    }

    fn get_output_shape(&self) -> &[usize] {
        &self.output_shape
    }

    fn get_name(&self) -> &str {
        &self.container.metadata.name
    }

    fn backend(&self) -> BackendKind {
        self.container.backend
    }
}

/// Constant-output model for tests and benchmarking server overhead.
///
/// Never used as a silent fallback: if a real model fails to load, loading fails.
#[derive(Debug)]
pub struct DummyModel {
    name: String,
    input_shape: Vec<usize>,
    output_shape: Vec<usize>,
    device: Device,
}

impl DummyModel {
    pub fn new(name: String, input_shape: Vec<usize>, output_shape: Vec<usize>) -> Self {
        Self {
            name,
            input_shape,
            output_shape,
            device: BackendKind::Cpu.device(),
        }
    }
}

impl BurnModel for DummyModel {
    fn predict(&self, input: Tensor<2>) -> Result<Tensor<2>> {
        let [batch_size, _] = input.dims();
        let output_size = self.output_shape.iter().product::<usize>();
        Ok(Tensor::from_data(
            TensorData::new(
                vec![0.5f32; batch_size * output_size],
                [batch_size, output_size],
            ),
            &self.device,
        ))
    }

    fn device(&self) -> &Device {
        &self.device
    }

    fn get_input_shape(&self) -> &[usize] {
        &self.input_shape
    }

    fn get_output_shape(&self) -> &[usize] {
        &self.output_shape
    }

    fn get_name(&self) -> &str {
        &self.name
    }

    fn backend(&self) -> BackendKind {
        BackendKind::Cpu
    }
}

/// A loaded model plus its metadata and runtime statistics
#[derive(Debug)]
pub struct Model {
    inner: Box<dyn BurnModel>,
    info: ModelInfo,
    model_type: &'static str,
    stats: Mutex<ModelStats>,
}

impl Model {
    pub fn new(inner: Box<dyn BurnModel>, model_type: &'static str, model_size_bytes: u64) -> Self {
        let info = ModelInfo {
            name: inner.get_name().to_string(),
            version: "1.0.0".to_string(),
            input_spec: TensorSpec {
                shape: inner.get_input_shape().to_vec(),
                dtype: "float32".to_string(),
                min_value: None,
                max_value: None,
            },
            output_spec: TensorSpec {
                shape: inner.get_output_shape().to_vec(),
                dtype: "float32".to_string(),
                min_value: None,
                max_value: None,
            },
            model_type: model_type.to_string(),
            backend: inner.backend().to_string(),
            created_at: Utc::now().to_rfc3339(),
            model_size_bytes,
        };

        Self {
            inner,
            info,
            model_type,
            stats: Mutex::new(ModelStats::new()),
        }
    }

    pub fn predict(&self, input: Vec<f32>) -> Result<Vec<f32>> {
        self.predict_batch(vec![input])
            .map(|mut batch| batch.pop().expect("batch of one yields one output"))
    }

    /// Run inference on a batch of flattened inputs.
    ///
    /// This is CPU/GPU-bound and blocking; call it from a blocking context
    /// (e.g. `tokio::task::spawn_blocking`) inside async code.
    pub fn predict_batch(&self, inputs: Vec<Vec<f32>>) -> Result<Vec<Vec<f32>>> {
        let batch_size = inputs.len();
        let flattened = match self.flatten_batch(inputs) {
            Ok(flattened) => flattened,
            Err(e) => {
                self.stats
                    .lock()
                    .unwrap_or_else(|e| e.into_inner())
                    .record_error();
                return Err(e);
            }
        };

        let output_size: usize = self.info.output_spec.shape.iter().product();
        Ok(self
            .predict_flat(flattened, batch_size)?
            .chunks_exact(output_size)
            .map(<[f32]>::to_vec)
            .collect())
    }

    /// Run inference on `batch_size` items stored back to back in one row-major buffer.
    ///
    /// Returns the outputs the same way: `batch_size * output_size` values. This avoids
    /// per-item allocations and copies, so it is the fastest entry point for large inputs.
    /// Blocking, like [`Model::predict_batch`].
    pub fn predict_flat(&self, input: Vec<f32>, batch_size: usize) -> Result<Vec<f32>> {
        self.predict_many(vec![FlatInput {
            data: input,
            batch_size,
        }])
        .pop()
        .expect("one request yields one result")
    }

    /// Run several independent requests as one forward pass and split the outputs back.
    ///
    /// Each request is validated on its own, so an invalid one gets its own error and
    /// the rest still run. Results come back in request order, each holding that
    /// request's `batch_size * output_size` values. The requests together must not
    /// exceed [`MAX_BATCH_SIZE`] items. Statistics count every request, plus one batch
    /// per forward pass. Blocking, like [`Model::predict_batch`].
    pub fn predict_many(&self, requests: Vec<FlatInput>) -> Vec<Result<Vec<f32>>> {
        let mut results: Vec<Option<Result<Vec<f32>>>> = Vec::with_capacity(requests.len());
        let mut valid = Vec::with_capacity(requests.len());
        let mut errors = 0;
        for (i, request) in requests.into_iter().enumerate() {
            match self.validate_flat(&request.data, request.batch_size) {
                Ok(()) => {
                    results.push(None);
                    valid.push((i, request));
                }
                Err(e) => {
                    results.push(Some(Err(e)));
                    errors += 1;
                }
            }
        }

        let start = Instant::now();
        let items: usize = valid.iter().map(|(_, r)| r.batch_size).sum();
        let output = match valid.as_mut_slice() {
            [] => None,
            // A lone request runs on its own buffer, without a copy
            [(_, only)] => Some(self.forward(std::mem::take(&mut only.data), items)),
            many => {
                let mut input = Vec::with_capacity(many.iter().map(|(_, r)| r.data.len()).sum());
                for (_, request) in many.iter_mut() {
                    input.append(&mut request.data);
                }
                Some(self.forward(input, items))
            }
        };
        let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0;

        let mut stats = self.stats.lock().unwrap_or_else(|e| e.into_inner());
        for _ in 0..errors {
            stats.record_error();
        }
        if let Some(output) = output {
            stats.record_batch(items);
            match output {
                Ok(output) => {
                    debug!(
                        "Batch inference completed in {:.2}ms, {} request(s), batch size: {}, backend: {}",
                        elapsed_ms,
                        valid.len(),
                        items,
                        self.info.backend
                    );
                    if let [(i, _)] = valid.as_slice() {
                        // A lone request takes the whole buffer
                        stats.record_success(elapsed_ms);
                        results[*i] = Some(Ok(output));
                    } else {
                        let output_size: usize = self.info.output_spec.shape.iter().product();
                        let mut rest = output.as_slice();
                        for (i, request) in &valid {
                            let (mine, tail) = rest.split_at(request.batch_size * output_size);
                            rest = tail;
                            stats.record_success(elapsed_ms);
                            results[*i] = Some(Ok(mine.to_vec()));
                        }
                    }
                }
                Err(message) => {
                    for (i, _) in &valid {
                        stats.record_error();
                        results[*i] =
                            Some(Err(ModelError::InferenceFailed(message.clone()).into()));
                    }
                }
            }
        }

        results
            .into_iter()
            .map(|r| r.expect("every request has a result"))
            .collect()
    }

    fn check_batch_size(batch_size: usize) -> Result<()> {
        if batch_size == 0 {
            return Err(ModelError::InputValidation {
                expected: vec![1],
                actual: vec![0],
            }
            .into());
        }
        if batch_size > MAX_BATCH_SIZE {
            return Err(ModelError::InputValidation {
                expected: vec![MAX_BATCH_SIZE],
                actual: vec![batch_size],
            }
            .into());
        }
        Ok(())
    }

    fn flatten_batch(&self, inputs: Vec<Vec<f32>>) -> Result<Vec<f32>> {
        Self::check_batch_size(inputs.len())?;
        let input_size: usize = self.info.input_spec.shape.iter().product();
        let mut flattened = Vec::with_capacity(inputs.len() * input_size);
        for input in inputs {
            self.validate_input_shape(&input)?;
            flattened.extend_from_slice(&input);
        }
        Ok(flattened)
    }

    /// Checks one request's flat input: batch size, length and finite values
    fn validate_flat(&self, input: &[f32], batch_size: usize) -> Result<()> {
        Self::check_batch_size(batch_size)?;

        let input_size: usize = self.info.input_spec.shape.iter().product();
        if input.len() != batch_size * input_size {
            let mut expected = vec![batch_size];
            expected.extend_from_slice(&self.info.input_spec.shape);
            return Err(ModelError::InputValidation {
                expected,
                actual: vec![input.len()],
            }
            .into());
        }
        if let Some(pos) = input.iter().position(|x| !x.is_finite()) {
            return Err(ModelError::InvalidDataType(format!(
                "Invalid value at batch[{}][{}]: {} (not finite)",
                pos / input_size,
                pos % input_size,
                input[pos]
            ))
            .into());
        }
        Ok(())
    }

    /// One forward pass over validated input; the error is the failure message
    /// Continue `previous`'s statistics, so a reloaded model's counters do not restart at zero
    pub fn adopt_stats(&self, previous: &Model) {
        let carried = previous.get_stats();
        *self.stats.lock().unwrap_or_else(|e| e.into_inner()) = carried;
    }

    /// Run one all-zero input through the model without touching the statistics, so the first
    /// real request does not pay for lazy initialization (weight transforms, buffer pools,
    /// page faults on the memory-mapped weights). Returns how long it took.
    pub fn warm_up(&self) -> std::result::Result<std::time::Duration, String> {
        let input_size: usize = self.info.input_spec.shape.iter().product();
        let start = Instant::now();
        self.forward(vec![0.0; input_size], 1)?;
        Ok(start.elapsed())
    }

    fn forward(&self, input: Vec<f32>, batch_size: usize) -> std::result::Result<Vec<f32>, String> {
        let input_size: usize = self.info.input_spec.shape.iter().product();
        let input_tensor = Tensor::from_data(
            TensorData::new(input, [batch_size, input_size]),
            self.inner.device(),
        );

        let output = self
            .inner
            .predict(input_tensor)
            .map_err(|e| e.to_string())?;

        let output_size: usize = self.info.output_spec.shape.iter().product();
        let [out_batch, out_size] = output.dims();
        if out_batch != batch_size || out_size != output_size {
            return Err(format!(
                "model returned shape [{out_batch}, {out_size}], expected [{batch_size}, {output_size}]"
            ));
        }

        output
            .into_data()
            .try_into_vec::<f32>()
            .map_err(|e| format!("{e:?}"))
    }

    pub fn get_info(&self) -> &ModelInfo {
        &self.info
    }

    pub fn get_stats(&self) -> ModelStats {
        self.stats.lock().unwrap_or_else(|e| e.into_inner()).clone()
    }

    pub fn validate_input_shape(&self, input: &[f32]) -> Result<()> {
        let expected_size: usize = self.info.input_spec.shape.iter().product();
        if input.len() != expected_size {
            return Err(ModelError::InputValidation {
                expected: self.info.input_spec.shape.clone(),
                actual: vec![input.len()],
            }
            .into());
        }
        Ok(())
    }

    pub fn get_backend_info(&self) -> String {
        self.info.backend.clone()
    }

    pub fn model_type(&self) -> &'static str {
        self.model_type
    }
}

/// Load a file-based MLP model from `<path>.bpk` + `<path>.json` on the CPU backend
pub fn load_model(path: &Path) -> Result<Model> {
    load_model_with_backend(path, BackendKind::Cpu)
}

/// Load a file-based MLP model from `<path>.bpk` + `<path>.json`
pub fn load_model_with_backend(path: &Path, backend: BackendKind) -> Result<Model> {
    info!("Loading model from: {:?}", path);

    // Accept `model`, `model.bpk` or `model.json`
    let base: PathBuf = match path.extension().and_then(|e| e.to_str()) {
        Some("bpk") | Some("json") => path.with_extension(""),
        _ => path.to_path_buf(),
    };
    let weights_path = base.with_extension("bpk");

    // A `sample` path without files on disk yields a freshly initialized sample model
    if !weights_path.exists() {
        let is_sample = base
            .file_name()
            .and_then(|s| s.to_str())
            .is_some_and(|s| s.contains("sample"));
        if is_sample {
            info!("Creating sample MLP model");
            let container = create_sample_model()?;
            return Ok(Model::new(Box::new(MlpModel::new(container)), "mlp", 0));
        }
        return Err(ModelError::FileNotFound(weights_path).into());
    }

    let model_size_bytes = std::fs::metadata(&weights_path)
        .map_err(|e| ModelError::LoadFailed {
            path: weights_path.clone(),
            source: Box::new(e),
        })?
        .len();

    let container = BurnModelContainer::load(&base, backend)?;
    Ok(Model::new(
        Box::new(MlpModel::new(container)),
        "mlp",
        model_size_bytes,
    ))
}

/// Load a built-in (compiled-in) model by name on the CPU backend, with default weight lookup
pub fn load_built_in_model(model_name: &str) -> Result<Model> {
    load_built_in_model_with_backend(model_name, BackendKind::Cpu, None)
}

/// Load a built-in model by name on the given backend.
///
/// `weights` overrides where the `.bpk` weight file is read from; see
/// [`crate::models::resolve_weights`] for the default lookup.
pub fn load_built_in_model_with_backend(
    model_name: &str,
    backend: BackendKind,
    weights: Option<&Path>,
) -> Result<Model> {
    let built_in = BuiltInModel::from_name(model_name)?;
    let weights = resolve_weights(built_in.name(), weights)?;
    match &weights {
        Some(path) => info!(
            "Loading built-in model {} ({} backend) with weights from {:?}",
            model_name, backend, path
        ),
        None => info!(
            "Loading built-in model {} ({} backend) with embedded weights",
            model_name, backend
        ),
    }
    let model_size_bytes = weights
        .as_deref()
        .and_then(|p| std::fs::metadata(p).ok())
        .map_or(0, |m| m.len());
    let inner = built_in.create_model(backend, weights.as_deref())?;
    Ok(Model::new(inner, "onnx", model_size_bytes))
}

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

    fn get_test_model() -> Model {
        let dummy = DummyModel::new("test".to_string(), vec![784], vec![10]);
        Model::new(Box::new(dummy), "dummy", 0)
    }

    #[test]
    fn test_model_info() {
        let model = get_test_model();
        let info = model.get_info();
        assert_eq!(info.model_type, "dummy");
        assert_eq!(info.backend, "cpu");
        assert_eq!(info.input_spec.shape, vec![784]);
        assert_eq!(info.output_spec.shape, vec![10]);
        assert_eq!(info.input_spec.dtype, "float32");
    }

    #[test]
    fn test_model_loading_invalid_path() {
        let result = load_model(Path::new("nonexistent_model"));
        assert!(matches!(
            result,
            Err(crate::error::FurnaceError::Model(ModelError::FileNotFound(
                _
            )))
        ));
    }

    #[test]
    fn test_inference_with_valid_input() {
        let model = get_test_model();
        let output = model.predict(vec![0.5; 784]).unwrap();
        assert_eq!(output.len(), 10);
    }

    #[test]
    fn test_batch_inference_splits_outputs() {
        let model = get_test_model();
        let outputs = model.predict_batch(vec![vec![0.5; 784]; 3]).unwrap();
        assert_eq!(outputs.len(), 3);
        assert!(outputs.iter().all(|o| o.len() == 10));
    }

    #[test]
    fn test_predict_flat_matches_batch() {
        let model = get_test_model();
        let flat = model.predict_flat(vec![0.5; 784 * 3], 3).unwrap();
        let batch = model.predict_batch(vec![vec![0.5; 784]; 3]).unwrap();
        assert_eq!(flat, batch.concat());
        assert_eq!(model.get_stats().success_count, 2);
    }

    #[test]
    fn test_predict_flat_rejects_bad_input() {
        let model = get_test_model();
        // Length not a multiple of the batch
        assert!(matches!(
            model.predict_flat(vec![0.5; 784 * 2 + 1], 2),
            Err(crate::error::FurnaceError::Model(
                ModelError::InputValidation { .. }
            ))
        ));
        assert!(model.predict_flat(vec![], 0).is_err());
        assert!(
            model
                .predict_flat(vec![0.5; 784], MAX_BATCH_SIZE + 1)
                .is_err()
        );

        // Non-finite values report the batch item and index
        let mut input = vec![0.5; 784 * 2];
        input[784 + 5] = f32::NAN;
        let err = model.predict_flat(input, 2).unwrap_err().to_string();
        assert!(err.contains("batch[1][5]"), "{err}");
        assert_eq!(model.get_stats().error_count, 4);
    }

    #[test]
    fn test_inference_with_invalid_shape() {
        let model = get_test_model();
        let result = model.predict(vec![0.5; 100]);
        assert!(matches!(
            result,
            Err(crate::error::FurnaceError::Model(
                ModelError::InputValidation { .. }
            ))
        ));
    }

    #[test]
    fn test_inference_rejects_non_finite() {
        let model = get_test_model();
        let mut input = vec![0.5; 784];
        input[42] = f32::NAN;
        assert!(matches!(
            model.predict(input),
            Err(crate::error::FurnaceError::Model(
                ModelError::InvalidDataType(_)
            ))
        ));
    }

    #[test]
    fn test_model_stats() {
        let model = get_test_model();
        let stats = model.get_stats();
        assert_eq!(stats.inference_count, 0);
        assert!(stats.last_inference_time.is_none());

        model.predict(vec![0.5; 784]).unwrap();
        let _ = model.predict(vec![0.5; 1]);

        let stats = model.get_stats();
        assert_eq!(stats.inference_count, 2);
        assert_eq!(stats.success_count, 1);
        assert_eq!(stats.error_count, 1);
        assert!(stats.last_inference_time.is_some());
        assert!(stats.min_inference_time_ms.is_some());
    }

    #[test]
    fn test_real_mlp_inference_is_not_constant() {
        let model = load_model(Path::new("sample_model")).unwrap();
        let out = model.predict(vec![0.3; 784]).unwrap();
        assert_eq!(out.len(), 10);
        assert!(out.iter().any(|&v| v != out[0]));
    }

    #[cfg(model_resnet18)]
    #[test]
    fn test_built_in_model_missing_weights_file() {
        let result = load_built_in_model_with_backend(
            "resnet18",
            BackendKind::Cpu,
            Some(Path::new("does/not/exist.bpk")),
        );
        assert!(matches!(
            result,
            Err(crate::error::FurnaceError::Model(ModelError::FileNotFound(
                _
            )))
        ));
    }

    #[cfg(model_resnet18)]
    #[test]
    fn test_resnet18_weights_from_explicit_path() {
        // Copy the generated weights elsewhere to prove the path override is honored
        let dir = tempfile::tempdir().unwrap();
        let src = Path::new(env!("OUT_DIR")).join("models/resnet18.bpk");
        if !src.exists() {
            return; // built with embed-weights: no separate file
        }
        let dst = dir.path().join("custom.bpk");
        std::fs::copy(&src, &dst).unwrap();

        let model =
            load_built_in_model_with_backend("resnet18", BackendKind::Cpu, Some(&dst)).unwrap();
        assert_eq!(
            model.get_info().model_size_bytes,
            std::fs::metadata(&dst).unwrap().len()
        );
        assert_eq!(model.predict(vec![0.5; 3 * 224 * 224]).unwrap().len(), 1000);
    }

    #[cfg(model_resnet18)]
    #[test]
    fn test_resnet18_inference() {
        let model = load_built_in_model("resnet18").unwrap();
        assert_eq!(model.get_info().input_spec.shape, vec![3, 224, 224]);
        let out = model.predict(vec![0.5; 3 * 224 * 224]).unwrap();
        assert_eq!(out.len(), 1000);
        assert!(out.iter().all(|v| v.is_finite()));
    }

    #[cfg(model_mnist)]
    #[test]
    fn test_mnist_inference() {
        let model = load_built_in_model("mnist").unwrap();
        assert_eq!(model.get_info().input_spec.shape, vec![1, 28, 28]);
        assert_eq!(model.get_info().output_spec.shape, vec![10]);
        let outputs = model
            .predict_batch(vec![vec![0.0; 28 * 28], vec![255.0; 28 * 28]])
            .unwrap();
        assert!(outputs.iter().all(|o| o.len() == 10));
        assert_ne!(outputs[0], outputs[1]);
    }
}