use crate::{
    compile::{
        block::{Compile, CompileContext, CompileValue},
        error::CompileError,
        syntax::{
            CallExpression, DefmacroExpression, DefunExpression, Expression, FunctionBody,
            LambdaExpression,
        },
    },
    vm::instruction::{ArgumentCount, Instruction},
};

fn builtin_identifier_callee(identifier: &str) -> Option<Instruction> {
    match identifier {
        "+" => Some(Instruction::Add),
        "-" => Some(Instruction::Sub),
        "*" | "·" => Some(Instruction::Mul),
        "/" | "÷" => Some(Instruction::Div),
        "%" => Some(Instruction::Mod),
        "=" => Some(Instruction::Eq),
        "/=" | "≠" => Some(Instruction::Ne),
        ">" => Some(Instruction::Gt),
        "<" => Some(Instruction::Lt),
        ">=" | "≥" => Some(Instruction::Ge),
        "<=" | "≤" => Some(Instruction::Le),
        "not" => Some(Instruction::Not),
        "and" | "&&" => Some(Instruction::And),
        "or" | "||" => Some(Instruction::Or),
        _ => None,
    }
}

impl Compile for FunctionBody {
    fn compile(&self, cx: &mut CompileContext) -> Result<CompileValue, CompileError> {
        for expression in self.head.iter() {
            cx.compile_statement(expression)?;
        }
        self.tail.compile(cx)
    }
}

impl Compile for CallExpression {
    fn compile(&self, cx: &mut CompileContext) -> Result<CompileValue, CompileError> {
        let argument_count = ArgumentCount::try_from(self.arguments.len())
            .map_err(|_| CompileError::TooManyArgumentsInCall)?;
        let instruction = if let Expression::Identifier(callee) = self.callee.as_ref() {
            if let Some(instruction) = builtin_identifier_callee(callee.as_ref()) {
                instruction
            } else {
                Instruction::Call
            }
        } else {
            Instruction::Call
        };
        if instruction == Instruction::Call {
            let callee = self.callee.compile(cx)?;
            cx.push(callee)?;
            cx.guard_argument();
        }
        for expression in self.arguments.iter() {
            let value = expression.compile(cx)?;
            cx.push(value)?;
            cx.guard_argument();
        }
        cx.emit(instruction);
        cx.emit(argument_count);
        let mut unguard_count = usize::from(argument_count);
        if instruction == Instruction::Call {
            unguard_count += 1;
        }
        for _ in 0..unguard_count {
            cx.unguard_argument();
        }
        Ok(CompileValue::Stack)
    }
}

impl Compile for LambdaExpression {
    fn compile(&self, cx: &mut CompileContext) -> Result<CompileValue, CompileError> {
        let function = cx.compile_function(
            Some("lambda".into()),
            self.docstring.clone(),
            &self.signature,
            &self.body,
        )?;
        Ok(CompileValue::LocalFunction(function))
    }
}

impl Compile for DefunExpression {
    fn compile(&self, cx: &mut CompileContext) -> Result<CompileValue, CompileError> {
        let function = cx.compile_function(
            Some(self.name.clone()),
            self.docstring.clone(),
            &self.signature,
            &self.body,
        )?;
        cx.compile_set_global(self.name.clone(), CompileValue::LocalFunction(function))?;
        Ok(CompileValue::Nil)
    }
}

impl Compile for DefmacroExpression {
    fn compile(&self, cx: &mut CompileContext) -> Result<CompileValue, CompileError> {
        let function = cx.compile_function(
            Some(self.name.clone()),
            self.docstring.clone(),
            &self.signature,
            &self.body,
        )?;
        cx.compile_declare_macro(self.name.clone(), CompileValue::LocalFunction(function))?;
        Ok(CompileValue::Nil)
    }
}

#[cfg(test)]
mod tests {
    use std::rc::Rc;

    use crate::{
        compile::{
            UpvalueDef,
            block::{CompileValue, test_compile},
            function::FunctionSignature,
            syntax::{
                Assignment, CallExpression, DefunExpression, Expression, FunctionBody,
                LambdaExpression, LetExpression,
            },
        },
        vm::instruction::{Instruction, LocalId},
    };

    #[test]
    fn test_lambda_capture_compile() {
        // (let (x 1)
        //      ((lambda (y) (+ x y)) 2))
        let e = Expression::Let(LetExpression {
            sequential: false,
            bindings: vec![Assignment {
                identifier: "x".into(),
                value: Rc::new(Expression::IntegerLiteral(1.into())),
            }],
            body: FunctionBody {
                head: vec![],
                tail: Rc::new(Expression::Call(CallExpression {
                    callee: Rc::new(Expression::Lambda(LambdaExpression {
                        docstring: None,
                        signature: FunctionSignature {
                            required_arguments: vec!["y".into()],
                            optional_arguments: vec![],
                            rest_argument: None,
                        },
                        body: FunctionBody {
                            head: vec![],
                            tail: Rc::new(Expression::Call(CallExpression {
                                callee: Rc::new(Expression::Identifier("+".into())),
                                arguments: vec![
                                    Rc::new(Expression::Identifier("x".into())),
                                    Rc::new(Expression::Identifier("y".into())),
                                ],
                            })),
                        },
                    })),
                    arguments: vec![Rc::new(Expression::IntegerLiteral(2.into()))],
                })),
            },
        });
        let (cx, _v) = test_compile(&e).unwrap();
        assert_eq!(cx.current, 0);
        assert_eq!(cx.function_blocks.len(), 2);
        let root_function = cx.function_blocks[0].to_bytecode().unwrap();
        let lambda_function = cx.function_blocks[1].to_bytecode().unwrap();

        // lambda
        assert_eq!(
            lambda_function.instructions.as_ref(),
            &[
                Instruction::GetUpvalue.into(),
                0,
                Instruction::GetLocal.into(),
                0,
                Instruction::Add.into(),
                2,
                Instruction::Return.into(),
            ]
        );
        assert_eq!(lambda_function.required_count, 1);
        assert_eq!(lambda_function.optional_count, 0);
        assert!(!lambda_function.has_rest);
        assert_eq!(
            lambda_function.upvalues.as_ref(),
            &[UpvalueDef {
                index: LocalId::from(0),
                is_local: true
            }]
        );

        // root
        assert_eq!(
            root_function.instructions.as_ref(),
            &[
                // (let ...
                // ... (x 1) ...
                Instruction::PushInteger.into(),
                1,
                0,
                // ... ((lambda (y) (+ x y)) 2) ...
                // (lambda ...)
                Instruction::PushConstant.into(),
                0,
                0,
                Instruction::MakeClosure.into(),
                // 2
                Instruction::PushInteger.into(),
                2,
                0,
                // ((lambda ...) 2)
                Instruction::Call.into(),
                1,
                // (let ... ) exit
                Instruction::SetTemp.into(),
                Instruction::CloseUpvalue.into(),
                Instruction::GetTemp.into(),
            ]
        );
    }

    #[test]
    fn test_lambda_compile() {
        let e = Expression::Lambda(LambdaExpression {
            docstring: None,
            signature: FunctionSignature {
                required_arguments: vec!["a".into()],
                optional_arguments: vec![],
                rest_argument: None,
            },
            body: FunctionBody {
                head: vec![],
                tail: Rc::new(Expression::IntegerLiteral(1234.into())),
            },
        });
        let (cx, v) = test_compile(&e).unwrap();
        assert_eq!(cx.current, 0);
        assert!(cx.function_blocks[0].instructions.is_empty());
        let CompileValue::LocalFunction(index) = v else {
            panic!("lambda did not compile to a function value")
        };
        let lambda_function = cx.function_blocks[index].to_bytecode().unwrap();
        assert_eq!(
            lambda_function.instructions.as_ref(),
            &[
                Instruction::PushInteger.into(),
                210,
                4,
                Instruction::Return.into(),
            ]
        );
    }

    #[test]
    fn test_compile_defun() {
        let e = Expression::Defun(DefunExpression {
            name: "my-function".into(),
            docstring: None,
            signature: FunctionSignature {
                required_arguments: vec!["a".into()],
                optional_arguments: vec![],
                rest_argument: None,
            },
            body: FunctionBody {
                head: vec![],
                tail: Rc::new(Expression::Identifier("a".into())),
            },
        });
        let (cx, _v) = test_compile(&e).unwrap();
        assert_eq!(cx.current, 0);
        assert_eq!(cx.function_blocks.len(), 2);
        let root_function = cx.function_blocks[0].to_bytecode().unwrap();
        let defun_function = cx.function_blocks[1].to_bytecode().unwrap();
        // identifier + function
        assert_eq!(root_function.constants.len(), 2);
        assert_eq!(
            root_function.instructions.as_ref(),
            &[
                // Identifier
                Instruction::PushConstant.into(),
                0,
                0,
                // Function
                Instruction::PushConstant.into(),
                1,
                0,
                Instruction::SetGlobal.into(),
            ]
        );
        // inner function
        assert!(defun_function.constants.is_empty());
        assert_eq!(defun_function.required_count, 1);
        assert_eq!(defun_function.optional_count, 0);
        assert!(!defun_function.has_rest);
        assert_eq!(
            defun_function.instructions.as_ref(),
            &[
                // a
                Instruction::GetLocal.into(),
                0,
                Instruction::Return.into(),
            ]
        );
    }
}
