// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Dynamo Error System
//!
//! This module provides a standardized error type for Dynamo with support for:
//! - Categorized error types via [`ErrorType`] enum
//! - Error chaining via the standard [`std::error::Error::source()`] method
//! - Serialization for network transmission via serde
//!
//! # DynamoError
//!
//! [`DynamoError`] is the standardized error type for Dynamo. It can be created
//! directly or converted from any [`std::error::Error`]:
//!
//! ```rust,ignore
//! use dynamo_runtime::error::{DynamoError, ErrorType};
//!
//! // Simple error
//! let err = DynamoError::msg("something failed");
//!
//! // Typed error with cause
//! let cause = std::io::Error::other("io error");
//! let err = DynamoError::builder()
//!     .error_type(ErrorType::Unknown)
//!     .message("operation failed")
//!     .cause(cause)
//!     .build();
//!
//! // Convert from any std::error::Error
//! let std_err = std::io::Error::other("io error");
//! let dynamo_err = DynamoError::from(Box::new(std_err) as Box<dyn std::error::Error>);
//! ```

use serde::{Deserialize, Serialize};
use std::fmt;

// ============================================================================
// ErrorType Enum
// ============================================================================

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ErrorType {
    /// Uncategorized or unknown error.
    Unknown,
    /// The request contains invalid input (e.g., prompt exceeds context length).
    InvalidArgument,
    /// Failed to establish a connection to a remote worker.
    CannotConnect,
    /// An established connection was lost unexpectedly.
    Disconnected,
    /// A connection or request timed out.
    ConnectionTimeout,
    /// The request was cancelled (e.g., client disconnected).
    Cancelled,
    /// Error originating from a backend engine.
    Backend(BackendError),
}

impl fmt::Display for ErrorType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ErrorType::Unknown => write!(f, "Unknown"),
            ErrorType::InvalidArgument => write!(f, "InvalidArgument"),
            ErrorType::CannotConnect => write!(f, "CannotConnect"),
            ErrorType::Disconnected => write!(f, "Disconnected"),
            ErrorType::ConnectionTimeout => write!(f, "ConnectionTimeout"),
            ErrorType::Cancelled => write!(f, "Cancelled"),
            ErrorType::Backend(sub) => write!(f, "Backend{sub}"),
        }
    }
}

/// Categorizes errors into a fixed set of standard types.
///
/// Consumers (e.g., the migration module) inspect the error type to decide
/// what action to take, rather than the error defining its own behavior.
/// Backend engine error subcategories.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum BackendError {
    /// Uncategorized or unknown backend error.
    Unknown,
    /// The request contains invalid input (e.g., prompt exceeds context length).
    InvalidArgument,
    /// Failed to establish a connection to a remote worker.
    CannotConnect,
    /// An established connection was lost unexpectedly.
    Disconnected,
    /// A connection or request timed out.
    ConnectionTimeout,
    /// The request was cancelled (e.g., client disconnected).
    Cancelled,
    /// The engine process has shut down or crashed.
    EngineShutdown,
    /// The response stream was terminated before completion (e.g., engine dropped mid-stream).
    StreamIncomplete,
}

impl fmt::Display for BackendError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BackendError::Unknown => write!(f, "Unknown"),
            BackendError::InvalidArgument => write!(f, "InvalidArgument"),
            BackendError::CannotConnect => write!(f, "CannotConnect"),
            BackendError::Disconnected => write!(f, "Disconnected"),
            BackendError::ConnectionTimeout => write!(f, "ConnectionTimeout"),
            BackendError::Cancelled => write!(f, "Cancelled"),
            BackendError::EngineShutdown => write!(f, "EngineShutdown"),
            BackendError::StreamIncomplete => write!(f, "StreamIncomplete"),
        }
    }
}

// ============================================================================
// DynamoError - The Standardized Error Type
// ============================================================================

/// The standardized error type for Dynamo.
///
/// `DynamoError` is a serializable, chainable error that:
/// - Carries an [`ErrorType`] for categorization
/// - Supports error chaining via [`std::error::Error::source()`]
/// - Is serializable for network transmission via `Annotated`
/// - Can be created from any [`std::error::Error`]
///
/// # Display
///
/// `Display` shows only the current error (standard Rust convention).
/// Use `source()` to walk the cause chain:
///
/// ```rust,ignore
/// let err = DynamoError::msg("outer");
/// println!("{}", err); // "Unknown: outer"
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DynamoError {
    error_type: ErrorType,
    message: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    caused_by: Option<Box<DynamoError>>,
}

impl DynamoError {
    /// Create a builder for constructing a `DynamoError`.
    pub fn builder() -> DynamoErrorBuilder {
        DynamoErrorBuilder::default()
    }

    /// Shorthand to create an `Unknown` error with just a message and no cause.
    pub fn msg(message: impl Into<String>) -> Self {
        Self::builder().message(message).build()
    }

    /// Returns the error type.
    pub fn error_type(&self) -> ErrorType {
        self.error_type
    }

    /// Returns the error message.
    pub fn message(&self) -> &str {
        &self.message
    }
}

impl fmt::Display for DynamoError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}: {}", self.error_type, self.message)
    }
}

impl std::error::Error for DynamoError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        self.caused_by
            .as_deref()
            .map(|e| e as &(dyn std::error::Error + 'static))
    }
}

/// Convert from a reference to any `std::error::Error`.
///
/// If the error is already a `DynamoError`, it is cloned. Otherwise, it is
/// wrapped as `ErrorType::Unknown` with the display string as the message.
/// The source chain is recursively converted, preserving `DynamoError` instances.
impl<'a> From<&'a (dyn std::error::Error + 'static)> for DynamoError {
    fn from(err: &'a (dyn std::error::Error + 'static)) -> Self {
        if let Some(dynamo_err) = err.downcast_ref::<DynamoError>() {
            return dynamo_err.clone();
        }

        Self {
            error_type: ErrorType::Unknown,
            message: err.to_string(),
            caused_by: err.source().map(|s| Box::new(DynamoError::from(s))),
        }
    }
}

/// Convert from an owned boxed `std::error::Error`.
///
/// If the error is already a `DynamoError`, ownership is taken without cloning.
/// Otherwise, falls back to the reference-based conversion.
impl From<Box<dyn std::error::Error + 'static>> for DynamoError {
    fn from(err: Box<dyn std::error::Error + 'static>) -> Self {
        match err.downcast::<DynamoError>() {
            Ok(dynamo_err) => *dynamo_err,
            Err(err) => DynamoError::from(&*err as &(dyn std::error::Error + 'static)),
        }
    }
}

// ============================================================================
// DynamoErrorBuilder
// ============================================================================

/// Builder for constructing a [`DynamoError`].
///
/// # Example
/// ```rust,ignore
/// let err = DynamoError::builder()
///     .error_type(ErrorType::Disconnected)
///     .message("worker lost")
///     .cause(some_io_error)
///     .build();
/// ```
#[derive(Default)]
pub struct DynamoErrorBuilder {
    error_type: Option<ErrorType>,
    message: Option<String>,
    caused_by: Option<Box<DynamoError>>,
}

impl DynamoErrorBuilder {
    /// Set the error type.
    pub fn error_type(mut self, error_type: ErrorType) -> Self {
        self.error_type = Some(error_type);
        self
    }

    /// Set the error message.
    pub fn message(mut self, message: impl Into<String>) -> Self {
        self.message = Some(message.into());
        self
    }

    /// Set the cause from any `std::error::Error`.
    ///
    /// If the cause is already a `DynamoError`, it is preserved as-is.
    /// Otherwise, it is converted to a `DynamoError` with `ErrorType::Unknown`.
    pub fn cause(mut self, cause: impl std::error::Error + 'static) -> Self {
        self.caused_by = Some(Box::new(DynamoError::from(
            &cause as &(dyn std::error::Error + 'static),
        )));
        self
    }

    /// Build the `DynamoError`.
    ///
    /// Defaults: `error_type` → `Unknown`, `message` → `""`, `cause` → `None`.
    pub fn build(self) -> DynamoError {
        DynamoError {
            error_type: self.error_type.unwrap_or(ErrorType::Unknown),
            message: self.message.unwrap_or_default(),
            caused_by: self.caused_by,
        }
    }
}

// ============================================================================
// Utility Functions
// ============================================================================

/// Check whether an error chain contains a specific set of error types
/// while not containing any of the excluded error types.
///
/// Walks the chain via `source()`, inspecting each error that can be downcast
/// to `DynamoError`. Returns `false` immediately if any error's type is in
/// `exclude_set`. Otherwise, returns `true` if at least one error's type is
/// in `match_set`. Errors that are not `DynamoError` are skipped.
pub fn match_error_chain(
    err: &(dyn std::error::Error + 'static),
    match_set: &[ErrorType],
    exclude_set: &[ErrorType],
) -> bool {
    let mut found = false;
    let mut current: Option<&(dyn std::error::Error + 'static)> = Some(err);

    while let Some(e) = current {
        if let Some(dynamo_err) = e.downcast_ref::<DynamoError>() {
            if exclude_set.contains(&dynamo_err.error_type()) {
                return false;
            }
            if match_set.contains(&dynamo_err.error_type()) {
                found = true;
            }
        }
        current = e.source();
    }

    found
}

// ============================================================================
// Tests
// ============================================================================

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

    // Compile-time assertions that DynamoError is std::error::Error + Send + Sync + 'static.
    // These fail at compile time if a future change breaks these guarantees.
    const _: () = {
        fn assert_stderror<T: std::error::Error>() {}
        fn assert_send<T: Send>() {}
        fn assert_sync<T: Sync>() {}
        fn assert_static<T: 'static>() {}
        fn assert_all() {
            assert_stderror::<DynamoError>();
            assert_send::<DynamoError>();
            assert_sync::<DynamoError>();
            assert_static::<DynamoError>();
        }
    };

    #[test]
    fn test_msg_constructor() {
        let err = DynamoError::msg("something failed");
        assert_eq!(err.error_type(), ErrorType::Unknown);
        assert_eq!(err.message(), "something failed");
        assert!(err.source().is_none());
    }

    #[test]
    fn test_new_constructor_with_cause() {
        let cause = std::io::Error::other("io error");
        let err = DynamoError::builder()
            .error_type(ErrorType::Unknown)
            .message("operation failed")
            .cause(cause)
            .build();

        assert_eq!(err.error_type(), ErrorType::Unknown);
        assert_eq!(err.message(), "operation failed");
        assert!(err.source().is_some());
    }

    #[test]
    fn test_display_shows_only_current_error() {
        let cause = std::io::Error::other("io error");
        let err = DynamoError::builder()
            .error_type(ErrorType::Unknown)
            .message("operation failed")
            .cause(cause)
            .build();

        // Display should only show the current error, not the chain
        assert_eq!(err.to_string(), "Unknown: operation failed");
    }

    #[test]
    fn test_source_chain() {
        let cause = std::io::Error::other("io error");
        let err = DynamoError::builder()
            .error_type(ErrorType::Unknown)
            .message("operation failed")
            .cause(cause)
            .build();

        // source() should return the cause
        let source = err.source().unwrap();
        assert!(source.to_string().contains("io error"));
    }

    #[test]
    fn test_from_boxed_std_error() {
        let std_err = std::io::Error::other("io error");
        let boxed: Box<dyn std::error::Error> = Box::new(std_err);
        let dynamo_err = DynamoError::from(boxed);

        assert_eq!(dynamo_err.error_type(), ErrorType::Unknown);
        assert_eq!(dynamo_err.message(), "io error");
    }

    #[test]
    fn test_from_boxed_takes_ownership_of_dynamo_error() {
        let inner = DynamoError::msg("original");
        let boxed: Box<dyn std::error::Error> = Box::new(inner);
        let dynamo_err = DynamoError::from(boxed);

        // Should take ownership, not clone or wrap
        assert_eq!(dynamo_err.error_type(), ErrorType::Unknown);
        assert_eq!(dynamo_err.message(), "original");
    }

    #[test]
    fn test_from_boxed_with_source_chain() {
        #[derive(Debug)]
        struct OuterError {
            source: std::io::Error,
        }

        impl fmt::Display for OuterError {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                write!(f, "outer error occurred")
            }
        }

        impl std::error::Error for OuterError {
            fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
                Some(&self.source)
            }
        }

        let inner = std::io::Error::new(std::io::ErrorKind::NotFound, "file not found");
        let outer = OuterError { source: inner };
        let boxed: Box<dyn std::error::Error> = Box::new(outer);
        let dynamo_err = DynamoError::from(boxed);

        assert_eq!(dynamo_err.message(), "outer error occurred");
        assert!(dynamo_err.source().is_some());

        let cause = dynamo_err.source().unwrap();
        assert!(cause.to_string().contains("file not found"));
    }

    #[test]
    fn test_serialization_roundtrip() {
        let cause = DynamoError::msg("inner cause");
        let err = DynamoError::builder()
            .error_type(ErrorType::Unknown)
            .message("outer error")
            .cause(cause)
            .build();

        let json = serde_json::to_string(&err).unwrap();
        let deserialized: DynamoError = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.error_type(), ErrorType::Unknown);
        assert_eq!(deserialized.message(), "outer error");
        assert!(deserialized.source().is_some());

        let cause = deserialized
            .source()
            .unwrap()
            .downcast_ref::<DynamoError>()
            .unwrap();
        assert_eq!(cause.message(), "inner cause");
    }

    #[test]
    fn test_error_type_display() {
        assert_eq!(ErrorType::Unknown.to_string(), "Unknown");
        assert_eq!(ErrorType::InvalidArgument.to_string(), "InvalidArgument");
        assert_eq!(ErrorType::CannotConnect.to_string(), "CannotConnect");
        assert_eq!(ErrorType::Disconnected.to_string(), "Disconnected");
        assert_eq!(
            ErrorType::ConnectionTimeout.to_string(),
            "ConnectionTimeout"
        );
        assert_eq!(ErrorType::Cancelled.to_string(), "Cancelled");
        assert_eq!(
            ErrorType::Backend(BackendError::Unknown).to_string(),
            "BackendUnknown"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::InvalidArgument).to_string(),
            "BackendInvalidArgument"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::CannotConnect).to_string(),
            "BackendCannotConnect"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::Disconnected).to_string(),
            "BackendDisconnected"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::ConnectionTimeout).to_string(),
            "BackendConnectionTimeout"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::Cancelled).to_string(),
            "BackendCancelled"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::EngineShutdown).to_string(),
            "BackendEngineShutdown"
        );
        assert_eq!(
            ErrorType::Backend(BackendError::StreamIncomplete).to_string(),
            "BackendStreamIncomplete"
        );
    }
}
