use std::collections::HashMap;

use cranelift::{
    codegen::{
        Context,
        ir::{BlockArg, immediates::Offset32},
    },
    prelude::*,
};
use cranelift_module::{Linkage, Module, default_libcall_names};
use cranelift_object::{ObjectBuilder, ObjectModule};
use ecow::{EcoString, eco_format};
use itertools::Itertools as _;

use crate::cranelift::mir::ast;

pub mod mir;

struct CodeGen<M: Module> {
    module: M,
    context: Context,
    function_context: FunctionBuilderContext,
}

pub fn compile(module: ast::Module<ast::CompleteType>) -> Vec<u8> {
    let mut flag_builder = settings::builder();

    flag_builder
        .set("is_pic", "true")
        .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 cranelift = {
        let module = ObjectModule::new(
            ObjectBuilder::new(isa, module.name.to_string(), default_libcall_names())
                .expect("unable to build object builder"),
        );

        CodeGen {
            function_context: FunctionBuilderContext::new(),
            context: module.make_context(),
            module,
        }
    };

    for function in module.functions {
        FunctionTranslator::new(&mut cranelift).translate(module.name.clone(), function);
    }

    let product = cranelift.module.finish();

    product.emit().expect("unable to emit finished product")
}

struct FunctionTranslator<'m, M: Module> {
    cg: &'m mut CodeGen<M>,
}

fn mangle_type_name(type_: &ast::CompleteType) -> EcoString {
    match type_ {
        ast::CompleteType::Int => "I".into(),
        ast::CompleteType::Float => "F".into(),
        ast::CompleteType::Bool => "B".into(),
        ast::CompleteType::String => "S".into(),
        ast::CompleteType::Func { arguments, returns } => {
            let arguments = arguments
                .iter()
                .map(mangle_type_name)
                .map(|s| eco_format!("{}{}", s.len(), s))
                .join("");

            let returns = mangle_type_name(returns);

            eco_format!("F{}{}R{}", arguments.len(), arguments, returns)
        }
        ast::CompleteType::Struct { elements } => {
            let elements = elements
                .iter()
                .map(mangle_type_name)
                .map(|s| eco_format!("{}{}", s.len(), s))
                .join("");

            eco_format!("C{}{}", elements.len(), elements)
        }
        ast::CompleteType::List(element_type) => {
            let element = mangle_type_name(element_type);

            eco_format!("L{}{}", element.len(), element)
        }
    }
}

fn mangled_function_name(
    module: EcoString,
    name: EcoString,
    arguments: &[ast::CompleteType],
    returns: &ast::CompleteType,
) -> EcoString {
    let arguments = arguments.iter().map(mangle_type_name).join(",");
    let returns = mangle_type_name(returns);

    eco_format!(
        "M{}{}_{}{}_A{}{}_R{}",
        module.len(),
        module.replace("/", "_"),
        name.len(),
        name,
        arguments.len(),
        arguments,
        returns,
    )
}

impl<'m, M: Module> FunctionTranslator<'m, M> {
    fn new(cg: &'m mut CodeGen<M>) -> Self {
        Self { cg }
    }

    fn translate(self, module: EcoString, function: ast::Function<ast::CompleteType>) {
        let arguments: Vec<_> = function
            .parameters
            .iter()
            .map(|parameter| parameter.type_.clone())
            .collect();

        let name = mangled_function_name(
            module,
            function.name.clone(),
            &arguments,
            &function.return_type,
        );

        self.cg.context.func.signature = self.signature_from_function(&function);

        let mut builder =
            FunctionBuilder::new(&mut self.cg.context.func, &mut self.cg.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 value = {
            let mut translator = ExpressionTranslator::new(&mut self.cg.module, &mut builder);

            for (index, parameter) in function.parameters.iter().enumerate() {
                if let Some(name) = &parameter.name {
                    let var_type = translate_type(&parameter.type_);
                    let var = translator.builder.declare_var(var_type);
                    _ = translator.vars.insert(name.clone(), var);

                    let val = translator.builder.block_params(entry)[index];
                    translator.builder.def_var(var, val);
                }
            }

            translator.translate(function.body)
        };

        _ = builder.ins().return_(&[value]);

        builder.finalize();

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

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

        self.cg.module.clear_context(&mut self.cg.context);
    }

    fn signature_from_function(&self, function: &ast::Function<ast::CompleteType>) -> Signature {
        let mut sig = self.cg.module.make_signature();

        for parameter in &function.parameters {
            sig.params
                .push(AbiParam::new(translate_type(&parameter.type_)));
        }

        sig.returns
            .push(AbiParam::new(translate_type(&function.return_type)));

        return sig;
    }
}

fn translate_type(type_: &ast::CompleteType) -> Type {
    match type_ {
        ast::CompleteType::Int => types::I64,
        ast::CompleteType::Float => types::F64,
        ast::CompleteType::Bool => types::I64, // wasteful but whatever
        ast::CompleteType::String => types::I64,
        ast::CompleteType::Func {
            arguments: _,
            returns: _,
        } => types::I64,
        ast::CompleteType::Struct { elements: _ } => types::I64,
        ast::CompleteType::List(_complete_type) => types::I64,
    }
}

struct ExpressionTranslator<'m, 'b, M: Module> {
    module: &'m mut M,
    builder: &'m mut FunctionBuilder<'b>,
    vars: HashMap<EcoString, Variable>,
}

impl<'m, 'b, M: Module> ExpressionTranslator<'m, 'b, M> {
    fn new(module: &'m mut M, builder: &'m mut FunctionBuilder<'b>) -> Self {
        Self {
            module,
            builder,
            vars: Default::default(),
        }
    }

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

        match expression {
            ast::Expression::Block(expressions) => expressions
                .into_iter()
                .map(|expression| self.translate(expression))
                .last()
                .expect("block should be non-empty"),
            ast::Expression::FunctionRef {
                module,
                name,
                arity: _,
                type_,
            } => {
                let ast::CompleteType::Func { arguments, returns } = type_ else {
                    unreachable!("function ref should have func type")
                };

                let name = mangled_function_name(module, name, &arguments, &returns);
                let sig = {
                    let mut sig = self.module.make_signature();
                    for argument in arguments {
                        sig.params.push(AbiParam::new(translate_type(&argument)));
                    }
                    sig.returns.push(AbiParam::new(translate_type(&returns)));
                    sig
                };

                let id = self
                    .module
                    .declare_function(&name, Linkage::Import, &sig)
                    .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)
            }
            ast::Expression::Var(var) => {
                let var = self
                    .vars
                    .get(&var.name)
                    .cloned()
                    .expect("variable to be defined");

                self.builder.use_var(var)
            }
            ast::Expression::Int { value } => {
                let value = i64::try_from(value).expect("integer value too large");

                self.builder.ins().iconst(types::I64, value)
            }
            ast::Expression::Float { value } => {
                let value: f64 = value.parse().expect("invalid float value");

                self.builder.ins().f64const(value)
            }
            ast::Expression::Bool { value } => match value {
                true => self.builder.ins().iconst(types::I64, 1),
                false => self.builder.ins().iconst(types::I64, 0),
            },
            ast::Expression::String { value: _ } => todo!(),
            ast::Expression::Equals { lhs, rhs } => match lhs.type_() {
                ast::CompleteType::Int => {
                    let lhs = self.translate(*lhs);
                    let rhs = self.translate(*rhs);

                    self.builder.ins().icmp(IntCC::Equal, lhs, rhs)
                }
                ast::CompleteType::Float => todo!(),
                ast::CompleteType::Bool => todo!(),
                ast::CompleteType::String => todo!(),
                ast::CompleteType::Func {
                    arguments: _,
                    returns: _,
                } => todo!(),
                ast::CompleteType::Struct { elements: _ } => todo!(),
                ast::CompleteType::List(_complete_type) => todo!(),
            },
            ast::Expression::NotEquals { lhs: _, rhs: _ } => todo!(),
            ast::Expression::IntGt { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().icmp(IntCC::SignedGreaterThan, lhs, rhs)
            }
            ast::Expression::IntGtEq { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder
                    .ins()
                    .icmp(IntCC::SignedGreaterThanOrEqual, lhs, rhs)
            }
            ast::Expression::IntLt { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().icmp(IntCC::SignedLessThan, lhs, rhs)
            }
            ast::Expression::IntLtEq { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder
                    .ins()
                    .icmp(IntCC::SignedLessThanOrEqual, lhs, rhs)
            }
            ast::Expression::IntAdd { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

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

                self.builder.ins().isub(lhs, rhs)
            }
            ast::Expression::IntMul { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().imul(lhs, rhs)
            }
            ast::Expression::IntDiv { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().sdiv(lhs, rhs)
            }
            ast::Expression::IntRem { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().srem(lhs, rhs)
            }
            ast::Expression::FloatGt { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fcmp(FloatCC::GreaterThan, lhs, rhs)
            }
            ast::Expression::FloatGtEq { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder
                    .ins()
                    .fcmp(FloatCC::GreaterThanOrEqual, lhs, rhs)
            }
            ast::Expression::FloatLt { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fcmp(FloatCC::LessThan, lhs, rhs)
            }
            ast::Expression::FloatLtEq { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fcmp(FloatCC::LessThanOrEqual, lhs, rhs)
            }
            ast::Expression::FloatAdd { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fadd(lhs, rhs)
            }
            ast::Expression::FloatSub { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fsub(lhs, rhs)
            }
            ast::Expression::FloatMul { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fmul(lhs, rhs)
            }
            ast::Expression::FloatDiv { lhs, rhs } => {
                let lhs = self.translate(*lhs);
                let rhs = self.translate(*rhs);

                self.builder.ins().fdiv(lhs, rhs)
            }
            ast::Expression::StringConcat { lhs: _, rhs: _ } => todo!(),
            ast::Expression::List {
                items: _,
                tail: _,
                type_: _,
            } => todo!(),
            ast::Expression::Struct {
                tag,
                items,
                type_: _,
            } => {
                let mut sig = self.module.make_signature();
                sig.params.push(AbiParam::new(types::I64));
                sig.returns.push(AbiParam::new(types::I64));

                let id = self
                    .module
                    .declare_function("malloc", Linkage::Import, &sig)
                    .expect("failed to declare malloc");

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

                let size_arg = self
                    .builder
                    .ins()
                    .iconst(types::I64, (items.len() as i64) * 8 + 8);

                let call = self.builder.ins().call(malloc, &[size_arg]);
                let ptr = self.builder.inst_results(call)[0];

                let tag = i64::from(tag.unwrap_or_default());
                let tag = self.builder.ins().iconst(types::I64, tag);
                _ = self
                    .builder
                    .ins()
                    .store(MemFlags::new(), tag, ptr, Offset32::new(0));

                for (index, item) in items.into_iter().enumerate() {
                    let index = index as i32 + 1;

                    let value = self.translate(item);

                    _ = self.builder.ins().store(
                        MemFlags::new(),
                        value,
                        ptr,
                        Offset32::new(index * 8),
                    );
                }

                ptr
            }
            ast::Expression::StructTag { value } => {
                let ptr = self.translate(*value);

                self.builder
                    .ins()
                    .load(types::I64, MemFlags::new(), ptr, Offset32::new(0))
            }
            ast::Expression::StructAccess {
                value,
                index,
                type_,
            } => {
                let ptr = self.translate(*value);

                self.builder.ins().load(
                    translate_type(&type_),
                    MemFlags::new(),
                    ptr,
                    Offset32::new(index as i32 * 8 + 8),
                )
            }
            ast::Expression::Set { name, value } => {
                let var = self.builder.declare_var(translate_type(&value.type_()));
                _ = self.vars.insert(name.name, var);

                let val = self.translate(*value);
                self.builder.def_var(var, val);
                self.builder.use_var(var)
            }
            ast::Expression::If { cond, then, else_ } => {
                let cond = self.translate(*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, translate_type(&then.type_()));

                _ = 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(*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(*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]
            }
            ast::Expression::Call {
                target,
                args,
                type_,
            } => {
                let mut signature = self.module.make_signature();

                for arg in &args {
                    signature
                        .params
                        .push(AbiParam::new(translate_type(&arg.type_())));
                }
                signature
                    .returns
                    .push(AbiParam::new(translate_type(&type_)));

                let target = self.translate(*target);
                let args: Vec<_> = args.into_iter().map(|arg| self.translate(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]
            }
            ast::Expression::Panic { type_, .. } => match type_ {
                ast::CompleteType::Func { .. }
                | ast::CompleteType::Struct { .. }
                | ast::CompleteType::List(_)
                | ast::CompleteType::String
                | ast::CompleteType::Bool
                | ast::CompleteType::Int => self.builder.ins().iconst(types::I64, 0),

                ast::CompleteType::Float => self.builder.ins().f64const(0.0),
            },
        }
    }
}