//! Error types for Burn integration

use thiserror::Error;

/// Result type for Burn operations
pub type Result<T> = std::result::Result<T, BurnError>;

/// Errors that can occur in Burn integration
#[derive(Debug, Error)]
pub enum BurnError {
    /// Failed to load model
    #[error("Failed to load model: {0}")]
    ModelLoad(String),

    /// Failed to create session
    #[error("Failed to create session: {0}")]
    SessionCreation(String),

    /// Inference failed
    #[error("Inference failed: {0}")]
    Inference(String),

    /// Tensor conversion error
    #[error("Tensor conversion error: {0}")]
    TensorConversion(String),

    /// Backend not available
    #[error("Backend not available: {0}")]
    BackendUnavailable(String),

    /// Shape mismatch
    #[error("Shape mismatch: expected {expected:?}, got {actual:?}")]
    ShapeMismatch {
        expected: Vec<usize>,
        actual: Vec<usize>,
    },

    /// Unsupported data type
    #[error("Unsupported data type: {0}")]
    UnsupportedDType(String),

    /// IO error
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    /// Serialization error
    #[error("Serialization error: {0}")]
    Serialization(String),

    /// Configuration error
    #[error("Configuration error: {0}")]
    Config(String),
}

impl BurnError {
    /// Create a model load error
    pub fn model_load(msg: impl Into<String>) -> Self {
        Self::ModelLoad(msg.into())
    }

    /// Create a tensor conversion error
    pub fn tensor_conversion(msg: impl Into<String>) -> Self {
        Self::TensorConversion(msg.into())
    }

    /// Create a shape mismatch error
    pub fn shape_mismatch(expected: Vec<usize>, actual: Vec<usize>) -> Self {
        Self::ShapeMismatch { expected, actual }
    }
}