use crate::sexpr::{Atom, Expr};

const PREC_EQUAL: u8 = 10;
const PREC_ADD: u8 = 20;
const PREC_MUL: u8 = 30;
const PREC_POW: u8 = 40;
const PREC_ATOM: u8 = 50;

pub fn to_latex(expr: &Expr) -> String {
    render_expr(expr, 0)
}

fn render_expr(expr: &Expr, parent_prec: u8) -> String {
    let (rendered, prec) = render_expr_inner(expr);
    if prec < parent_prec {
        format!("\\left({}\\right)", rendered)
    } else {
        rendered
    }
}

fn render_expr_inner(expr: &Expr) -> (String, u8) {
    match expr {
        Expr::Atom(atom) => (render_atom(atom), PREC_ATOM),
        Expr::OrderedList(items) => {
            let pieces: Vec<String> = items.iter().map(|item| render_expr(item, 0)).collect();
            (format!("\\left[{}\\right]", pieces.join(", ")), PREC_ATOM)
        }
        Expr::Application(items) => render_application(items),
    }
}

fn render_application(items: &[Expr]) -> (String, u8) {
    let Some(Expr::Atom(Atom::Builtin(head))) = items.first() else {
        let pieces: Vec<String> = items.iter().map(|item| render_expr(item, 0)).collect();
        return (format!("\\left({}\\right)", pieces.join(" ")), PREC_ATOM);
    };

    let args = &items[1..];

    match head.as_str() {
        "+" => {
            let pieces: Vec<String> = args.iter().map(|arg| render_expr(arg, PREC_ADD)).collect();
            (pieces.join(" + "), PREC_ADD)
        }
        "*" => {
            let pieces: Vec<String> = args.iter().map(|arg| render_expr(arg, PREC_MUL)).collect();
            (pieces.join(" \\cdot "), PREC_MUL)
        }
        "=" => {
            let pieces: Vec<String> = args.iter().map(|arg| render_expr(arg, PREC_EQUAL)).collect();
            (pieces.join(" = "), PREC_EQUAL)
        }
        "sin" => unary_func("\\sin", args),
        "cos" => unary_func("\\cos", args),
        "d/dx" => {
            if args.len() == 1 {
                (
                    format!(
                        "\\frac{{d}}{{dx}}\\left({}\\right)",
                        render_expr(&args[0], 0)
                    ),
                    PREC_ATOM,
                )
            } else {
                fallback_call(head, args)
            }
        }
        _ => {
            if let Some(power) = power_from_head(head) {
                if args.len() == 1 {
                    let base = render_expr(&args[0], PREC_POW);
                    (format!("{}^{{{}}}", base, power), PREC_POW)
                } else {
                    fallback_call(head, args)
                }
            } else {
                fallback_call(head, args)
            }
        }
    }
}

fn unary_func(name: &str, args: &[Expr]) -> (String, u8) {
    if args.len() == 1 {
        (format!("{}\\left({}\\right)", name, render_expr(&args[0], 0)), PREC_ATOM)
    } else {
        fallback_call(name, args)
    }
}

fn fallback_call(head: &str, args: &[Expr]) -> (String, u8) {
    let pieces: Vec<String> = args.iter().map(|arg| render_expr(arg, 0)).collect();
    (
        format!(
            "{}\\left({}\\right)",
            latex_symbol(head),
            pieces.join(", ")
        ),
        PREC_ATOM,
    )
}

fn power_from_head(head: &str) -> Option<&str> {
    head.strip_prefix('^')
        .filter(|suffix| !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()))
}

fn render_atom(atom: &Atom) -> String {
    match atom {
        Atom::Number(n) => n.to_string(),
        Atom::Builtin(name) => latex_symbol(name),
        Atom::Variable(var) => latex_text(&format!("{}{}", var.r#type.r#type_prefix(), var.index)),
    }
}

fn latex_symbol(input: &str) -> String {
    input.replace('_', "\\_")
}

fn latex_text(input: &str) -> String {
    let escaped = input
        .replace('\\', "\\textbackslash{}")
        .replace('{', "\\{")
        .replace('}', "\\}")
        .replace('$', "\\$")
        .replace('_', "\\_")
        .replace('&', "\\&")
        .replace('#', "\\#")
        .replace('%', "\\%")
        .replace('^', "\\^")
        .replace('~', "\\~{}");
    format!("\\text{{{}}}", escaped)
}

trait VariablePrefix {
    fn r#type_prefix(&self) -> &'static str;
}

impl VariablePrefix for crate::sexpr::VariableType {
    fn r#type_prefix(&self) -> &'static str {
        match self {
            crate::sexpr::VariableType::Number => "$",
            crate::sexpr::VariableType::Expr => "@",
            crate::sexpr::VariableType::NonNumberExpr => "!",
            crate::sexpr::VariableType::Ellipsis => "..",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sexpr::{Number, Variable, VariableType};

    fn num(n: f64) -> Expr {
        Expr::Atom(Atom::Number(Number(n)))
    }

    fn sym(s: &str) -> Expr {
        Expr::Atom(Atom::Builtin(s.to_string()))
    }

    fn app(items: Vec<Expr>) -> Expr {
        Expr::Application(items)
    }

    #[test]
    fn renders_basic_arithmetic() {
        let expr = app(vec![sym("+"), num(2.0), app(vec![sym("*"), num(3.0), sym("x")])]);
        assert_eq!(to_latex(&expr), "2 + 3 \\cdot x");
    }

    #[test]
    fn renders_derivative_and_power() {
        let expr = app(vec![sym("d/dx"), app(vec![sym("^2"), sym("x")])]);
        assert_eq!(to_latex(&expr), "\\frac{d}{dx}\\left(x^{2}\\right)");
    }

    #[test]
    fn renders_ordered_list() {
        let expr = Expr::OrderedList(vec![sym("x"), num(1.0)]);
        assert_eq!(to_latex(&expr), "\\left[x, 1\\right]");
    }

    #[test]
    fn renders_variable_atom() {
        let var = Expr::Atom(Atom::Variable(Variable {
            r#type: VariableType::NonNumberExpr,
            index: "a".to_string(),
        }));
        assert_eq!(to_latex(&var), "\\text{!a}");
    }
}