use egg::{Id, RecExpr};
use std::str::FromStr;

use crate::diag::Diag;
use crate::ir::{
    BinaryOp, Builtin, Const, DType, DTypeKind, Node, NodeId, ParamId, TypedDag, TypedDagError,
    UnaryOp, promote,
};
use crate::opt::OptError;
use crate::opt::error::opt_promotion;
use crate::opt::lang::MathLang;
use crate::opt::strip::TypeMap;

fn invalid_node(idx: usize, enode: &MathLang) -> OptError {
    OptError::InvalidExpr {
        msg: format!("invalid node #{idx}: {enode:?}"),
    }
}

fn expect_num(
    expr: &RecExpr<MathLang>,
    id: Id,
    idx: usize,
    enode: &MathLang,
) -> Result<f64, OptError> {
    match expr[id] {
        MathLang::Num(v) => Ok(v.into_inner()),
        _ => Err(invalid_node(idx, enode)),
    }
}

fn expect_int(
    expr: &RecExpr<MathLang>,
    id: Id,
    idx: usize,
    enode: &MathLang,
) -> Result<i64, OptError> {
    match expr[id] {
        MathLang::Int(v) => Ok(v),
        _ => Err(invalid_node(idx, enode)),
    }
}

fn expect_token(
    expr: &RecExpr<MathLang>,
    id: Id,
    idx: usize,
    enode: &MathLang,
) -> Result<DType, OptError> {
    match expr[id] {
        MathLang::TypeToken(sym) => {
            DType::from_str(sym.as_str()).map_err(|_| invalid_node(idx, enode))
        }
        _ => Err(invalid_node(idx, enode)),
    }
}

fn expect_node_id(
    id_map: &[Option<NodeId>],
    id: Id,
    idx: usize,
    enode: &MathLang,
) -> Result<NodeId, OptError> {
    id_map
        .get(usize::from(id))
        .and_then(|mapped| *mapped)
        .ok_or_else(|| invalid_node(idx, enode))
}

fn try_promote(
    diags: &mut Vec<Diag>,
    where_: &str,
    a: DType,
    b: DType,
    idx: usize,
    enode: &MathLang,
) -> Result<DType, OptError> {
    let out = promote(a, b).map_err(|_| invalid_node(idx, enode))?;
    if a != out || b != out {
        diags.push(opt_promotion(where_, idx, a, b, out));
    }
    Ok(out)
}

/// Prefer [`TypeMap::node_types`]`[idx]` (one entry per enode from [`crate::opt::strip::strip_types`]).
/// After e-graph extraction the expression may grow; for **leaf** nodes only we fall back to
/// conservative defaults instead of silently coercing every unknown node to `f64`.
fn dtype_for_restored_leaf(
    type_map: &TypeMap,
    idx: usize,
    enode: &MathLang,
) -> Result<DType, OptError> {
    if let Some(dt) = type_map.node_types.get(idx).copied() {
        return Ok(dt);
    }
    match enode {
        MathLang::Int(_) => Ok(DType::I64),
        MathLang::Num(_) => Ok(DType::F64),
        MathLang::Complex(_) => Ok(DType::C128),
        _ => Err(OptError::InvalidExpr {
            msg: format!("missing dtype for restored node #{idx} ({enode:?})"),
        }),
    }
}

pub fn restore_types_with_info(
    expr: &RecExpr<MathLang>,
    type_map: &TypeMap,
    params: &[(String, DType)],
) -> Result<(TypedDag, Vec<Diag>), OptError> {
    let mut diags: Vec<Diag> = Vec::new();
    let mut nodes: Vec<Node> = Vec::with_capacity(expr.as_ref().len());
    let mut dtypes: Vec<DType> = Vec::with_capacity(expr.as_ref().len());
    let mut id_map: Vec<Option<NodeId>> = vec![None; expr.as_ref().len()];

    for (idx, enode) in expr.as_ref().iter().enumerate() {
        let current = Id::from(idx);
        let converted: Option<(Node, DType)> = match enode {
            MathLang::TypeToken(_) => None,
            MathLang::Bool(v) => Some((Node::Const(Const::Bool(*v)), DType::BOOL)),
            MathLang::Int(v) => {
                let dtype = dtype_for_restored_leaf(type_map, idx, enode)?;
                match dtype.kind {
                    DTypeKind::UInt | DTypeKind::SInt => Some((Node::Const(Const::Int(*v)), dtype)),
                    DTypeKind::Float(_) => Some((Node::Const(Const::Float(*v as f64)), dtype)),
                    DTypeKind::Complex(_) => {
                        Some((Node::Const(Const::Complex(*v as f64, 0.0)), dtype))
                    }
                    DTypeKind::Bool => Some((Node::Const(Const::Bool(*v != 0)), DType::BOOL)),
                }
            }
            MathLang::Num(v) => {
                let dtype = dtype_for_restored_leaf(type_map, idx, enode)?;
                let value = v.into_inner();
                let const_node = match dtype.kind {
                    DTypeKind::Float(_) => Const::Float(value),
                    // Complex constants should be represented via MathLang::Complex,
                    // but fall back to a purely real complex here if it ever happens.
                    DTypeKind::Complex(_) => Const::Complex(value, 0.0),
                    DTypeKind::Bool => Const::Bool(value != 0.0),
                    DTypeKind::UInt | DTypeKind::SInt => Const::Int(value as i64),
                };
                Some((Node::Const(const_node), dtype))
            }
            MathLang::Complex([re, im]) => {
                let re = expect_num(expr, *re, idx, enode)?;
                let im = expect_num(expr, *im, idx, enode)?;
                let dtype = dtype_for_restored_leaf(type_map, idx, enode)?;
                if !matches!(dtype.kind, DTypeKind::Complex(_)) {
                    return Err(invalid_node(idx, enode));
                }
                Some((Node::Const(Const::Complex(re, im)), dtype))
            }
            MathLang::Param([pid]) => {
                let raw = expect_int(expr, *pid, idx, enode)?;
                if raw < 0 {
                    return Err(invalid_node(idx, enode));
                }
                let param_idx = raw as usize;
                let dtype = params
                    .get(param_idx)
                    .map(|(_, dt)| *dt)
                    .ok_or_else(|| invalid_node(idx, enode))?;
                Some((Node::Param(ParamId(param_idx)), dtype))
            }
            MathLang::Cast([tok, child]) => {
                let dtype = expect_token(expr, *tok, idx, enode)?;
                let child = expect_node_id(&id_map, *child, idx, enode)?;
                Some((Node::Cast(dtype, child, false), dtype))
            }
            MathLang::Neg([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Unary(UnaryOp::Neg, a), dtypes[a.0]))
            }
            MathLang::Not([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Unary(UnaryOp::Not, a), DType::BOOL))
            }
            MathLang::Add([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::Add, a, b),
                    try_promote(
                        &mut diags,
                        "restore:Add",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Sub([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::Sub, a, b),
                    try_promote(
                        &mut diags,
                        "restore:Sub",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Mul([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::Mul, a, b),
                    try_promote(
                        &mut diags,
                        "restore:Mul",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Div([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::Div, a, b),
                    try_promote(
                        &mut diags,
                        "restore:Div",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::FloorDiv([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::FloorDiv, a, b),
                    try_promote(
                        &mut diags,
                        "restore:FloorDiv",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Rem([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::Rem, a, b),
                    try_promote(
                        &mut diags,
                        "restore:Rem",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Pow([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((
                    Node::Binary(BinaryOp::Pow, a, b),
                    try_promote(
                        &mut diags,
                        "restore:Pow",
                        dtypes[a.0],
                        dtypes[b.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Eq([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Eq, a, b), DType::BOOL))
            }
            MathLang::Ne([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Ne, a, b), DType::BOOL))
            }
            MathLang::Lt([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Lt, a, b), DType::BOOL))
            }
            MathLang::Le([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Le, a, b), DType::BOOL))
            }
            MathLang::Gt([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Gt, a, b), DType::BOOL))
            }
            MathLang::Ge([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Ge, a, b), DType::BOOL))
            }
            MathLang::And([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::And, a, b), DType::BOOL))
            }
            MathLang::Or([a, b]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                let b = expect_node_id(&id_map, *b, idx, enode)?;
                Some((Node::Binary(BinaryOp::Or, a, b), DType::BOOL))
            }
            MathLang::IfElse([c, t, e]) => {
                let c = expect_node_id(&id_map, *c, idx, enode)?;
                let t = expect_node_id(&id_map, *t, idx, enode)?;
                let e = expect_node_id(&id_map, *e, idx, enode)?;
                Some((
                    Node::IfElse(c, t, e),
                    try_promote(
                        &mut diags,
                        "restore:IfElse",
                        dtypes[t.0],
                        dtypes[e.0],
                        idx,
                        enode,
                    )?,
                ))
            }
            MathLang::Exp([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Exp, vec![a]), dtypes[a.0]))
            }
            MathLang::Log([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Log, vec![a]), dtypes[a.0]))
            }
            MathLang::Sqrt([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Sqrt, vec![a]), dtypes[a.0]))
            }
            MathLang::Abs([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Abs, vec![a]), dtypes[a.0]))
            }
            MathLang::Sin([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Sin, vec![a]), dtypes[a.0]))
            }
            MathLang::Cos([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Cos, vec![a]), dtypes[a.0]))
            }
            MathLang::Tanh([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Tanh, vec![a]), dtypes[a.0]))
            }
            MathLang::Sigmoid([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Sigmoid, vec![a]), dtypes[a.0]))
            }
            MathLang::Relu([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Relu, vec![a]), dtypes[a.0]))
            }
            MathLang::Softplus([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Softplus, vec![a]), dtypes[a.0]))
            }
            MathLang::Floor([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Floor, vec![a]), dtypes[a.0]))
            }
            MathLang::Ceil([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::Ceil, vec![a]), dtypes[a.0]))
            }
            MathLang::FusedExpSigmoid([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::ExpSigmoid, vec![a]), dtypes[a.0]))
            }
            MathLang::FusedSigmoidLog([a]) => {
                let a = expect_node_id(&id_map, *a, idx, enode)?;
                Some((Node::Call(Builtin::SigmoidLog, vec![a]), dtypes[a.0]))
            }
        };

        if let Some((node, dtype)) = converted {
            let nid = nodes.len();
            nodes.push(node);
            dtypes.push(dtype);
            id_map[usize::from(current)] = Some(NodeId(nid));
        }
    }

    let root_expr_id = expr
        .as_ref()
        .len()
        .checked_sub(1)
        .map(Id::from)
        .ok_or_else(|| OptError::InvalidExpr {
            msg: "cannot restore empty expression".to_string(),
        })?;
    let root_node = expect_node_id(
        &id_map,
        root_expr_id,
        usize::from(root_expr_id),
        &expr[root_expr_id],
    )?;
    let params = params.to_vec();

    let dag =
        TypedDag::try_new(nodes, dtypes, params, vec![root_node]).map_err(|e: TypedDagError| {
            OptError::InvalidExpr {
                msg: format!("restored TypedDag invalid: {e}"),
            }
        })?;
    Ok((dag, diags))
}

pub fn restore_types(
    expr: &RecExpr<MathLang>,
    type_map: &TypeMap,
    params: &[(String, DType)],
) -> Result<TypedDag, OptError> {
    restore_types_with_info(expr, type_map, params).map(|(dag, _)| dag)
}

#[cfg(test)]
mod tests {
    use super::restore_types;
    use crate::opt::strip::strip_types;
    use crate::{Value, interp_eval, lower, parse_dsl};

    fn assert_values_close(lhs: &[Value], rhs: &[Value]) {
        fn as_f64(v: &Value) -> Option<f64> {
            match v {
                Value::F16(x) => Some(x.to_f64()),
                Value::BF16(x) => Some(x.to_f64()),
                Value::F32(x) => Some(*x as f64),
                Value::F64(x) => Some(*x),
                _ => None,
            }
        }

        assert_eq!(lhs.len(), rhs.len(), "output length mismatch");
        for (l, r) in lhs.iter().zip(rhs.iter()) {
            if let (Some(a), Some(b)) = (as_f64(l), as_f64(r)) {
                assert!((a - b).abs() <= 1e-12, "float mismatch: {a} vs {b}");
            } else {
                assert_eq!(l, r);
            }
        }
    }

    #[test]
    fn strip_restore_roundtrip_keeps_semantics() {
        let src = r#"
        |x: f32, y: i32|
        cast<f64>(x + y)
        "#;
        let func = parse_dsl(src).expect("parse should succeed");
        let dag = lower(&func).expect("lower should succeed").dag;

        let (expr, type_map) = strip_types(&dag).expect("strip should succeed");
        let restored =
            restore_types(&expr, &type_map, &dag.params).expect("restore should succeed");

        let inputs = [Value::F64(1.25), Value::SInt { bits: 64, value: 2 }];
        let base = interp_eval(&dag, &inputs).expect("baseline eval should succeed");
        let after = interp_eval(&restored, &inputs).expect("restored eval should succeed");
        assert_values_close(&base, &after);
        assert_eq!(restored.params, dag.params);
    }
}
