//! Node representation for ONNX operations.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::error::Result;
use crate::onnx_proto::{AttributeProto, NodeProto};

/// ONNX operator kinds.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum NodeKind {
    // Arithmetic
    /// Element-wise addition of two tensors.
    Add,
    /// Element-wise subtraction of two tensors.
    Sub,
    /// Element-wise multiplication of two tensors.
    Mul,
    /// Element-wise division of two tensors.
    Div,
    /// Element-wise negation of a tensor.
    Neg,
    /// Element-wise absolute value of a tensor.
    Abs,
    /// Element-wise square root of a tensor.
    Sqrt,
    /// Element-wise power operation.
    Pow,
    /// Element-wise exponential function.
    Exp,
    /// Element-wise natural logarithm.
    Log,
    /// Element-wise floor rounding.
    Floor,
    /// Element-wise ceiling rounding.
    Ceil,
    /// Element-wise rounding to nearest integer.
    Round,

    // Matrix operations
    /// Matrix multiplication of two tensors.
    MatMul,
    /// General matrix multiplication with optional bias.
    Gemm,

    // Activations
    /// Rectified Linear Unit activation function.
    Relu,
    /// Leaky Rectified Linear Unit activation function.
    LeakyRelu,
    /// Sigmoid activation function.
    Sigmoid,
    /// Hyperbolic tangent activation function.
    Tanh,
    /// Softmax activation function.
    Softmax,
    /// Log-softmax activation function.
    LogSoftmax,
    /// Gaussian Error Linear Unit activation function.
    Gelu,
    /// Sigmoid Linear Unit (Swish) activation function.
    Silu,
    /// Hard sigmoid activation function.
    HardSigmoid,
    /// Hard swish activation function.
    HardSwish,
    /// Exponential Linear Unit activation function.
    Elu,
    /// Scaled Exponential Linear Unit activation function.
    Selu,
    /// Continuously Differentiable Exponential Linear Unit.
    Celu,
    /// Softplus activation function.
    Softplus,
    /// Softsign activation function.
    Softsign,
    /// Mish activation function.
    Mish,

    // Normalization
    /// Batch normalization operation.
    BatchNormalization,
    /// Layer normalization operation.
    LayerNormalization,
    /// Instance normalization operation.
    InstanceNormalization,
    /// Group normalization operation.
    GroupNormalization,
    /// Lp normalization operation.
    LpNormalization,

    // Convolution
    /// Convolution operation.
    Conv,
    /// Transposed convolution (deconvolution) operation.
    ConvTranspose,

    // Pooling
    /// Max pooling operation.
    MaxPool,
    /// Average pooling operation.
    AveragePool,
    /// Global average pooling over spatial dimensions.
    GlobalAveragePool,
    /// Global max pooling over spatial dimensions.
    GlobalMaxPool,
    /// Global Lp pooling over spatial dimensions.
    GlobalLpPool,

    // Reduction
    /// Sum reduction along specified axes.
    ReduceSum,
    /// Mean reduction along specified axes.
    ReduceMean,
    /// Max reduction along specified axes.
    ReduceMax,
    /// Min reduction along specified axes.
    ReduceMin,
    /// Product reduction along specified axes.
    ReduceProd,
    /// L1 norm reduction along specified axes.
    ReduceL1,
    /// L2 norm reduction along specified axes.
    ReduceL2,
    /// Log of sum reduction along specified axes.
    ReduceLogSum,
    /// Log of sum of exponentials reduction along specified axes.
    ReduceLogSumExp,
    /// Sum of squares reduction along specified axes.
    ReduceSumSquare,

    // Shape operations
    /// Reshape tensor to a new shape.
    Reshape,
    /// Transpose tensor dimensions.
    Transpose,
    /// Flatten tensor to 2D.
    Flatten,
    /// Remove dimensions of size 1.
    Squeeze,
    /// Insert dimensions of size 1.
    Unsqueeze,
    /// Concatenate tensors along an axis.
    Concat,
    /// Split tensor into multiple tensors.
    Split,
    /// Extract a slice from a tensor.
    Slice,
    /// Gather elements along an axis by index.
    Gather,
    /// Gather elements by index along specified axis.
    GatherElements,
    /// Gather slices from tensor using N-dimensional indices.
    GatherND,
    /// Scatter elements along an axis by index.
    Scatter,
    /// Scatter elements by index along specified axis.
    ScatterElements,
    /// Scatter slices into tensor using N-dimensional indices.
    ScatterND,
    /// Broadcast tensor to a larger shape.
    Expand,
    /// Tile tensor by repeating along dimensions.
    Tile,
    /// Pad tensor with constant or edge values.
    Pad,

    // Comparison
    /// Element-wise equality comparison.
    Equal,
    /// Element-wise greater-than comparison.
    Greater,
    /// Element-wise greater-than-or-equal comparison.
    GreaterOrEqual,
    /// Element-wise less-than comparison.
    Less,
    /// Element-wise less-than-or-equal comparison.
    LessOrEqual,
    /// Element-wise logical NOT.
    Not,
    /// Element-wise logical AND.
    And,
    /// Element-wise logical OR.
    Or,
    /// Element-wise logical XOR.
    Xor,
    /// Conditional selection based on a mask tensor.
    Where,

    // Type conversion
    /// Cast tensor to a different data type.
    Cast,
    /// Cast tensor to the same type as another tensor.
    CastLike,

    // Constants
    /// Produce a constant tensor.
    Constant,
    /// Generate a tensor of a given shape filled with a constant.
    ConstantOfShape,
    /// Get the shape of a tensor as a 1D tensor.
    Shape,
    /// Get the total number of elements in a tensor.
    Size,

    // RNN
    /// Long Short-Term Memory recurrent layer.
    LSTM,
    /// Gated Recurrent Unit recurrent layer.
    GRU,
    /// Simple recurrent neural network layer.
    RNN,

    // Attention
    /// Single-head attention mechanism.
    Attention,
    /// Multi-head attention mechanism.
    MultiHeadAttention,

    // Misc
    /// Dropout regularization layer.
    Dropout,
    /// Identity operation (pass-through).
    Identity,
    /// Clip values to a specified range.
    Clip,
    /// Gauss error function.
    Erf,
    /// Element-wise sine function.
    Sin,
    /// Element-wise cosine function.
    Cos,
    /// Element-wise tangent function.
    Tan,
    /// Element-wise arcsine function.
    Asin,
    /// Element-wise arccosine function.
    Acos,
    /// Element-wise arctangent function.
    Atan,
    /// Element-wise hyperbolic sine function.
    Sinh,
    /// Element-wise hyperbolic cosine function.
    Cosh,
    /// Element-wise inverse hyperbolic sine function.
    Asinh,
    /// Element-wise inverse hyperbolic cosine function.
    Acosh,
    /// Element-wise inverse hyperbolic tangent function.
    Atanh,
    /// Element-wise sign function.
    Sign,
    /// Element-wise reciprocal (1/x).
    Reciprocal,
    /// Element-wise minimum of input tensors.
    Min,
    /// Element-wise maximum of input tensors.
    Max,
    /// Element-wise mean of input tensors.
    Mean,
    /// Element-wise sum of input tensors.
    Sum,

    // Embedding
    /// Embedding lookup operation.
    Embedding,
    /// One-hot encoding operation.
    OneHot,

    // Resize
    /// Resize tensor using interpolation.
    Resize,
    /// Upsample tensor (deprecated, use Resize).
    Upsample,

    // Custom/Unknown
    /// Custom or unsupported operator with its name.
    Custom(String),
}

impl NodeKind {
    /// Parse from ONNX op_type string.
    pub fn from_op_type(op_type: &str) -> Self {
        match op_type {
            // Arithmetic
            "Add" => NodeKind::Add,
            "Sub" => NodeKind::Sub,
            "Mul" => NodeKind::Mul,
            "Div" => NodeKind::Div,
            "Neg" => NodeKind::Neg,
            "Abs" => NodeKind::Abs,
            "Sqrt" => NodeKind::Sqrt,
            "Pow" => NodeKind::Pow,
            "Exp" => NodeKind::Exp,
            "Log" => NodeKind::Log,
            "Floor" => NodeKind::Floor,
            "Ceil" => NodeKind::Ceil,
            "Round" => NodeKind::Round,

            // Matrix
            "MatMul" => NodeKind::MatMul,
            "Gemm" => NodeKind::Gemm,

            // Activations
            "Relu" => NodeKind::Relu,
            "LeakyRelu" => NodeKind::LeakyRelu,
            "Sigmoid" => NodeKind::Sigmoid,
            "Tanh" => NodeKind::Tanh,
            "Softmax" => NodeKind::Softmax,
            "LogSoftmax" => NodeKind::LogSoftmax,
            "Gelu" => NodeKind::Gelu,
            "Silu" => NodeKind::Silu,
            "HardSigmoid" => NodeKind::HardSigmoid,
            "HardSwish" => NodeKind::HardSwish,
            "Elu" => NodeKind::Elu,
            "Selu" => NodeKind::Selu,
            "Celu" => NodeKind::Celu,
            "Softplus" => NodeKind::Softplus,
            "Softsign" => NodeKind::Softsign,
            "Mish" => NodeKind::Mish,

            // Normalization
            "BatchNormalization" => NodeKind::BatchNormalization,
            "LayerNormalization" => NodeKind::LayerNormalization,
            "InstanceNormalization" => NodeKind::InstanceNormalization,
            "GroupNormalization" => NodeKind::GroupNormalization,
            "LpNormalization" => NodeKind::LpNormalization,

            // Convolution
            "Conv" => NodeKind::Conv,
            "ConvTranspose" => NodeKind::ConvTranspose,

            // Pooling
            "MaxPool" => NodeKind::MaxPool,
            "AveragePool" => NodeKind::AveragePool,
            "GlobalAveragePool" => NodeKind::GlobalAveragePool,
            "GlobalMaxPool" => NodeKind::GlobalMaxPool,
            "GlobalLpPool" => NodeKind::GlobalLpPool,

            // Reduction
            "ReduceSum" => NodeKind::ReduceSum,
            "ReduceMean" => NodeKind::ReduceMean,
            "ReduceMax" => NodeKind::ReduceMax,
            "ReduceMin" => NodeKind::ReduceMin,
            "ReduceProd" => NodeKind::ReduceProd,
            "ReduceL1" => NodeKind::ReduceL1,
            "ReduceL2" => NodeKind::ReduceL2,
            "ReduceLogSum" => NodeKind::ReduceLogSum,
            "ReduceLogSumExp" => NodeKind::ReduceLogSumExp,
            "ReduceSumSquare" => NodeKind::ReduceSumSquare,

            // Shape
            "Reshape" => NodeKind::Reshape,
            "Transpose" => NodeKind::Transpose,
            "Flatten" => NodeKind::Flatten,
            "Squeeze" => NodeKind::Squeeze,
            "Unsqueeze" => NodeKind::Unsqueeze,
            "Concat" => NodeKind::Concat,
            "Split" => NodeKind::Split,
            "Slice" => NodeKind::Slice,
            "Gather" => NodeKind::Gather,
            "GatherElements" => NodeKind::GatherElements,
            "GatherND" => NodeKind::GatherND,
            "Scatter" => NodeKind::Scatter,
            "ScatterElements" => NodeKind::ScatterElements,
            "ScatterND" => NodeKind::ScatterND,
            "Expand" => NodeKind::Expand,
            "Tile" => NodeKind::Tile,
            "Pad" => NodeKind::Pad,

            // Comparison
            "Equal" => NodeKind::Equal,
            "Greater" => NodeKind::Greater,
            "GreaterOrEqual" => NodeKind::GreaterOrEqual,
            "Less" => NodeKind::Less,
            "LessOrEqual" => NodeKind::LessOrEqual,
            "Not" => NodeKind::Not,
            "And" => NodeKind::And,
            "Or" => NodeKind::Or,
            "Xor" => NodeKind::Xor,
            "Where" => NodeKind::Where,

            // Type
            "Cast" => NodeKind::Cast,
            "CastLike" => NodeKind::CastLike,

            // Constants
            "Constant" => NodeKind::Constant,
            "ConstantOfShape" => NodeKind::ConstantOfShape,
            "Shape" => NodeKind::Shape,
            "Size" => NodeKind::Size,

            // RNN
            "LSTM" => NodeKind::LSTM,
            "GRU" => NodeKind::GRU,
            "RNN" => NodeKind::RNN,

            // Attention
            "Attention" => NodeKind::Attention,
            "MultiHeadAttention" => NodeKind::MultiHeadAttention,

            // Misc
            "Dropout" => NodeKind::Dropout,
            "Identity" => NodeKind::Identity,
            "Clip" => NodeKind::Clip,
            "Erf" => NodeKind::Erf,
            "Sin" => NodeKind::Sin,
            "Cos" => NodeKind::Cos,
            "Tan" => NodeKind::Tan,
            "Asin" => NodeKind::Asin,
            "Acos" => NodeKind::Acos,
            "Atan" => NodeKind::Atan,
            "Sinh" => NodeKind::Sinh,
            "Cosh" => NodeKind::Cosh,
            "Asinh" => NodeKind::Asinh,
            "Acosh" => NodeKind::Acosh,
            "Atanh" => NodeKind::Atanh,
            "Sign" => NodeKind::Sign,
            "Reciprocal" => NodeKind::Reciprocal,
            "Min" => NodeKind::Min,
            "Max" => NodeKind::Max,
            "Mean" => NodeKind::Mean,
            "Sum" => NodeKind::Sum,

            // Embedding
            "Embedding" => NodeKind::Embedding,
            "OneHot" => NodeKind::OneHot,

            // Resize
            "Resize" => NodeKind::Resize,
            "Upsample" => NodeKind::Upsample,

            // Unknown
            other => NodeKind::Custom(other.to_string()),
        }
    }

    /// Check if this is a supported operator.
    pub fn is_supported(&self) -> bool {
        !matches!(self, NodeKind::Custom(_))
    }
}

/// Attribute value types.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AttributeValue {
    /// Floating point value.
    Float(f32),
    /// Integer value.
    Int(i64),
    /// String value.
    String(String),
    /// List of floats.
    Floats(Vec<f32>),
    /// List of integers.
    Ints(Vec<i64>),
    /// List of strings.
    Strings(Vec<String>),
}

/// A node in the ONNX graph.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Node {
    /// Node name (may be empty).
    pub name: String,
    /// Operator kind.
    pub kind: NodeKind,
    /// Input tensor names.
    pub inputs: Vec<String>,
    /// Output tensor names.
    pub outputs: Vec<String>,
    /// Node attributes.
    pub attributes: HashMap<String, AttributeValue>,
}

impl Node {
    /// Create from ONNX NodeProto.
    pub fn from_proto(proto: &NodeProto) -> Result<Self> {
        let name = proto.name.clone();
        let kind = NodeKind::from_op_type(&proto.op_type);
        let inputs = proto.input.clone();
        let outputs = proto.output.clone();
        let attributes = Self::parse_attributes(&proto.attribute)?;

        Ok(Self {
            name,
            kind,
            inputs,
            outputs,
            attributes,
        })
    }

    fn parse_attributes(attrs: &[AttributeProto]) -> Result<HashMap<String, AttributeValue>> {
        let mut result = HashMap::new();

        for attr in attrs {
            let value = Self::parse_attribute_value(attr)?;
            result.insert(attr.name.clone(), value);
        }

        Ok(result)
    }

    fn parse_attribute_value(attr: &AttributeProto) -> Result<AttributeValue> {
        // Check type field first, then fall back to presence of data
        match attr.r#type {
            1 => Ok(AttributeValue::Float(attr.f)),
            2 => Ok(AttributeValue::Int(attr.i)),
            3 => Ok(AttributeValue::String(
                String::from_utf8_lossy(&attr.s).to_string(),
            )),
            6 => Ok(AttributeValue::Floats(attr.floats.clone())),
            7 => Ok(AttributeValue::Ints(attr.ints.clone())),
            8 => Ok(AttributeValue::Strings(
                attr.strings
                    .iter()
                    .map(|s| String::from_utf8_lossy(s).to_string())
                    .collect(),
            )),
            _ => {
                // Fall back to checking which field has data
                if attr.f != 0.0 {
                    Ok(AttributeValue::Float(attr.f))
                } else if attr.i != 0 {
                    Ok(AttributeValue::Int(attr.i))
                } else if !attr.s.is_empty() {
                    Ok(AttributeValue::String(
                        String::from_utf8_lossy(&attr.s).to_string(),
                    ))
                } else if !attr.floats.is_empty() {
                    Ok(AttributeValue::Floats(attr.floats.clone()))
                } else if !attr.ints.is_empty() {
                    Ok(AttributeValue::Ints(attr.ints.clone()))
                } else {
                    // Default to int 0
                    Ok(AttributeValue::Int(0))
                }
            }
        }
    }

    /// Get an integer attribute.
    pub fn get_int(&self, name: &str) -> Option<i64> {
        match self.attributes.get(name) {
            Some(AttributeValue::Int(v)) => Some(*v),
            _ => None,
        }
    }

    /// Get a float attribute.
    pub fn get_float(&self, name: &str) -> Option<f32> {
        match self.attributes.get(name) {
            Some(AttributeValue::Float(v)) => Some(*v),
            _ => None,
        }
    }

    /// Get a string attribute.
    pub fn get_string(&self, name: &str) -> Option<&str> {
        match self.attributes.get(name) {
            Some(AttributeValue::String(v)) => Some(v),
            _ => None,
        }
    }

    /// Get an integer list attribute.
    pub fn get_ints(&self, name: &str) -> Option<&[i64]> {
        match self.attributes.get(name) {
            Some(AttributeValue::Ints(v)) => Some(v),
            _ => None,
        }
    }

    /// Get a float list attribute.
    pub fn get_floats(&self, name: &str) -> Option<&[f32]> {
        match self.attributes.get(name) {
            Some(AttributeValue::Floats(v)) => Some(v),
            _ => None,
        }
    }
}