use cord_parse::ast::*;
use crate::tree::SdfNode;
use std::collections::HashMap;

#[derive(Debug)]
pub struct LowerError {
    pub msg: String,
}

impl std::fmt::Display for LowerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SDF lowering error: {}", self.msg)
    }
}

impl std::error::Error for LowerError {}

#[derive(Debug, Clone)]
enum Value {
    Scalar(f64),
    Array(Vec<f64>),
}

type Env = HashMap<String, Value>;
type ModuleDefs<'a> = HashMap<String, &'a ModuleDef>;

fn env_get_scalar(env: &Env, name: &str) -> Option<f64> {
    match env.get(name)? {
        Value::Scalar(v) => Some(*v),
        _ => None,
    }
}

fn env_insert_scalar(env: &mut Env, name: String, val: f64) {
    env.insert(name, Value::Scalar(val));
}

/// Lower a parsed SCAD program into an SDF tree.
/// Multiple top-level statements become an implicit union.
pub fn lower_program(program: &Program) -> Result<SdfNode, LowerError> {
    let mut env = Env::new();
    let mut modules = ModuleDefs::new();
    collect_module_defs(&program.statements, &mut modules);
    let nodes = lower_statements(&program.statements, &mut env, &modules)?;

    match nodes.len() {
        0 => Err(LowerError { msg: "empty scene".into() }),
        1 => Ok(nodes.into_iter().next().unwrap()),
        _ => Ok(SdfNode::Union(nodes)),
    }
}

fn collect_module_defs<'a>(stmts: &'a [Statement], modules: &mut ModuleDefs<'a>) {
    for stmt in stmts {
        if let Statement::ModuleDef(def) = stmt {
            modules.insert(def.name.clone(), def);
        }
    }
}

fn lower_statements(stmts: &[Statement], env: &mut Env, modules: &ModuleDefs) -> Result<Vec<SdfNode>, LowerError> {
    let mut nodes = Vec::new();
    for stmt in stmts {
        if let Some(node) = lower_statement(stmt, env, modules)? {
            nodes.push(node);
        }
    }
    Ok(nodes)
}

fn lower_statement(stmt: &Statement, env: &mut Env, modules: &ModuleDefs) -> Result<Option<SdfNode>, LowerError> {
    match stmt {
        Statement::ModuleCall(call) => lower_module_call(call, env, modules).map(Some),
        Statement::BooleanOp(bop) => lower_boolean(bop, env, modules).map(Some),
        Statement::Assignment(asgn) => {
            if let Some(val) = eval_value(&asgn.value, env) {
                env.insert(asgn.name.clone(), val);
            }
            Ok(None)
        }
        Statement::ModuleDef(_) => Ok(None),
        Statement::ForLoop(fl) => lower_for_loop(fl, env, modules).map(Some),
        Statement::IfElse(ie) => lower_if_else(ie, env, modules),
    }
}

/// Unroll a for loop into a union of independent instances.
///
/// Each iteration gets its own copy of the environment with the loop
/// variable bound. Since iterations share no state, every unrolled
/// instance is an independent parallel branch.
fn lower_for_loop(fl: &ForLoop, env: &mut Env, modules: &ModuleDefs) -> Result<SdfNode, LowerError> {
    let values = expand_range(&fl.range, env)?;

    let mut branches = Vec::with_capacity(values.len());
    for val in values {
        let mut inner_env = env.clone();
        env_insert_scalar(&mut inner_env, fl.var.clone(), val);
        let nodes = lower_statements(&fl.body, &mut inner_env, modules)?;
        match nodes.len() {
            0 => {}
            1 => branches.push(nodes.into_iter().next().unwrap()),
            _ => branches.push(SdfNode::Union(nodes)),
        }
    }

    match branches.len() {
        0 => Err(LowerError { msg: "for loop produced no geometry".into() }),
        1 => Ok(branches.into_iter().next().unwrap()),
        _ => Ok(SdfNode::Union(branches)),
    }
}

/// Expand a for-loop range into concrete values.
fn expand_range(range: &ForRange, env: &Env) -> Result<Vec<f64>, LowerError> {
    match range {
        ForRange::Range { start, step, end } => {
            let s = eval_expr_env(start, env)
                .ok_or_else(|| LowerError { msg: "for range start must be constant".into() })?;
            let e = eval_expr_env(end, env)
                .ok_or_else(|| LowerError { msg: "for range end must be constant".into() })?;
            let step_val = step.as_ref()
                .and_then(|st| eval_expr_env(st, env))
                .unwrap_or(1.0);

            if step_val.abs() < 1e-15 {
                return Err(LowerError { msg: "for range step cannot be zero".into() });
            }

            let mut values = Vec::new();
            let mut v = s;
            if step_val > 0.0 {
                while v <= e + 1e-10 {
                    values.push(v);
                    v += step_val;
                }
            } else {
                while v >= e - 1e-10 {
                    values.push(v);
                    v += step_val;
                }
            }

            if values.len() > 10000 {
                return Err(LowerError { msg: "for loop exceeds 10000 iterations".into() });
            }

            Ok(values)
        }
        ForRange::List(exprs) => {
            exprs.iter()
                .map(|e| eval_expr_env(e, env)
                    .ok_or_else(|| LowerError { msg: "for list element must be constant".into() }))
                .collect()
        }
    }
}

/// Lower if/else by evaluating the condition at lowering time.
fn lower_if_else(ie: &IfElse, env: &mut Env, modules: &ModuleDefs) -> Result<Option<SdfNode>, LowerError> {
    if let Some(cond) = eval_expr_env(&ie.condition, env) {
        if cond != 0.0 {
            let nodes = lower_statements(&ie.then_body, env, modules)?;
            match nodes.len() {
                0 => Ok(None),
                1 => Ok(Some(nodes.into_iter().next().unwrap())),
                _ => Ok(Some(SdfNode::Union(nodes))),
            }
        } else if !ie.else_body.is_empty() {
            let nodes = lower_statements(&ie.else_body, env, modules)?;
            match nodes.len() {
                0 => Ok(None),
                1 => Ok(Some(nodes.into_iter().next().unwrap())),
                _ => Ok(Some(SdfNode::Union(nodes))),
            }
        } else {
            Ok(None)
        }
    } else {
        let mut nodes = lower_statements(&ie.then_body, env, modules)?;
        nodes.extend(lower_statements(&ie.else_body, env, modules)?);
        match nodes.len() {
            0 => Ok(None),
            1 => Ok(Some(nodes.into_iter().next().unwrap())),
            _ => Ok(Some(SdfNode::Union(nodes))),
        }
    }
}

/// Evaluate an expression to a full Value (scalar or array).
fn eval_value(expr: &Expr, env: &Env) -> Option<Value> {
    match expr {
        Expr::Ident(name) => env.get(name).cloned(),
        Expr::Vector(elems) => {
            let vals: Option<Vec<f64>> = elems.iter().map(|e| eval_expr_env(e, env)).collect();
            Some(Value::Array(vals?))
        }
        Expr::Index { array, index } => {
            let idx = eval_expr_env(index, env)? as usize;
            match eval_value(array, env)? {
                Value::Array(arr) => arr.get(idx).map(|v| Value::Scalar(*v)),
                _ => None,
            }
        }
        _ => eval_expr_env(expr, env).map(Value::Scalar),
    }
}

/// Evaluate a constant expression to a scalar f64.
fn eval_expr_env(expr: &Expr, env: &Env) -> Option<f64> {
    match expr {
        Expr::Number(n) => Some(*n),
        Expr::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
        Expr::Ident(name) => env_get_scalar(env, name),
        Expr::UnaryOp { op: UnaryOp::Neg, operand } => eval_expr_env(operand, env).map(|n| -n),
        Expr::UnaryOp { op: UnaryOp::Not, operand } => {
            eval_expr_env(operand, env).map(|n| if n == 0.0 { 1.0 } else { 0.0 })
        }
        Expr::BinaryOp { op, left, right } => {
            let l = eval_expr_env(left, env)?;
            let r = eval_expr_env(right, env)?;
            Some(match op {
                BinaryOp::Add => l + r,
                BinaryOp::Sub => l - r,
                BinaryOp::Mul => l * r,
                BinaryOp::Div => if r != 0.0 { l / r } else { return None },
                BinaryOp::Mod => if r != 0.0 { l % r } else { return None },
                BinaryOp::Lt => if l < r { 1.0 } else { 0.0 },
                BinaryOp::Le => if l <= r { 1.0 } else { 0.0 },
                BinaryOp::Gt => if l > r { 1.0 } else { 0.0 },
                BinaryOp::Ge => if l >= r { 1.0 } else { 0.0 },
                BinaryOp::Eq => if (l - r).abs() < 1e-10 { 1.0 } else { 0.0 },
                BinaryOp::Ne => if (l - r).abs() >= 1e-10 { 1.0 } else { 0.0 },
                BinaryOp::And => if l != 0.0 && r != 0.0 { 1.0 } else { 0.0 },
                BinaryOp::Or => if l != 0.0 || r != 0.0 { 1.0 } else { 0.0 },
            })
        }
        Expr::Ternary { cond, then_expr, else_expr } => {
            let c = eval_expr_env(cond, env)?;
            if c != 0.0 {
                eval_expr_env(then_expr, env)
            } else {
                eval_expr_env(else_expr, env)
            }
        }
        Expr::FnCall { name, args } => {
            match name.as_str() {
                "sin" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::sin),
                "cos" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::cos),
                "abs" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::abs),
                "sqrt" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::sqrt),
                "floor" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::floor),
                "ceil" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::ceil),
                "round" => args.first().and_then(|a| eval_expr_env(&a.value, env)).map(f64::round),
                "min" => {
                    let a = args.first().and_then(|a| eval_expr_env(&a.value, env))?;
                    let b = args.get(1).and_then(|a| eval_expr_env(&a.value, env))?;
                    Some(a.min(b))
                }
                "max" => {
                    let a = args.first().and_then(|a| eval_expr_env(&a.value, env))?;
                    let b = args.get(1).and_then(|a| eval_expr_env(&a.value, env))?;
                    Some(a.max(b))
                }
                _ => None,
            }
        }
        Expr::Index { array, index } => {
            let idx = eval_expr_env(index, env)? as usize;
            match eval_value(array, env)? {
                Value::Array(arr) => arr.get(idx).copied(),
                Value::Scalar(v) if idx == 0 => Some(v),
                _ => None,
            }
        }
        _ => None,
    }
}

fn lower_module_call(call: &ModuleCall, env: &mut Env, modules: &ModuleDefs) -> Result<SdfNode, LowerError> {
    let child_nodes = lower_statements(&call.children, env, modules)?;

    match call.name.as_str() {
        "sphere" => {
            let r = get_f64(&call.args, "r", 0, env)?.unwrap_or(1.0);
            Ok(SdfNode::Sphere { radius: r })
        }
        "cube" => {
            let size = get_cube_size(&call.args, env)?;
            let center = get_named_bool(&call.args, "center")?.unwrap_or(false);
            if center {
                Ok(SdfNode::Box { half_extents: [size[0] / 2.0, size[1] / 2.0, size[2] / 2.0] })
            } else {
                Ok(SdfNode::Translate {
                    offset: [size[0] / 2.0, size[1] / 2.0, size[2] / 2.0],
                    child: Box::new(SdfNode::Box {
                        half_extents: [size[0] / 2.0, size[1] / 2.0, size[2] / 2.0],
                    }),
                })
            }
        }
        "cylinder" => {
            let h = get_f64(&call.args, "h", 0, env)?.unwrap_or(1.0);
            let r = get_f64(&call.args, "r", 1, env)?.unwrap_or(1.0);
            let center = get_named_bool(&call.args, "center")?.unwrap_or(false);
            let node = SdfNode::Cylinder { radius: r, height: h };
            if center {
                Ok(node)
            } else {
                Ok(SdfNode::Translate {
                    offset: [0.0, 0.0, h / 2.0],
                    child: Box::new(node),
                })
            }
        }
        "translate" => {
            let v = get_vec3(&call.args, 0, env)?;
            let child = require_single_child(child_nodes, "translate")?;
            Ok(SdfNode::Translate { offset: v, child: Box::new(child) })
        }
        "rotate" => {
            let v = get_vec3(&call.args, 0, env)?;
            let child = require_single_child(child_nodes, "rotate")?;
            let mut node = child;
            if v[0] != 0.0 {
                node = SdfNode::Rotate { axis: [1.0, 0.0, 0.0], angle_deg: v[0], child: Box::new(node) };
            }
            if v[1] != 0.0 {
                node = SdfNode::Rotate { axis: [0.0, 1.0, 0.0], angle_deg: v[1], child: Box::new(node) };
            }
            if v[2] != 0.0 {
                node = SdfNode::Rotate { axis: [0.0, 0.0, 1.0], angle_deg: v[2], child: Box::new(node) };
            }
            Ok(node)
        }
        "scale" => {
            let child = require_single_child(child_nodes, "scale")?;
            if let Some(v) = try_get_vec3(&call.args, 0, env) {
                Ok(SdfNode::Scale { factor: v, child: Box::new(child) })
            } else {
                let s = get_f64(&call.args, "v", 0, env)?.unwrap_or(1.0);
                Ok(SdfNode::Scale { factor: [s, s, s], child: Box::new(child) })
            }
        }
        name => {
            if let Some(def) = modules.get(name) {
                let mut child_env = env.clone();
                for (i, param) in def.params.iter().enumerate() {
                    let val = call.args.get(i)
                        .and_then(|a| eval_value(&a.value, &child_env))
                        .or_else(|| param.default.as_ref().and_then(|d| eval_value(d, &child_env)));
                    if let Some(v) = val {
                        child_env.insert(param.name.clone(), v);
                    }
                }
                let nodes = lower_statements(&def.body, &mut child_env, modules)?;
                match nodes.len() {
                    0 => Err(LowerError { msg: format!("module {name}() produced no geometry") }),
                    1 => Ok(nodes.into_iter().next().unwrap()),
                    _ => Ok(SdfNode::Union(nodes)),
                }
            } else {
                Err(LowerError { msg: format!("unknown module: {name}") })
            }
        }
    }
}

fn lower_boolean(bop: &BooleanOp, env: &mut Env, modules: &ModuleDefs) -> Result<SdfNode, LowerError> {
    let children = lower_statements(&bop.children, env, modules)?;
    if children.is_empty() {
        return Err(LowerError { msg: "boolean operation with no children".into() });
    }

    match bop.op {
        BooleanKind::Union => Ok(SdfNode::Union(children)),
        BooleanKind::Intersection => Ok(SdfNode::Intersection(children)),
        BooleanKind::Difference => {
            let mut iter = children.into_iter();
            let base = iter.next().unwrap();
            let subtract: Vec<_> = iter.collect();
            if subtract.is_empty() {
                Ok(base)
            } else {
                Ok(SdfNode::Difference { base: Box::new(base), subtract })
            }
        }
    }
}

fn require_single_child(mut nodes: Vec<SdfNode>, op: &str) -> Result<SdfNode, LowerError> {
    match nodes.len() {
        0 => Err(LowerError { msg: format!("{op}() requires a child") }),
        1 => Ok(nodes.remove(0)),
        _ => Ok(SdfNode::Union(nodes)),
    }
}

fn eval_const_bool(expr: &Expr) -> Option<bool> {
    match expr {
        Expr::Bool(b) => Some(*b),
        _ => None,
    }
}

fn get_f64(args: &[Argument], name: &str, pos: usize, env: &Env) -> Result<Option<f64>, LowerError> {
    for arg in args {
        if arg.name.as_deref() == Some(name) {
            return eval_expr_env(&arg.value, env)
                .map(Some)
                .ok_or_else(|| LowerError { msg: format!("argument '{name}' must be a number") });
        }
    }
    if let Some(arg) = args.get(pos) {
        if arg.name.is_none() {
            return Ok(eval_expr_env(&arg.value, env));
        }
    }
    Ok(None)
}

fn get_named_bool(args: &[Argument], name: &str) -> Result<Option<bool>, LowerError> {
    for arg in args {
        if arg.name.as_deref() == Some(name) {
            return eval_const_bool(&arg.value)
                .map(Some)
                .ok_or_else(|| LowerError { msg: format!("argument '{name}' must be a boolean") });
        }
    }
    Ok(None)
}

fn get_vec3(args: &[Argument], pos: usize, env: &Env) -> Result<[f64; 3], LowerError> {
    try_get_vec3(args, pos, env)
        .ok_or_else(|| LowerError { msg: format!("expected [x,y,z] vector at position {pos}") })
}

fn try_get_vec3(args: &[Argument], pos: usize, env: &Env) -> Option<[f64; 3]> {
    let arg = args.get(pos)?;
    if arg.name.is_some() {
        return None;
    }
    if let Expr::Vector(elems) = &arg.value {
        if elems.len() >= 3 {
            let x = eval_expr_env(&elems[0], env)?;
            let y = eval_expr_env(&elems[1], env)?;
            let z = eval_expr_env(&elems[2], env)?;
            return Some([x, y, z]);
        }
    }
    None
}

fn get_cube_size(args: &[Argument], env: &Env) -> Result<[f64; 3], LowerError> {
    if let Some(arg) = args.first() {
        if arg.name.is_none() || arg.name.as_deref() == Some("size") {
            if let Expr::Vector(elems) = &arg.value {
                if elems.len() >= 3 {
                    if let (Some(x), Some(y), Some(z)) = (
                        eval_expr_env(&elems[0], env),
                        eval_expr_env(&elems[1], env),
                        eval_expr_env(&elems[2], env),
                    ) {
                        return Ok([x, y, z]);
                    }
                }
            }
            if let Some(s) = eval_expr_env(&arg.value, env) {
                return Ok([s, s, s]);
            }
        }
    }
    Ok([1.0, 1.0, 1.0])
}

#[cfg(test)]
mod tests {
    use super::*;
    use cord_parse::lexer::Lexer;
    use cord_parse::parser::Parser;

    fn parse_and_lower(src: &str) -> SdfNode {
        let tokens = Lexer::new(src).tokenize().unwrap();
        let program = Parser::new(tokens).parse_program().unwrap();
        lower_program(&program).unwrap()
    }

    fn count_union_children(node: &SdfNode) -> usize {
        match node {
            SdfNode::Union(children) => children.len(),
            _ => 1,
        }
    }

    #[test]
    fn basic_sphere() {
        let node = parse_and_lower("sphere(5);");
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 5.0).abs() < 1e-10));
    }

    #[test]
    fn variable_in_args() {
        let node = parse_and_lower("r = 3; sphere(r);");
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 3.0).abs() < 1e-10));
    }

    #[test]
    fn for_loop_unrolls() {
        let node = parse_and_lower(
            "for (i = [0:3]) translate([i*10, 0, 0]) sphere(1);"
        );
        assert_eq!(count_union_children(&node), 4);
    }

    #[test]
    fn for_loop_with_step() {
        let node = parse_and_lower(
            "for (i = [0:2:6]) sphere(i);"
        );
        assert_eq!(count_union_children(&node), 4);
    }

    #[test]
    fn for_loop_explicit_list() {
        let node = parse_and_lower(
            "for (x = [1, 5, 10]) translate([x, 0, 0]) sphere(1);"
        );
        assert_eq!(count_union_children(&node), 3);
    }

    #[test]
    fn if_constant_true() {
        let node = parse_and_lower("if (true) sphere(1);");
        assert!(matches!(node, SdfNode::Sphere { .. }));
    }

    #[test]
    fn if_constant_false_with_else() {
        let node = parse_and_lower("if (false) sphere(1); else cube(2);");
        assert!(matches!(node, SdfNode::Translate { .. }));
    }

    #[test]
    fn if_constant_false_no_else() {
        let src = "if (false) sphere(1);";
        let tokens = Lexer::new(src).tokenize().unwrap();
        let program = Parser::new(tokens).parse_program().unwrap();
        let result = lower_program(&program);
        assert!(result.is_err());
    }

    #[test]
    fn variable_condition_includes_both_branches() {
        let src = "x = 1; if (x) sphere(1); else cube(2);";
        let node = parse_and_lower(src);
        assert!(matches!(node, SdfNode::Sphere { .. }));
    }

    #[test]
    fn nested_for_loops() {
        let node = parse_and_lower(
            "for (i = [0:2]) for (j = [0:2]) translate([i*10, j*10, 0]) sphere(1);"
        );
        assert_eq!(count_union_children(&node), 3);
        if let SdfNode::Union(outer) = &node {
            for child in outer {
                assert_eq!(count_union_children(child), 3);
            }
        }
    }

    #[test]
    fn ternary_in_expression() {
        let node = parse_and_lower("x = 5; sphere(x > 3 ? 10 : 1);");
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 10.0).abs() < 1e-10));
    }

    #[test]
    fn expr_with_math_functions() {
        let node = parse_and_lower("sphere(sqrt(4));");
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 2.0).abs() < 1e-10));
    }

    #[test]
    fn for_with_variable_bounds() {
        let node = parse_and_lower("n = 3; for (i = [0:n]) sphere(i);");
        assert_eq!(count_union_children(&node), 4);
    }

    #[test]
    fn difference_op() {
        let node = parse_and_lower("difference() { cube(10, center=true); sphere(5); }");
        assert!(matches!(node, SdfNode::Difference { .. }));
    }

    #[test]
    fn array_subscript() {
        let node = parse_and_lower("arr = [1, 5, 10]; sphere(arr[1]);");
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 5.0).abs() < 1e-10));
    }

    #[test]
    fn array_subscript_in_for_loop() {
        let node = parse_and_lower(
            "vals = [2, 4, 6]; for (i = [0:2]) sphere(vals[i]);"
        );
        assert_eq!(count_union_children(&node), 3);
    }

    #[test]
    fn user_defined_module() {
        let node = parse_and_lower(
            "module ball() { sphere(5); } ball();"
        );
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 5.0).abs() < 1e-10));
    }

    #[test]
    fn user_module_with_env() {
        let node = parse_and_lower(
            "r = 3; module ball() { sphere(r); } ball();"
        );
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 3.0).abs() < 1e-10));
    }
}
