use thiserror::Error;

/// Errors that can occur in diffusion model operations
#[derive(Error, Debug, Clone)]
pub enum DiffusionError {
    #[error("Invalid timestep: {timestep}, must be in range [0, {max_timesteps}]")]
    InvalidTimestep { timestep: f32, max_timesteps: u32 },

    #[error("Invalid noise schedule: {reason}")]
    InvalidNoiseSchedule { reason: String },

    #[error("Model architecture error: {details}")]
    ModelArchitecture { details: String },

    #[error("Tensor dimension mismatch: expected {expected:?}, got {actual:?}")]
    DimensionMismatch {
        expected: Vec<usize>,
        actual: Vec<usize>,
    },

    #[error("Scheduler error: {message}")]
    Scheduler { message: String },

    #[error("UNet forward pass failed: {reason}")]
    UNetForward { reason: String },

    #[error("DiT forward pass failed: {reason}")]
    DiTForward { reason: String },

    #[error("Runtime error: {message}")]
    Runtime { message: String },

    #[error("Tensor error: {0}")]
    TensorError(String),

    #[error("Tensor operation failed: {details}")]
    TensorOperation { details: String },

    #[error("Tensor error: {0}")]
    Tensor(rtx_tensor::TensorError),

    #[error("Not implemented: {0}")]
    NotImplemented(&'static str),
}

pub type Result<T> = std::result::Result<T, DiffusionError>;

impl From<rtx_tensor::TensorError> for DiffusionError {
    fn from(err: rtx_tensor::TensorError) -> Self {
        DiffusionError::Tensor(err)
    }
}