use crate::diag::{Diag, DiagCode};
use crate::ir::{BinaryOp, Builtin, Const, DType, DTypeKind, Node, NodeId, TypedDag, UnaryOp};

#[derive(Debug, Clone)]
pub struct ConstFoldResult {
    pub dag: TypedDag,
    pub diags: Vec<Diag>,
}

fn const_as_bool(c: &Const) -> Option<bool> {
    match c {
        Const::Bool(v) => Some(*v),
        _ => None,
    }
}

fn const_as_f64(c: &Const) -> Option<f64> {
    match c {
        Const::Bool(v) => Some(if *v { 1.0 } else { 0.0 }),
        Const::Int(v) => Some(*v as f64),
        Const::Float(v) => Some(*v),
        Const::Complex(_, _) => None,
    }
}

fn const_as_complex(c: &Const) -> Option<(f64, f64)> {
    match c {
        Const::Bool(v) => Some((if *v { 1.0 } else { 0.0 }, 0.0)),
        Const::Int(v) => Some((*v as f64, 0.0)),
        Const::Float(v) => Some((*v, 0.0)),
        Const::Complex(re, im) => Some((*re, *im)),
    }
}

fn cast_const_to_dtype(c: &Const, dtype: DType) -> Option<Const> {
    match dtype.kind {
        DTypeKind::Bool => const_as_bool(c).map(Const::Bool),
        DTypeKind::UInt | DTypeKind::SInt => const_as_f64(c).map(|v| Const::Int(v as i64)),
        DTypeKind::Float(_) => const_as_f64(c).map(Const::Float),
        DTypeKind::Complex(_) => const_as_complex(c).map(|(re, im)| Const::Complex(re, im)),
    }
}

fn const_is_zero(c: &Const) -> bool {
    match c {
        Const::Bool(v) => !*v,
        Const::Int(v) => *v == 0,
        Const::Float(v) => *v == 0.0,
        Const::Complex(re, im) => *re == 0.0 && *im == 0.0,
    }
}

fn const_is_one(c: &Const) -> bool {
    match c {
        Const::Bool(v) => *v,
        Const::Int(v) => *v == 1,
        Const::Float(v) => *v == 1.0,
        Const::Complex(re, im) => *re == 1.0 && *im == 0.0,
    }
}

fn fold_unary(op: UnaryOp, x: &Const, dtype: DType) -> Option<Const> {
    let raw = match op {
        UnaryOp::Neg => match x {
            Const::Int(v) => Const::Int(-v),
            Const::Float(v) => Const::Float(-v),
            Const::Complex(re, im) => Const::Complex(-re, -im),
            Const::Bool(_) => return None,
        },
        UnaryOp::Not => Const::Bool(!const_as_bool(x)?),
    };
    cast_const_to_dtype(&raw, dtype)
}

fn int_div_floor(lhs: i64, rhs: i64) -> Option<i64> {
    if rhs == 0 {
        return None;
    }
    if lhs == i64::MIN && rhs == -1 {
        return None;
    }
    let q = lhs / rhs;
    let r = lhs % rhs;
    if r == 0 {
        Some(q)
    } else if (r > 0) != (rhs > 0) {
        Some(q - 1)
    } else {
        Some(q)
    }
}

fn int_rem_floor(lhs: i64, rhs: i64) -> Option<i64> {
    let q = int_div_floor(lhs, rhs)?;
    Some(lhs - q * rhs)
}

fn fold_binary(op: BinaryOp, lhs: &Const, rhs: &Const, dtype: DType) -> Option<Const> {
    let raw = match op {
        BinaryOp::And => Const::Bool(const_as_bool(lhs)? && const_as_bool(rhs)?),
        BinaryOp::Or => Const::Bool(const_as_bool(lhs)? || const_as_bool(rhs)?),
        BinaryOp::Eq => Const::Bool(const_as_f64(lhs)? == const_as_f64(rhs)?),
        BinaryOp::Ne => Const::Bool(const_as_f64(lhs)? != const_as_f64(rhs)?),
        BinaryOp::Lt => Const::Bool(const_as_f64(lhs)? < const_as_f64(rhs)?),
        BinaryOp::Le => Const::Bool(const_as_f64(lhs)? <= const_as_f64(rhs)?),
        BinaryOp::Gt => Const::Bool(const_as_f64(lhs)? > const_as_f64(rhs)?),
        BinaryOp::Ge => Const::Bool(const_as_f64(lhs)? >= const_as_f64(rhs)?),
        BinaryOp::Add => match (lhs, rhs) {
            (Const::Complex(lr, li), Const::Complex(rr, ri)) => Const::Complex(lr + rr, li + ri),
            _ => Const::Float(const_as_f64(lhs)? + const_as_f64(rhs)?),
        },
        BinaryOp::Sub => match (lhs, rhs) {
            (Const::Complex(lr, li), Const::Complex(rr, ri)) => Const::Complex(lr - rr, li - ri),
            _ => Const::Float(const_as_f64(lhs)? - const_as_f64(rhs)?),
        },
        BinaryOp::Mul => match (lhs, rhs) {
            (Const::Complex(lr, li), Const::Complex(rr, ri)) => {
                Const::Complex(lr * rr - li * ri, lr * ri + li * rr)
            }
            _ => Const::Float(const_as_f64(lhs)? * const_as_f64(rhs)?),
        },
        BinaryOp::Div => match (lhs, rhs) {
            (Const::Complex(lr, li), Const::Complex(rr, ri)) => {
                let denom = rr * rr + ri * ri;
                Const::Complex((lr * rr + li * ri) / denom, (li * rr - lr * ri) / denom)
            }
            _ => Const::Float(const_as_f64(lhs)? / const_as_f64(rhs)?),
        },
        BinaryOp::FloorDiv => match (lhs, rhs) {
            (Const::Int(l), Const::Int(r)) => Const::Int(int_div_floor(*l, *r)?),
            _ => return None,
        },
        BinaryOp::Rem => match (lhs, rhs) {
            (Const::Int(l), Const::Int(r)) => Const::Int(int_rem_floor(*l, *r)?),
            _ => return None,
        },
        BinaryOp::Pow => Const::Float(const_as_f64(lhs)?.powf(const_as_f64(rhs)?)),
    };
    cast_const_to_dtype(&raw, dtype)
}

fn fold_call(builtin: Builtin, x: &Const, dtype: DType) -> Option<Const> {
    let xf = const_as_f64(x)?;
    let raw = match builtin {
        Builtin::Exp => Const::Float(xf.exp()),
        Builtin::Log => Const::Float(xf.ln()),
        Builtin::Sqrt => Const::Float(xf.sqrt()),
        Builtin::Abs => Const::Float(xf.abs()),
        Builtin::Sin => Const::Float(xf.sin()),
        Builtin::Cos => Const::Float(xf.cos()),
        Builtin::Tanh => Const::Float(xf.tanh()),
        Builtin::Sigmoid => Const::Float(1.0 / (1.0 + (-xf).exp())),
        Builtin::Relu => Const::Float(if xf > 0.0 { xf } else { 0.0 }),
        Builtin::Softplus => Const::Float((1.0 + xf.exp()).ln()),
        Builtin::Floor => Const::Float(xf.floor()),
        Builtin::Ceil => Const::Float(xf.ceil()),
        Builtin::ExpSigmoid => Const::Float(1.0 / (1.0 + (-xf.exp()).exp())),
        Builtin::SigmoidLog => Const::Float((1.0 / (1.0 + (-xf).exp())).ln()),
    };
    cast_const_to_dtype(&raw, dtype)
}

fn const_of(nodes: &[Node], node_id: NodeId) -> Option<&Const> {
    match nodes.get(node_id.0) {
        Some(Node::Const(c)) => Some(c),
        _ => None,
    }
}

fn push_node(nodes: &mut Vec<Node>, dtypes: &mut Vec<DType>, node: Node, dtype: DType) -> NodeId {
    let nid = nodes.len();
    nodes.push(node);
    dtypes.push(dtype);
    NodeId(nid)
}

fn int_bounds(dtype: DType) -> Option<(i128, i128)> {
    match (dtype.kind, dtype.bits) {
        (DTypeKind::UInt, 8) => Some((0, u8::MAX as i128)),
        (DTypeKind::UInt, 16) => Some((0, u16::MAX as i128)),
        (DTypeKind::UInt, 32) => Some((0, u32::MAX as i128)),
        (DTypeKind::UInt, 64) => Some((0, u64::MAX as i128)),
        (DTypeKind::SInt, 8) => Some((i8::MIN as i128, i8::MAX as i128)),
        (DTypeKind::SInt, 16) => Some((i16::MIN as i128, i16::MAX as i128)),
        (DTypeKind::SInt, 32) => Some((i32::MIN as i128, i32::MAX as i128)),
        (DTypeKind::SInt, 64) => Some((i64::MIN as i128, i64::MAX as i128)),
        _ => None,
    }
}

fn infer_wider_int_dtype(value: i128, preferred_unsigned: bool) -> Option<DType> {
    if preferred_unsigned && value >= 0 {
        if value <= u8::MAX as i128 {
            Some(DType::U8)
        } else if value <= u16::MAX as i128 {
            Some(DType::U16)
        } else if value <= u32::MAX as i128 {
            Some(DType::U32)
        } else if value <= u64::MAX as i128 {
            Some(DType::U64)
        } else {
            None
        }
    } else if value >= i8::MIN as i128 && value <= i8::MAX as i128 {
        Some(DType::I8)
    } else if value >= i16::MIN as i128 && value <= i16::MAX as i128 {
        Some(DType::I16)
    } else if value >= i32::MIN as i128 && value <= i32::MAX as i128 {
        Some(DType::I32)
    } else if value >= i64::MIN as i128 && value <= i64::MAX as i128 {
        Some(DType::I64)
    } else {
        None
    }
}

fn maybe_fold_integer_with_overflow(
    op: BinaryOp,
    lhs: &Const,
    rhs: &Const,
    expected: DType,
    node_id: usize,
    diags: &mut Vec<Diag>,
) -> Option<(Const, DType)> {
    let lhs_i = match lhs {
        Const::Int(v) => *v as i128,
        _ => return None,
    };
    let rhs_i = match rhs {
        Const::Int(v) => *v as i128,
        _ => return None,
    };
    let result = match op {
        BinaryOp::Add => lhs_i + rhs_i,
        BinaryOp::Sub => lhs_i - rhs_i,
        BinaryOp::Mul => lhs_i * rhs_i,
        _ => return None,
    };

    let (min, max) = int_bounds(expected)?;
    if result >= min && result <= max {
        return Some((Const::Int(result as i64), expected));
    }

    let preferred_unsigned = matches!(expected.kind, DTypeKind::UInt);
    let widened = infer_wider_int_dtype(result, preferred_unsigned).unwrap_or(DType::I64);
    let op_sym = match op {
        BinaryOp::Add => "+",
        BinaryOp::Sub => "-",
        BinaryOp::Mul => "*",
        _ => "?",
    };
    diags.push(Diag::new(
        DiagCode::ConstFoldUnsafe,
        format!(
            "const_fold node#{node_id}: \
             integer overflow: \
             {lhs_i} {op_sym} {rhs_i} = {result}; \
             promoted {expected} -> {widened}; \
             write result directly `{result}` or wrapping result"
        ),
        None,
        None,
    ));

    if result >= i64::MIN as i128 && result <= i64::MAX as i128 {
        Some((Const::Int(result as i64), widened))
    } else {
        Some((Const::Float(result as f64), DType::F64))
    }
}

pub fn const_fold(dag: &TypedDag) -> ConstFoldResult {
    let mut nodes: Vec<Node> = Vec::with_capacity(dag.nodes.len());
    let mut dtypes: Vec<DType> = Vec::with_capacity(dag.dtypes.len());
    let mut id_map: Vec<NodeId> = vec![NodeId(0); dag.nodes.len()];
    let mut diags: Vec<Diag> = Vec::new();

    for (old_id, node) in dag.nodes.iter().enumerate() {
        let expected = dag.dtypes[old_id];
        let remap_id = |id: NodeId| id_map[id.0];

        let mapped_node = match node {
            Node::Param(pid) => Node::Param(*pid),
            Node::Const(c) => Node::Const(c.clone()),
            Node::Unary(op, a) => Node::Unary(*op, remap_id(*a)),
            Node::Binary(op, a, b) => Node::Binary(*op, remap_id(*a), remap_id(*b)),
            Node::IfElse(c, t, e) => Node::IfElse(remap_id(*c), remap_id(*t), remap_id(*e)),
            Node::Call(name, args) => {
                Node::Call(*name, args.iter().map(|id| remap_id(*id)).collect())
            }
            Node::Cast(dtype, child, explicit) => Node::Cast(*dtype, remap_id(*child), *explicit),
        };

        let mut alias_to: Option<NodeId> = None;
        let folded_const = match &mapped_node {
            Node::Unary(op, a) => const_of(&nodes, *a)
                .and_then(|x| fold_unary(*op, x, expected).map(|c| (c, expected))),
            Node::Binary(BinaryOp::Pow, a, b) => {
                let lhs = const_of(&nodes, *a);
                let rhs = const_of(&nodes, *b);
                if let Some(r) = rhs {
                    if const_is_one(r) {
                        alias_to = Some(*a);
                        None
                    } else if const_is_zero(r) {
                        cast_const_to_dtype(&Const::Float(1.0), expected).map(|c| (c, expected))
                    } else {
                        lhs.and_then(|l| {
                            fold_binary(BinaryOp::Pow, l, r, expected).map(|c| (c, expected))
                        })
                    }
                } else {
                    None
                }
            }
            Node::Binary(BinaryOp::And, a, b) => {
                let lhs = const_of(&nodes, *a);
                let rhs = const_of(&nodes, *b);
                match (lhs, rhs) {
                    (Some(Const::Bool(true)), _) => {
                        alias_to = Some(*b);
                        None
                    }
                    (_, Some(Const::Bool(true))) => {
                        alias_to = Some(*a);
                        None
                    }
                    (Some(l), Some(r)) => {
                        fold_binary(BinaryOp::And, l, r, expected).map(|c| (c, expected))
                    }
                    _ => None,
                }
            }
            Node::Binary(BinaryOp::Or, a, b) => {
                let lhs = const_of(&nodes, *a);
                let rhs = const_of(&nodes, *b);
                match (lhs, rhs) {
                    (Some(Const::Bool(false)), _) => {
                        alias_to = Some(*b);
                        None
                    }
                    (_, Some(Const::Bool(false))) => {
                        alias_to = Some(*a);
                        None
                    }
                    (Some(l), Some(r)) => {
                        fold_binary(BinaryOp::Or, l, r, expected).map(|c| (c, expected))
                    }
                    _ => None,
                }
            }
            Node::Binary(op, a, b) => {
                let lhs = const_of(&nodes, *a);
                let rhs = const_of(&nodes, *b);
                match (lhs, rhs) {
                    (Some(l), Some(r)) => {
                        if let Some((c, d)) = maybe_fold_integer_with_overflow(
                            *op, l, r, expected, old_id, &mut diags,
                        ) {
                            Some((c, d))
                        } else {
                            fold_binary(*op, l, r, expected).map(|c| (c, expected))
                        }
                    }
                    _ => None,
                }
            }
            Node::IfElse(c, t, e) => match const_of(&nodes, *c).and_then(const_as_bool) {
                Some(true) => {
                    alias_to = Some(*t);
                    None
                }
                Some(false) => {
                    alias_to = Some(*e);
                    None
                }
                None => None,
            },
            Node::Call(builtin, args) => {
                if args.len() == 1 {
                    if *builtin == Builtin::Exp {
                        if let Some(x) = const_of(&nodes, args[0]) {
                            if const_is_zero(x) {
                                cast_const_to_dtype(&Const::Float(1.0), expected)
                                    .map(|c| (c, expected))
                            } else {
                                fold_call(*builtin, x, expected).map(|c| (c, expected))
                            }
                        } else {
                            None
                        }
                    } else if *builtin == Builtin::Log {
                        if let Some(x) = const_of(&nodes, args[0]) {
                            if const_is_one(x) {
                                cast_const_to_dtype(&Const::Float(0.0), expected)
                                    .map(|c| (c, expected))
                            } else {
                                fold_call(*builtin, x, expected).map(|c| (c, expected))
                            }
                        } else {
                            None
                        }
                    } else {
                        const_of(&nodes, args[0])
                            .and_then(|x| fold_call(*builtin, x, expected).map(|c| (c, expected)))
                    }
                } else {
                    None
                }
            }
            Node::Cast(dtype, child, _) => const_of(&nodes, *child)
                .and_then(|x| cast_const_to_dtype(x, *dtype).map(|c| (c, *dtype))),
            _ => None,
        };

        let new_id = if let Some(alias) = alias_to {
            alias
        } else if let Some((c, dtype)) = folded_const {
            push_node(&mut nodes, &mut dtypes, Node::Const(c), dtype)
        } else {
            push_node(&mut nodes, &mut dtypes, mapped_node, expected)
        };
        id_map[old_id] = new_id;
    }

    let outputs = dag.outputs.iter().map(|id| id_map[id.0]).collect();
    let folded = TypedDag::try_new(nodes, dtypes, dag.params.clone(), outputs)
        .unwrap_or_else(|e| panic!("const_fold produced invalid TypedDag: {e}"));
    ConstFoldResult { dag: folded, diags }
}

#[cfg(test)]
mod tests {
    use super::const_fold;
    use crate::ir::DType;
    use crate::{lower, parse_dsl};

    #[test]
    fn const_fold_emits_integer_overflow_warning() {
        let func = parse_dsl("||\n255@u8 + 255@u8").expect("parse");
        let r = lower(&func).expect("lower");
        let folded = const_fold(&r.dag);
        assert!(
            folded
                .diags
                .iter()
                .any(|d| d.message.contains("integer overflow")),
            "const fold should emit overflow warning"
        );
        assert_eq!(folded.dag.dtypes[folded.dag.outputs[0].0], DType::U16);
    }

    #[test]
    fn const_fold_floor_div_and_rem_match_spec() {
        let func = parse_dsl("||\n-7@i32 // 3@i32").expect("parse");
        let r = lower(&func).expect("lower");
        let folded = const_fold(&r.dag);
        let out_dtype = folded.dag.dtypes[folded.dag.outputs[0].0];
        assert_eq!(out_dtype, DType::I32);

        let func = parse_dsl("||\n-7@i32 % 3@i32").expect("parse");
        let r = lower(&func).expect("lower");
        let folded = const_fold(&r.dag);
        let out_dtype = folded.dag.dtypes[folded.dag.outputs[0].0];
        assert_eq!(out_dtype, DType::I32);
    }

    #[test]
    fn const_fold_floor_and_ceil_preserve_float_dtype() {
        let func = parse_dsl("||\nfloor(1.9@f64)").expect("parse");
        let r = lower(&func).expect("lower");
        let folded = const_fold(&r.dag);
        assert_eq!(folded.dag.dtypes[folded.dag.outputs[0].0], DType::F64);

        let func = parse_dsl("||\nceil(1.1@f64)").expect("parse");
        let r = lower(&func).expect("lower");
        let folded = const_fold(&r.dag);
        assert_eq!(folded.dag.dtypes[folded.dag.outputs[0].0], DType::F64);
    }
}
