use std::path::Path;
use std::result;

use inkwell::basic_block::BasicBlock;
use inkwell::builder::Builder;
pub use inkwell::context::Context;
use inkwell::module::Module;
use inkwell::support::LLVMString;
use inkwell::types::FunctionType;
use inkwell::values::{BasicValueEnum, FunctionValue};
use inkwell::IntPredicate;
use thiserror::Error;

use crate::ast::{BinaryOperator, Decl, Expr, Fun, Ident, Literal, UnaryOperator};
use crate::common::env::Env;

#[derive(Debug, PartialEq, Eq, Error)]
pub enum Error {
    #[error("Undefined variable {0}")]
    UndefinedVariable(Ident<'static>),

    #[error("LLVM Error: {0}")]
    LLVMError(String),
}

impl From<LLVMString> for Error {
    fn from(s: LLVMString) -> Self {
        Self::LLVMError(s.to_string())
    }
}

pub type Result<T> = result::Result<T, Error>;

pub struct Codegen<'ctx, 'ast> {
    context: &'ctx Context,
    pub module: Module<'ctx>,
    builder: Builder<'ctx>,
    env: Env<'ast, BasicValueEnum<'ctx>>,
    function: Option<FunctionValue<'ctx>>,
}

impl<'ctx, 'ast> Codegen<'ctx, 'ast> {
    pub fn new(context: &'ctx Context, module_name: &str) -> Self {
        let module = context.create_module(module_name);
        let builder = context.create_builder();
        Self {
            context,
            module,
            builder,
            env: Default::default(),
            function: None,
        }
    }

    pub fn new_function<'a>(
        &'a mut self,
        name: &str,
        ty: FunctionType<'ctx>,
    ) -> &'a FunctionValue<'ctx> {
        self.function = Some(self.module.add_function(name, ty, None));
        let basic_block = self.append_basic_block("entry");
        self.builder.position_at_end(basic_block);
        self.function.as_ref().unwrap()
    }

    pub fn finish_function(&self, res: &BasicValueEnum<'ctx>) {
        self.builder.build_return(Some(res));
    }

    pub fn append_basic_block(&self, name: &str) -> BasicBlock<'ctx> {
        self.context
            .append_basic_block(self.function.unwrap(), name)
    }

    pub fn codegen_expr(&mut self, expr: &'ast Expr<'ast>) -> Result<BasicValueEnum<'ctx>> {
        match expr {
            Expr::Ident(id) => self
                .env
                .resolve(id)
                .cloned()
                .ok_or_else(|| Error::UndefinedVariable(id.to_owned())),
            Expr::Literal(Literal::Int(i)) => {
                let ty = self.context.i64_type();
                Ok(BasicValueEnum::IntValue(ty.const_int(*i, false)))
            }
            Expr::UnaryOp { op, rhs } => {
                let rhs = self.codegen_expr(rhs)?;
                match op {
                    UnaryOperator::Not => unimplemented!(),
                    UnaryOperator::Neg => Ok(BasicValueEnum::IntValue(
                        self.builder.build_int_neg(rhs.into_int_value(), "neg"),
                    )),
                }
            }
            Expr::BinaryOp { lhs, op, rhs } => {
                let lhs = self.codegen_expr(lhs)?;
                let rhs = self.codegen_expr(rhs)?;
                match op {
                    BinaryOperator::Add => {
                        Ok(BasicValueEnum::IntValue(self.builder.build_int_add(
                            lhs.into_int_value(),
                            rhs.into_int_value(),
                            "add",
                        )))
                    }
                    BinaryOperator::Sub => {
                        Ok(BasicValueEnum::IntValue(self.builder.build_int_sub(
                            lhs.into_int_value(),
                            rhs.into_int_value(),
                            "add",
                        )))
                    }
                    BinaryOperator::Mul => {
                        Ok(BasicValueEnum::IntValue(self.builder.build_int_sub(
                            lhs.into_int_value(),
                            rhs.into_int_value(),
                            "add",
                        )))
                    }
                    BinaryOperator::Div => {
                        Ok(BasicValueEnum::IntValue(self.builder.build_int_signed_div(
                            lhs.into_int_value(),
                            rhs.into_int_value(),
                            "add",
                        )))
                    }
                    BinaryOperator::Pow => unimplemented!(),
                    BinaryOperator::Equ => {
                        Ok(BasicValueEnum::IntValue(self.builder.build_int_compare(
                            IntPredicate::EQ,
                            lhs.into_int_value(),
                            rhs.into_int_value(),
                            "eq",
                        )))
                    }
                    BinaryOperator::Neq => todo!(),
                }
            }
            Expr::Let { bindings, body } => {
                self.env.push();
                for (id, val) in bindings {
                    let val = self.codegen_expr(val)?;
                    self.env.set(id, val);
                }
                let res = self.codegen_expr(body);
                self.env.pop();
                res
            }
            Expr::If {
                condition,
                then,
                else_,
            } => {
                let then_block = self.append_basic_block("then");
                let else_block = self.append_basic_block("else");
                let join_block = self.append_basic_block("join");
                let condition = self.codegen_expr(condition)?;
                self.builder.build_conditional_branch(
                    condition.into_int_value(),
                    then_block,
                    else_block,
                );
                self.builder.position_at_end(then_block);
                let then_res = self.codegen_expr(then)?;
                self.builder.build_unconditional_branch(join_block);

                self.builder.position_at_end(else_block);
                let else_res = self.codegen_expr(else_)?;
                self.builder.build_unconditional_branch(join_block);

                self.builder.position_at_end(join_block);
                let phi = self.builder.build_phi(self.context.i64_type(), "join");
                phi.add_incoming(&[(&then_res, then_block), (&else_res, else_block)]);
                Ok(phi.as_basic_value())
            }
        }
    }

    pub fn codegen_decl(&mut self, decl: &'ast Decl<'ast>) -> Result<()> {
        match decl {
            Decl::Fun(Fun { name, args, body }) => {
                let i64_type = self.context.i64_type();
                self.new_function(
                    name.into(),
                    i64_type.fn_type(
                        args.iter()
                            .map(|_| i64_type.into())
                            .collect::<Vec<_>>()
                            .as_slice(),
                        false,
                    ),
                );
                self.env.push();
                for (i, arg) in args.iter().enumerate() {
                    self.env
                        .set(arg, self.function.unwrap().get_nth_param(i as u32).unwrap());
                }
                let res = self.codegen_expr(body)?;
                self.env.pop();
                self.finish_function(&res);
                Ok(())
            }
        }
    }

    pub fn codegen_main(&mut self, expr: &'ast Expr<'ast>) -> Result<()> {
        self.new_function("main", self.context.i64_type().fn_type(&[], false));
        let res = self.codegen_expr(expr)?;
        self.finish_function(&res);
        Ok(())
    }

    pub fn print_to_file<P>(&self, path: P) -> Result<()>
    where
        P: AsRef<Path>,
    {
        Ok(self.module.print_to_file(path)?)
    }

    pub fn binary_to_file<P>(&self, path: P) -> Result<()>
    where
        P: AsRef<Path>,
    {
        if self.module.write_bitcode_to_path(path.as_ref()) {
            Ok(())
        } else {
            Err(Error::LLVMError(
                "Error writing bitcode to output path".to_owned(),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use inkwell::execution_engine::JitFunction;
    use inkwell::OptimizationLevel;

    use super::*;

    fn jit_eval<T>(expr: &str) -> anyhow::Result<T> {
        let expr = crate::parser::expr(expr).unwrap().1;

        let context = Context::create();
        let mut codegen = Codegen::new(&context, "test");
        let execution_engine = codegen
            .module
            .create_jit_execution_engine(OptimizationLevel::None)
            .unwrap();

        codegen.new_function("test", context.i64_type().fn_type(&[], false));
        let res = codegen.codegen_expr(&expr)?;
        codegen.finish_function(&res);

        unsafe {
            let fun: JitFunction<unsafe extern "C" fn() -> T> =
                execution_engine.get_function("test")?;
            Ok(fun.call())
        }
    }

    #[test]
    fn add_literals() {
        assert_eq!(jit_eval::<i64>("1 + 2").unwrap(), 3);
    }

    #[test]
    fn variable_shadowing() {
        assert_eq!(
            jit_eval::<i64>("let x = 1 in (let x = 2 in x) + x").unwrap(),
            3
        );
    }

    #[test]
    fn eq() {
        assert_eq!(
            jit_eval::<i64>("let x = 1 in if x == 1 then 2 else 4").unwrap(),
            2
        );
    }
}