use std::fmt::{Display, Formatter};
use std::str::FromStr;

use crate::ir::dtype::DType;
use crate::ir::ops::{BinaryOp, UnaryOp};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId(pub usize);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ParamId(pub usize);

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

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

#[derive(Debug, Clone, PartialEq)]
pub enum Const {
    Bool(bool),
    Int(i64),
    Float(f64),
    Complex(f64, f64),
}

/// Known built-in functions in the Gonidium IR.
///
/// Surface builtins (`exp` through `softplus`, plus `floor` / `ceil`) are
/// available in `.gon` source
/// syntax. Fused builtins (`ExpSigmoid`, `SigmoidLog`) are synthesised by
/// the optimizer and never appear in user-written code.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Builtin {
    // ── Surface builtins (available in .gon source) ──
    Exp,
    Log,
    Sqrt,
    Abs,
    Sin,
    Cos,
    Tanh,
    Sigmoid,
    Relu,
    Softplus,
    Floor,
    Ceil,
    // ── Fused builtins (optimizer-synthesised only) ──
    ExpSigmoid,
    SigmoidLog,
}

impl Builtin {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Exp => "exp",
            Self::Log => "log",
            Self::Sqrt => "sqrt",
            Self::Abs => "abs",
            Self::Sin => "sin",
            Self::Cos => "cos",
            Self::Tanh => "tanh",
            Self::Sigmoid => "sigmoid",
            Self::Relu => "relu",
            Self::Softplus => "softplus",
            Self::Floor => "floor",
            Self::Ceil => "ceil",
            Self::ExpSigmoid => "exp_sigmoid",
            Self::SigmoidLog => "sigmoid_log",
        }
    }

    /// Number of arguments this builtin expects.
    pub const fn arity(self) -> usize {
        // All current builtins are unary.
        match self {
            Self::Exp
            | Self::Log
            | Self::Sqrt
            | Self::Abs
            | Self::Sin
            | Self::Cos
            | Self::Tanh
            | Self::Sigmoid
            | Self::Relu
            | Self::Softplus
            | Self::Floor
            | Self::Ceil
            | Self::ExpSigmoid
            | Self::SigmoidLog => 1,
        }
    }

    /// Returns `true` for builtins available in surface `.gon` syntax.
    pub const fn is_surface(self) -> bool {
        matches!(
            self,
            Self::Exp
                | Self::Log
                | Self::Sqrt
                | Self::Abs
                | Self::Sin
                | Self::Cos
                | Self::Tanh
                | Self::Sigmoid
                | Self::Relu
                | Self::Softplus
                | Self::Floor
                | Self::Ceil
        )
    }

    /// Single source of truth for surface-call builtins (user-facing `.gon` syntax).
    pub const SURFACE: [Builtin; 12] = [
        Builtin::Exp,
        Builtin::Log,
        Builtin::Sqrt,
        Builtin::Abs,
        Builtin::Sin,
        Builtin::Cos,
        Builtin::Tanh,
        Builtin::Sigmoid,
        Builtin::Relu,
        Builtin::Softplus,
        Builtin::Floor,
        Builtin::Ceil,
    ];
}

impl Display for Builtin {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

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

    #[test]
    fn surface_builtins_match_is_surface() {
        for &b in &Builtin::SURFACE {
            assert!(b.is_surface());
        }
    }
}

impl FromStr for Builtin {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "exp" => Ok(Self::Exp),
            "log" => Ok(Self::Log),
            "sqrt" => Ok(Self::Sqrt),
            "abs" => Ok(Self::Abs),
            "sin" => Ok(Self::Sin),
            "cos" => Ok(Self::Cos),
            "tanh" => Ok(Self::Tanh),
            "sigmoid" => Ok(Self::Sigmoid),
            "relu" => Ok(Self::Relu),
            "softplus" => Ok(Self::Softplus),
            "floor" => Ok(Self::Floor),
            "ceil" => Ok(Self::Ceil),
            "exp_sigmoid" => Ok(Self::ExpSigmoid),
            "sigmoid_log" => Ok(Self::SigmoidLog),
            _ => Err(()),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Node {
    Param(ParamId),
    Const(Const),
    Unary(UnaryOp, NodeId),
    Binary(BinaryOp, NodeId, NodeId),
    IfElse(NodeId, NodeId, NodeId),
    Call(Builtin, Vec<NodeId>),
    /// `Cast(target_dtype, child, is_explicit)` — `is_explicit` is `true` for
    /// casts written in source (`cast<T>(x)` or `x@T`), `false` for implicit
    /// promotions inserted by lowering or optimisation.
    Cast(DType, NodeId, bool),
}
