// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Comprehensive error handling for FEA operations.
//!
//! This module provides robust error types covering all aspects of finite element analysis,
//! including mesh operations, assembly, solving, and GPU operations.

use thiserror::Error;

/// Result type alias for FEA operations.
pub type FeaResult<T> = Result<T, FeaError>;

/// Comprehensive error types for finite element analysis operations.
#[derive(Error, Debug)]
pub enum FeaError {
    /// Mesh-related errors
    #[error("Mesh error: {0}")]
    Mesh(#[from] MeshError),

    /// Element-related errors
    #[error("Element error: {0}")]
    Element(#[from] ElementError),

    /// Material model errors
    #[error("Material error: {0}")]
    Material(#[from] MaterialError),

    /// Assembly process errors
    #[error("Assembly error: {0}")]
    Assembly(#[from] AssemblyError),

    /// Boundary condition errors
    #[error("Boundary condition error: {0}")]
    BoundaryCondition(#[from] BoundaryConditionError),

    /// Analysis errors
    #[error("Analysis error: {0}")]
    Analysis(#[from] AnalysisError),

    /// Solver errors
    #[error("Solver error: {0}")]
    Solver(#[from] SolverError),

    /// GPU/CUDA operation errors
    #[error("GPU error: {0}")]
    Gpu(#[from] GpuError),

    /// GPU kernel specific errors
    #[error("Kernel error: {0}")]
    Kernel(#[from] KernelError),

    /// Memory management errors
    #[error("Memory error: {0}")]
    Memory(#[from] MemoryError),

    /// General computational errors
    #[error("Computation error: {0}")]
    Computation(#[from] ComputationError),

    /// I/O and serialization errors
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    /// Generic error for wrapping external errors
    #[error("External error: {0}")]
    External(#[from] anyhow::Error),

    /// Invalid operation error
    #[error("Invalid operation: {message}")]
    InvalidOperation { message: String },

    /// Invalid input error
    #[error("Invalid input: {0}")]
    InvalidInput(String),

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

/// Mesh-specific error types.
#[derive(Error, Debug)]
pub enum MeshError {
    #[error("Invalid mesh topology: {message}")]
    InvalidTopology { message: String },

    #[error("Node index {index} out of bounds (max: {max_index})")]
    NodeIndexOutOfBounds { index: usize, max_index: usize },

    #[error("Element index {index} out of bounds (max: {max_index})")]
    ElementIndexOutOfBounds { index: usize, max_index: usize },

    #[error("Invalid element type: expected {expected}, found {found}")]
    InvalidElementType { expected: String, found: String },

    #[error("Mesh connectivity is invalid: {message}")]
    InvalidConnectivity { message: String },

    #[error("Mesh refinement failed: {reason}")]
    RefinementFailed { reason: String },

    #[error("Mesh partitioning failed: {reason}")]
    PartitioningFailed { reason: String },

    #[error("Empty mesh: no nodes or elements")]
    EmptyMesh,

    #[error(
        "Inconsistent mesh dimensions: nodes have {node_dim}D coordinates but elements require {element_dim}D"
    )]
    InconsistentDimensions { node_dim: usize, element_dim: usize },

    #[error("Invalid element: {0}")]
    InvalidElement(String),

    #[error("Node {node_id} not found")]
    NodeNotFound { node_id: usize },

    #[error("Element {element_id} not found")]
    ElementNotFound { element_id: usize },

    #[error("Set {set_name} not found")]
    SetNotFound { set_name: String },

    #[error("Invalid dimension: {0}")]
    InvalidDimension(String),
}

/// Element computation error types.
#[derive(Error, Debug)]
pub enum ElementError {
    #[error("Shape function evaluation failed at coordinates ({xi}, {eta}, {zeta})")]
    ShapeFunctionEvaluation { xi: f64, eta: f64, zeta: f64 },

    #[error("Jacobian computation failed: determinant is {det} (near zero or negative)")]
    JacobianSingular { det: f64 },

    #[error("Invalid quadrature rule: {rule} is not supported for element type {element_type}")]
    InvalidQuadratureRule { rule: String, element_type: String },

    #[error("Element matrix computation failed: {reason}")]
    MatrixComputationFailed { reason: String },

    #[error("Invalid element geometry: {message}")]
    InvalidGeometry { message: String },

    #[error("Unsupported element type: {element_type}")]
    UnsupportedElementType { element_type: String },

    #[error("Integration point {point} out of range [0, {max_points}]")]
    IntegrationPointOutOfRange { point: usize, max_points: usize },

    #[error("Unsupported quadrature order {requested_order} for {element_type} (max: {max_order})")]
    UnsupportedQuadratureOrder {
        element_type: String,
        requested_order: usize,
        max_order: usize,
    },

    #[error("Invalid quadrature: {reason}")]
    InvalidQuadrature { reason: String },

    #[error("GPU required for operation: {operation}")]
    GpuRequired { operation: String },
}

/// Material model error types.
#[derive(Error, Debug)]
pub enum MaterialError {
    #[error("Invalid material property: {property} = {value} is out of valid range [{min}, {max}]")]
    InvalidProperty {
        property: String,
        value: f64,
        min: f64,
        max: f64,
    },

    #[error("Material model convergence failed after {iterations} iterations")]
    ConvergenceFailed { iterations: usize },

    #[error("Unsupported material model: {model}")]
    UnsupportedModel { model: String },

    #[error("Material state update failed: {reason}")]
    StateUpdateFailed { reason: String },

    #[error("Invalid stress/strain state: {message}")]
    InvalidState { message: String },

    #[error("Material parameter {parameter} not found")]
    ParameterNotFound { parameter: String },

    #[error("Plastic consistency condition violated: f = {yield_function_value}")]
    PlasticConsistencyViolated { yield_function_value: f64 },
}

/// Assembly process error types.
#[derive(Error, Debug)]
pub enum AssemblyError {
    #[error(
        "Matrix dimension mismatch: expected {expected_rows}x{expected_cols}, got {actual_rows}x{actual_cols}"
    )]
    MatrixDimensionMismatch {
        expected_rows: usize,
        expected_cols: usize,
        actual_rows: usize,
        actual_cols: usize,
    },

    #[error("Sparse matrix assembly failed: {reason}")]
    SparseAssemblyFailed { reason: String },

    #[error("Global DOF index {index} out of bounds (max: {max_index})")]
    GlobalDofOutOfBounds { index: usize, max_index: usize },

    #[error("Element contribution assembly failed for element {element_id}: {reason}")]
    ElementContributionFailed { element_id: usize, reason: String },

    #[error("Constraint assembly failed: {reason}")]
    ConstraintAssemblyFailed { reason: String },

    #[error("Matrix pattern inconsistency: {message}")]
    MatrixPatternInconsistent { message: String },

    #[error("Node {node_id} not found")]
    NodeNotFound { node_id: usize },

    #[error("Material {material_id} not found")]
    MaterialNotFound { material_id: usize },

    #[error("Dimension mismatch: expected {expected}, got {actual}")]
    DimensionMismatch { expected: usize, actual: usize },

    #[error("Singular matrix detected")]
    SingularMatrix,
}

/// Boundary condition error types.
pub type BoundaryError = BoundaryConditionError;

/// Analysis error types.
#[derive(Error, Debug)]
pub enum AnalysisError {
    #[error("Analysis setup failed: {reason}")]
    SetupFailed { reason: String },

    #[error("Analysis convergence failed after {iterations} iterations")]
    ConvergenceFailed { iterations: usize },

    #[error("Time stepping failed at time {time}: {reason}")]
    TimeSteppingFailed { time: f64, reason: String },

    #[error("Invalid analysis parameters: {message}")]
    InvalidParameters { message: String },

    #[error("Analysis not supported: {analysis_type}")]
    UnsupportedAnalysis { analysis_type: String },

    #[error("Invalid configuration: {0}")]
    InvalidConfiguration(String),
}

/// Boundary condition error types.
#[derive(Error, Debug)]
pub enum BoundaryConditionError {
    #[error("Boundary condition applied to invalid node {node_id}")]
    InvalidNode { node_id: usize },

    #[error("Boundary condition applied to invalid DOF {dof} (valid range: 0-{max_dof})")]
    InvalidDof { dof: usize, max_dof: usize },

    #[error("Conflicting boundary conditions on node {node_id}, DOF {dof}")]
    ConflictingConditions { node_id: usize, dof: usize },

    #[error("Invalid boundary surface specification: {message}")]
    InvalidBoundarySurface { message: String },

    #[error("Traction boundary condition requires surface normal vector")]
    MissingSurfaceNormal,

    #[error("Invalid load pattern: {pattern}")]
    InvalidLoadPattern { pattern: String },

    #[error("DOF {dof} out of bounds (max: {max_dof})")]
    DofOutOfBounds { dof: usize, max_dof: usize },

    #[error("History data mismatch: {times_len} times vs {values_len} values")]
    HistoryDataMismatch { times_len: usize, values_len: usize },

    #[error("Node {node_id} not found")]
    NodeNotFound { node_id: usize },

    #[error("Element {element_id} not found")]
    ElementNotFound { element_id: usize },
}

/// Solver error types.
#[derive(Error, Debug)]
pub enum SolverError {
    #[error("Matrix is singular: condition number = {condition_number}")]
    SingularMatrix { condition_number: f64 },

    #[error(
        "Iterative solver failed to converge after {iterations} iterations (residual: {residual})"
    )]
    ConvergenceFailed { iterations: usize, residual: f64 },

    #[error("Direct solver failed: {reason}")]
    DirectSolverFailed { reason: String },

    #[error("Preconditioner setup failed: {reason}")]
    PreconditionerFailed { reason: String },

    #[error("Invalid solver parameters: {message}")]
    InvalidParameters { message: String },

    #[error("Solver not supported for matrix type: {matrix_type}")]
    UnsupportedMatrixType { matrix_type: String },

    #[error("Numerical instability detected: {message}")]
    NumericalInstability { message: String },

    #[error("Device is not available or unsupported")]
    DeviceUnavailable,

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

    #[error("Matrix dimension mismatch: {matrix_rows}x{matrix_cols} vs RHS {rhs_rows}x{rhs_cols}")]
    DimensionMismatch {
        matrix_rows: usize,
        matrix_cols: usize,
        rhs_rows: usize,
        rhs_cols: usize,
    },

    #[error("Matrix is not symmetric")]
    MatrixNotSymmetric,

    #[error("Solve operation failed: {reason}")]
    SolveError { reason: String },

    #[error("Factorization failed: {reason}")]
    FactorizationFailed { reason: String },

    #[error("Factorization required but not performed")]
    FactorizationRequired,

    #[error("GPU not available: {0}")]
    GpuNotAvailable(String),

    #[error("Preconditioner not setup")]
    PreconditionerNotSetup,
}

/// GPU operation error types.
#[derive(Error, Debug)]
pub enum GpuError {
    #[error("CUDA runtime error: {message}")]
    CudaRuntime { message: String },

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

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

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

    #[error("Kernel launch failed: {kernel_name}")]
    KernelLaunchFailed { kernel_name: String },

    #[error("GPU memory allocation failed: requested {requested} bytes")]
    MemoryAllocationFailed { requested: usize },

    #[error("GPU-CPU memory transfer failed: {direction}")]
    MemoryTransferFailed { direction: String },

    #[error("GPU device not available or unsupported")]
    DeviceUnavailable,

    #[error("GPU context error: {message}")]
    ContextError { message: String },
}

/// GPU kernel specific error types.
#[derive(Error, Debug)]
pub enum KernelError {
    #[error("Kernel compilation failed for {kernel_name}: {reason}")]
    CompilationFailed { kernel_name: String, reason: String },

    #[error("Kernel not found: {kernel_name} - {reason}")]
    KernelNotFound { kernel_name: String, reason: String },

    #[error("Kernel launch failed for {kernel_name}: {reason}")]
    LaunchFailed { kernel_name: String, reason: String },

    #[error("Device initialization failed: {reason}")]
    DeviceInitializationFailed { reason: String },

    #[error("Stream creation failed: {reason}")]
    StreamCreationFailed { reason: String },

    #[error("Synchronization failed: {reason}")]
    SynchronizationFailed { reason: String },

    #[error("Memory allocation failed: size {size} bytes - {reason}")]
    MemoryAllocationFailed { size: usize, reason: String },

    #[error("Memory copy failed: {direction} - {reason}")]
    MemoryCopyFailed { direction: String, reason: String },

    #[error("Device {device_id} not found")]
    DeviceNotFound { device_id: usize },

    #[error("Invalid kernel parameters: {message}")]
    InvalidParameters { message: String },

    #[error("Kernel execution timeout: {kernel_name}")]
    ExecutionTimeout { kernel_name: String },
}

/// Memory management error types.
#[derive(Error, Debug)]
pub enum MemoryError {
    #[error("Memory allocation failed: requested {size} bytes")]
    AllocationFailed { size: usize },

    #[error("Memory alignment error: required {required}, got {actual}")]
    AlignmentError { required: usize, actual: usize },

    #[error("Buffer overflow: accessing index {index} in buffer of size {size}")]
    BufferOverflow { index: usize, size: usize },

    #[error("Memory pool exhausted: {pool_name}")]
    PoolExhausted { pool_name: String },

    #[error("Invalid memory layout: {message}")]
    InvalidLayout { message: String },
}

/// General computation error types.
#[derive(Error, Debug)]
pub enum ComputationError {
    #[error("Numerical overflow in computation: {operation}")]
    NumericalOverflow { operation: String },

    #[error("Numerical underflow in computation: {operation}")]
    NumericalUnderflow { operation: String },

    #[error("Division by zero in {context}")]
    DivisionByZero { context: String },

    #[error("Invalid floating point result: {value} in {context}")]
    InvalidFloatingPoint { value: f64, context: String },

    #[error("Algorithm failed to converge: {algorithm}")]
    AlgorithmConvergence { algorithm: String },

    #[error("Dimension mismatch: expected {expected}, got {actual}")]
    DimensionMismatch { expected: String, actual: String },
}

impl FeaError {
    /// Check if the error is recoverable.
    pub fn is_recoverable(&self) -> bool {
        match self {
            Self::Solver(SolverError::ConvergenceFailed { .. }) => true,
            Self::Material(MaterialError::ConvergenceFailed { .. }) => true,
            Self::Gpu(GpuError::MemoryAllocationFailed { .. }) => true,
            _ => false,
        }
    }

    /// Get the error severity level.
    pub fn severity(&self) -> ErrorSeverity {
        match self {
            Self::Mesh(MeshError::EmptyMesh) => ErrorSeverity::Critical,
            Self::Solver(SolverError::SingularMatrix { .. }) => ErrorSeverity::Critical,
            Self::Gpu(GpuError::DeviceUnavailable) => ErrorSeverity::Critical,
            Self::Memory(MemoryError::AllocationFailed { .. }) => ErrorSeverity::High,
            Self::Element(ElementError::JacobianSingular { .. }) => ErrorSeverity::High,
            Self::Solver(SolverError::ConvergenceFailed { .. }) => ErrorSeverity::Medium,
            Self::Material(MaterialError::ConvergenceFailed { .. }) => ErrorSeverity::Medium,
            _ => ErrorSeverity::Low,
        }
    }
}

/// Error severity levels for proper error handling and logging.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorSeverity {
    Critical, // System cannot continue
    High,     // Major functionality affected
    Medium,   // Recoverable with potential data loss
    Low,      // Minor issues, operation can continue
}

/// Helper trait for error context enhancement.
pub trait FeaErrorExt<T> {
    /// Add context to a mesh error.
    fn with_mesh_context(self, context: &str) -> FeaResult<T>;

    /// Add element context.
    fn with_element_context(self, element_id: usize) -> FeaResult<T>;

    /// Add material context.
    fn with_material_context(self, material_id: usize) -> FeaResult<T>;
}

impl<T> FeaErrorExt<T> for FeaResult<T> {
    fn with_mesh_context(self, context: &str) -> Self {
        self.map_err(|e| match e {
            FeaError::Mesh(mesh_err) => FeaError::Mesh(match mesh_err {
                MeshError::InvalidTopology { message } => MeshError::InvalidTopology {
                    message: format!("{context}: {message}"),
                },
                other => other,
            }),
            other => other,
        })
    }

    fn with_element_context(self, element_id: usize) -> Self {
        self.map_err(|e| match e {
            FeaError::Element(elem_err) => FeaError::Element(match elem_err {
                ElementError::MatrixComputationFailed { reason } => {
                    ElementError::MatrixComputationFailed {
                        reason: format!("Element {element_id}: {reason}"),
                    }
                }
                other => other,
            }),
            other => other,
        })
    }

    fn with_material_context(self, material_id: usize) -> Self {
        self.map_err(|e| match e {
            FeaError::Material(mat_err) => FeaError::Material(match mat_err {
                MaterialError::StateUpdateFailed { reason } => MaterialError::StateUpdateFailed {
                    reason: format!("Material {material_id}: {reason}"),
                },
                other => other,
            }),
            other => other,
        })
    }
}

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

    #[test]
    fn test_error_severity() {
        let error = FeaError::Mesh(MeshError::EmptyMesh);
        assert_eq!(error.severity(), ErrorSeverity::Critical);

        let error = FeaError::Solver(SolverError::ConvergenceFailed {
            iterations: 100,
            residual: 1e-6,
        });
        assert_eq!(error.severity(), ErrorSeverity::Medium);
    }

    #[test]
    fn test_error_recoverability() {
        let error = FeaError::Solver(SolverError::ConvergenceFailed {
            iterations: 100,
            residual: 1e-6,
        });
        assert!(error.is_recoverable());

        let error = FeaError::Mesh(MeshError::EmptyMesh);
        assert!(!error.is_recoverable());
    }

    #[test]
    fn test_error_context() {
        let result: FeaResult<()> = Err(FeaError::Mesh(MeshError::InvalidTopology {
            message: "original message".to_string(),
        }));

        let with_context = result.with_mesh_context("preprocessing");
        if let Err(FeaError::Mesh(MeshError::InvalidTopology { message })) = with_context {
            assert!(message.contains("preprocessing"));
            assert!(message.contains("original message"));
        } else {
            panic!("Expected mesh error with context");
        }
    }
}