use crate::backend::Backend;
use crate::ir::{BinaryOp, Builtin, Const, DType, DTypeKind, Node, TypedDag, UnaryOp};

/// Python reserved words and invalid identifiers: use `p{idx}` so generated `def` lines stay valid.
fn py_param_name(name: &str, param_idx: usize) -> String {
    const KW: &[&str] = &[
        "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class",
        "continue", "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
        "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return",
        "try", "while", "with", "yield",
    ];
    let ok = !name.is_empty()
        && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
        && !name.chars().next().is_some_and(|c| c.is_ascii_digit())
        && !KW.contains(&name);
    if ok {
        name.to_string()
    } else {
        format!("p{param_idx}")
    }
}

fn py_cast_expr(src: DType, dst: DType, expr: &str) -> String {
    if src == dst {
        return expr.to_string();
    }

    match dst.kind {
        DTypeKind::Bool => format!("bool({expr})"),
        DTypeKind::UInt | DTypeKind::SInt => format!("int({expr})"),
        DTypeKind::Float(_) => format!("float({expr})"),
        DTypeKind::Complex(_) => match src.kind {
            DTypeKind::Complex(_) => expr.to_string(),
            _ => format!("complex({expr}, 0.0)"),
        },
    }
}

fn py_const(c: &Const) -> String {
    match c {
        Const::Bool(v) => {
            if *v {
                "True".to_string()
            } else {
                "False".to_string()
            }
        }
        Const::Int(v) => v.to_string(),
        Const::Float(v) => {
            if v.is_nan() {
                "float('nan')".to_string()
            } else if *v == f64::INFINITY {
                "float('inf')".to_string()
            } else if *v == f64::NEG_INFINITY {
                "float('-inf')".to_string()
            } else {
                format!("{v}")
            }
        }
        Const::Complex(re, im) => format!("complex({re}, {im})"),
    }
}

fn py_bin(op: BinaryOp) -> &'static str {
    match op {
        BinaryOp::Add => "+",
        BinaryOp::Sub => "-",
        BinaryOp::Mul => "*",
        BinaryOp::Div => "/",
        BinaryOp::FloorDiv => "//",
        BinaryOp::Rem => "%",
        BinaryOp::Pow => "**",
        BinaryOp::Eq => "==",
        BinaryOp::Ne => "!=",
        BinaryOp::Lt => "<",
        BinaryOp::Le => "<=",
        BinaryOp::Gt => ">",
        BinaryOp::Ge => ">=",
        BinaryOp::And => "and",
        BinaryOp::Or => "or",
    }
}

/// Default backend that emits plain Python kernels.
pub struct PythonBackend;

impl Backend for PythonBackend {
    type Output = String;

    fn emit_kernel(&self, dag: &TypedDag, fn_name: &str) -> String {
        let params = dag
            .params
            .iter()
            .enumerate()
            .map(|(i, (name, _))| py_param_name(name, i))
            .collect::<Vec<_>>()
            .join(", ");

        let mut out = String::new();
        out.push_str("import math\n\n");
        out.push_str(&format!("def {fn_name}({params}):\n"));

        // Only emit nodes that are actually needed to compute the outputs.
        let mut used = vec![false; dag.nodes.len()];
        let mut stack = dag.outputs.clone();
        while let Some(id) = stack.pop() {
            if used[id.0] {
                continue;
            }
            used[id.0] = true;
            match &dag.nodes[id.0] {
                Node::Param(_) | Node::Const(_) => {}
                Node::Unary(_, a) => stack.push(*a),
                Node::Binary(_, a, b) => {
                    stack.push(*a);
                    stack.push(*b);
                }
                Node::IfElse(c, t, e) => {
                    stack.push(*c);
                    stack.push(*t);
                    stack.push(*e);
                }
                Node::Call(_, args) => {
                    for a in args {
                        stack.push(*a);
                    }
                }
                Node::Cast(_, a, _) => stack.push(*a),
            }
        }

        for (idx, node) in dag.nodes.iter().enumerate() {
            if !used[idx] {
                continue;
            }
            let line = match node {
                Node::Param(i) => {
                    let name = dag
                        .params
                        .get(i.0)
                        .map(|(n, _)| py_param_name(n, i.0))
                        .expect("invalid Param index in TypedDag");
                    format!("    t{idx} = {name}\n")
                }
                Node::Const(c) => format!("    t{idx} = {}\n", py_const(c)),
                Node::Unary(op, a) => match op {
                    UnaryOp::Neg => format!("    t{idx} = -t{}\n", a.0),
                    UnaryOp::Not => format!("    t{idx} = (not t{})\n", a.0),
                },
                Node::Binary(op, a, b) => {
                    let sym = py_bin(*op);
                    format!("    t{idx} = (t{} {sym} t{})\n", a.0, b.0)
                }
                Node::IfElse(c, t, e) => {
                    format!("    t{idx} = (t{} if t{} else t{})\n", t.0, c.0, e.0)
                }
                Node::Call(builtin, args) => {
                    let a0 = args
                        .first()
                        .copied()
                        .expect("Call node must have at least one argument")
                        .0;
                    let expr = match builtin {
                        Builtin::Exp => format!("math.exp(t{a0})"),
                        Builtin::Log => format!("math.log(t{a0})"),
                        Builtin::Sqrt => format!("math.sqrt(t{a0})"),
                        Builtin::Abs => format!("abs(t{a0})"),
                        Builtin::Sin => format!("math.sin(t{a0})"),
                        Builtin::Cos => format!("math.cos(t{a0})"),
                        Builtin::Tanh => format!("math.tanh(t{a0})"),
                        Builtin::Sigmoid => format!("(1.0 / (1.0 + math.exp(-t{a0})))"),
                        Builtin::Relu => format!("(t{a0} if t{a0} > 0 else 0.0)"),
                        Builtin::Softplus => format!("math.log(1.0 + math.exp(t{a0}))"),
                        Builtin::Floor => format!("float(math.floor(t{a0}))"),
                        Builtin::Ceil => format!("float(math.ceil(t{a0}))"),
                        Builtin::ExpSigmoid => {
                            format!("(1.0 / (1.0 + math.exp(-math.exp(t{a0}))))")
                        }
                        Builtin::SigmoidLog => {
                            format!("math.log(1.0 / (1.0 + math.exp(-t{a0})))")
                        }
                    };
                    format!("    t{idx} = {expr}\n")
                }
                Node::Cast(dtype, a, _) => {
                    let src = dag.dtypes[a.0];
                    let expr = py_cast_expr(src, *dtype, &format!("t{}", a.0));
                    format!("    t{idx} = {expr}\n")
                }
            };
            out.push_str(&line);
        }

        if dag.outputs.len() == 1 {
            out.push_str(&format!("    return t{}\n", dag.outputs[0].0));
        } else {
            let refs = dag
                .outputs
                .iter()
                .map(|id| format!("t{}", id.0))
                .collect::<Vec<_>>()
                .join(", ");
            out.push_str(&format!("    return ({refs})\n"));
        }

        out
    }
}
