use std::collections::HashMap;

use cranelift::{codegen::ir::BlockArg, prelude::*};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{FuncOrDataId, Linkage, Module, default_libcall_names};
use ecow::EcoString;
use num_bigint::BigInt;

pub mod mir;

pub fn compile(m: mir::ast::Module) {
    let mut flag_builder = settings::builder();

    flag_builder
        .set("is_pic", "false")
        .expect("unable to set `is_pic` to false");

    let isa = cranelift_native::builder()
        .unwrap_or_else(|reason| panic!("host machine not support: {reason}"))
        .finish(settings::Flags::new(flag_builder))
        .expect("unable to build isa target");

    let mut module = JITModule::new(JITBuilder::with_isa(isa, default_libcall_names()));
    let mut context = module.make_context();
    let mut function_context = FunctionBuilderContext::new();

    for function in m.functions {
        let sig = &mut context.func.signature;

        for _ in &function.parameters {
            sig.params.push(AbiParam::new(types::I64));
        }

        sig.returns.push(AbiParam::new(types::I64));

        let mut builder = FunctionBuilder::new(&mut context.func, &mut function_context);

        let entry = builder.create_block();
        builder.append_block_params_for_function_params(entry);
        builder.switch_to_block(entry);
        builder.seal_block(entry);

        let mut translator = FunctionTranslator {
            builder,
            module: &mut module,
            variables: HashMap::new(),
        };

        for (index, parameter) in function.parameters.into_iter().enumerate() {
            if let Some(name) = parameter.name {
                let var = translator.make_var(name);
                let val = translator.builder.block_params(entry)[index];

                translator.builder.def_var(var, val);
            }
        }

        let output = translator.translate_expression(function.body);

        _ = translator.builder.ins().return_(&[output]);
        translator.builder.finalize();

        let id = module
            .declare_function(&function.name, Linkage::Export, &context.func.signature)
            .expect("unable to declare function");

        module
            .define_function(id, &mut context)
            .expect("unable to define function");

        module.clear_context(&mut context);
    }

    module
        .finalize_definitions()
        .expect("unable to finalize definitions");

    let FuncOrDataId::Func(main) = module.get_name("main").expect("expected a main function")
    else {
        panic!("expected main to be a function")
    };

    let code = module.get_finalized_function(main);
    #[allow(unsafe_code)]
    let main = unsafe { std::mem::transmute::<_, extern "C" fn() -> i64>(code) };

    println!("{}", main());
}

struct FunctionTranslator<'a, M: Module> {
    builder: FunctionBuilder<'a>,
    module: &'a mut M,
    variables: HashMap<EcoString, Variable>,
}

impl<'a, M: Module> FunctionTranslator<'a, M> {
    fn make_var(&mut self, name: EcoString) -> Variable {
        let variable = self.builder.declare_var(types::I64);
        _ = self.variables.insert(name, variable);
        variable
    }

    fn get_var(&mut self, name: &EcoString) -> Variable {
        self.variables
            .get(name)
            .cloned()
            .unwrap_or_else(|| panic!("variable '{}' not found", name))
    }

    fn make_int(&mut self, value: BigInt) -> Value {
        let value: i64 = value.try_into().expect("bigint out of representable range");

        self.builder.ins().iconst(types::I64, value)
    }

    fn translate_expression(&mut self, expression: mir::ast::Expression) -> Value {
        _ = self.module;

        match expression {
            mir::ast::Expression::Block(expressions) => expressions
                .into_iter()
                .map(|expression| self.translate_expression(expression))
                .last()
                .expect("block shouldn't be empty"),
            mir::ast::Expression::FunctionRef {
                module: _,
                name,
                arity,
            } => {
                let mut signature = self.module.make_signature();

                for _ in 0..arity {
                    signature.params.push(AbiParam::new(types::I64));
                }
                signature.returns.push(AbiParam::new(types::I64));

                let id = self
                    .module
                    .declare_function(&name, Linkage::Import, &signature)
                    .expect("unable to declare function for import");

                let func = self.module.declare_func_in_func(id, self.builder.func);

                self.builder.ins().func_addr(types::I64, func)
            }
            mir::ast::Expression::Var(var) => {
                let var = self.get_var(&var.name);

                self.builder.use_var(var)
            }
            mir::ast::Expression::Int { value } => self.make_int(value),
            mir::ast::Expression::Float { value: _ } => todo!(),
            mir::ast::Expression::Bool { value: _ } => todo!(),
            mir::ast::Expression::String { value: _ } => todo!(),
            mir::ast::Expression::Equals { lhs, rhs } => {
                let lhs = self.translate_expression(*lhs);
                let rhs = self.translate_expression(*rhs);

                self.builder.ins().icmp(IntCC::Equal, lhs, rhs)
            }
            mir::ast::Expression::NotEquals { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntGt { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntGtEq { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntLt { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntLtEq { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntAdd { lhs, rhs } => {
                let lhs = self.translate_expression(*lhs);
                let rhs = self.translate_expression(*rhs);

                self.builder.ins().iadd(lhs, rhs)
            }
            mir::ast::Expression::IntSub { lhs, rhs } => {
                let lhs = self.translate_expression(*lhs);
                let rhs = self.translate_expression(*rhs);

                self.builder.ins().isub(lhs, rhs)
            }
            mir::ast::Expression::IntMul { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntDiv { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::IntRem { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatGt { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatGtEq { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatLt { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatLtEq { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatAdd { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatSub { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatMul { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::FloatDiv { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::StringConcat { lhs: _, rhs: _ } => todo!(),
            mir::ast::Expression::List { items: _, tail: _ } => todo!(),
            mir::ast::Expression::Struct { tag: _, items: _ } => todo!(),
            mir::ast::Expression::StructAccess { value: _, index: _ } => todo!(),
            mir::ast::Expression::Set { name, value } => {
                let var = self.make_var(name.name);
                let val = self.translate_expression(*value);

                self.builder.def_var(var, val);
                self.builder.use_var(var)
            }
            mir::ast::Expression::If { cond, then, else_ } => {
                let cond = self.translate_expression(*cond);

                let then_block = self.builder.create_block();
                let else_block = self.builder.create_block();
                let merge_block = self.builder.create_block();

                _ = self.builder.append_block_param(merge_block, types::I64);

                _ = self
                    .builder
                    .ins()
                    .brif(cond, then_block, &[], else_block, &[]);

                self.builder.switch_to_block(then_block);
                self.builder.seal_block(then_block);

                let then_return = self.translate_expression(*then);
                _ = self
                    .builder
                    .ins()
                    .jump(merge_block, &[BlockArg::Value(then_return)]);

                self.builder.switch_to_block(else_block);
                self.builder.seal_block(else_block);

                let else_return = self.translate_expression(*else_);
                _ = self
                    .builder
                    .ins()
                    .jump(merge_block, &[BlockArg::Value(else_return)]);

                self.builder.switch_to_block(merge_block);
                self.builder.seal_block(merge_block);

                self.builder.block_params(merge_block)[0]
            }
            mir::ast::Expression::Call { target, args } => {
                let mut signature = self.module.make_signature();

                for _ in 0..args.len() {
                    signature.params.push(AbiParam::new(types::I64));
                }
                signature.returns.push(AbiParam::new(types::I64));

                let target = self.translate_expression(*target);
                let args: Vec<_> = args
                    .into_iter()
                    .map(|arg| self.translate_expression(arg))
                    .collect();

                let sig = self.builder.import_signature(signature);
                let call = self.builder.ins().call_indirect(sig, target, &args);

                self.builder.inst_results(call)[0]
            }
        }
    }
}