//! This module defines the `JITCompile` trait and its implementations for
//! various IR types. It provides the translation from LIR to Cranelift IR.

use cranelift::codegen::ir::{self, StackSlot};
use cranelift::prelude::*;

use nixjit_eval::Value;
use nixjit_ir::*;
use nixjit_lir::Lir;

use super::{Context, JITContext};

/// A trait for compiling IR nodes to Cranelift IR.
///
/// This trait defines how different IR nodes should be compiled to
/// Cranelift IR instructions that can be executed by the JIT compiler.
pub trait JITCompile<Ctx: JITContext> {
    /// Compiles the IR node to Cranelift IR.
    ///
    /// # Arguments
    /// * `ctx` - The compilation context
    /// * `engine` - The evaluation context value
    /// * `env` - The environment value
    ///
    /// # Returns
    /// A stack slot containing the compiled result
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot;
}

impl<Ctx: JITContext> JITCompile<Ctx> for ExprId {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Lir {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for AttrSet {
    /// Compiles an attribute set to Cranelift IR.
    ///
    /// This creates a new attribute set and compiles all static attributes into it.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        let attrs = ctx.create_attrs();
        for (k, v) in self.stcs.iter() {
            let v = v.compile(ctx, engine, env);
            ctx.push_attr(attrs, k, v);
        }
        ctx.finalize_attrs(attrs)
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for List {
    /// Compiles a list to Cranelift IR.
    ///
    /// This creates a new list by compiling all items and storing them in an array.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        let array = ctx.alloc_array(self.items.len());
        for (i, item) in self.items.iter().enumerate() {
            let item = item.compile(ctx, engine, env);
            let tag = ctx.builder.ins().stack_load(types::I64, item, 0);
            let val0 = ctx.builder.ins().stack_load(types::I64, item, 8);
            let val1 = ctx.builder.ins().stack_load(types::I64, item, 16);
            let val2 = ctx.builder.ins().stack_load(types::I64, item, 24);
            ctx.builder
                .ins()
                .store(MemFlags::new(), tag, array, i as i32 * 32);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val0, array, i as i32 * 32 + 8);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val1, array, i as i32 * 32 + 16);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val2, array, i as i32 * 32 + 24);
        }
        ctx.create_list(array, self.items.len())
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for HasAttr {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for BinOp {
    /// Compiles a binary operation to Cranelift IR.
    ///
    /// This implementation handles various binary operations like addition, subtraction,
    /// division, logical AND/OR, and equality checks. It generates code that checks
    /// the types of operands and performs the appropriate operation.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        use BinOpKind::*;
        let lhs = self.lhs.compile(ctx, engine, env);
        let rhs = self.rhs.compile(ctx, engine, env);
        let lhs_tag = ctx.get_tag(lhs);
        let rhs_tag = ctx.get_tag(rhs);
        let eq = ctx.builder.ins().icmp(IntCC::Equal, lhs_tag, rhs_tag);

        let eq_block = ctx.builder.create_block();
        let neq_block = ctx.builder.create_block();
        let exit_block = ctx.builder.create_block();
        ctx.builder.ins().brif(eq, eq_block, [], neq_block, []);

        match self.kind {
            Add => {
                ctx.builder.switch_to_block(eq_block);
                let default_block = ctx.builder.create_block();
                let int_block = ctx.builder.create_block();
                let float_block = ctx.builder.create_block();
                let float_check_block = ctx.builder.create_block();

                let is_int = ctx
                    .builder
                    .ins()
                    .icmp_imm(IntCC::Equal, lhs_tag, Value::INT as i64);
                ctx.builder
                    .ins()
                    .brif(is_int, int_block, [], float_check_block, []);

                ctx.builder.switch_to_block(int_block);
                let lhs_value = ctx.get_small_value(types::I64, lhs);
                let rhs_value = ctx.get_small_value(types::I64, rhs);
                let result = ctx.builder.ins().iadd(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, &[]);

                // FIXME: Non-float
                ctx.builder.switch_to_block(float_check_block);
                let is_float =
                    ctx.builder
                        .ins()
                        .icmp_imm(IntCC::Equal, lhs_tag, Value::FLOAT as i64);
                ctx.builder
                    .ins()
                    .brif(is_float, float_block, [], default_block, []);

                ctx.builder.switch_to_block(float_block);
                let lhs_value = ctx.get_small_value(types::F64, lhs);
                let rhs_value = ctx.get_small_value(types::F64, rhs);
                let result = ctx.builder.ins().fadd(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, &[]);

                // FIXME: finish this
                ctx.builder.switch_to_block(default_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.switch_to_block(neq_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.seal_block(default_block);
                ctx.builder.seal_block(int_block);
                ctx.builder.seal_block(float_check_block);
                ctx.builder.seal_block(float_block);
            }
            Sub => {
                ctx.builder.switch_to_block(eq_block);
                let default_block = ctx.builder.create_block();
                let int_block = ctx.builder.create_block();
                let float_block = ctx.builder.create_block();
                let float_check_block = ctx.builder.create_block();

                let is_int = ctx
                    .builder
                    .ins()
                    .icmp_imm(IntCC::Equal, lhs_tag, Value::INT as i64);
                ctx.builder
                    .ins()
                    .brif(is_int, int_block, [], float_check_block, []);

                ctx.builder.switch_to_block(int_block);
                let lhs_value = ctx.get_small_value(types::I64, lhs);
                let rhs_value = ctx.get_small_value(types::I64, rhs);
                let result = ctx.builder.ins().isub(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, &[]);

                // FIXME: Non-float
                ctx.builder.switch_to_block(float_check_block);
                let is_float =
                    ctx.builder
                        .ins()
                        .icmp_imm(IntCC::Equal, lhs_tag, Value::FLOAT as i64);
                ctx.builder
                    .ins()
                    .brif(is_float, float_block, [], default_block, []);

                ctx.builder.switch_to_block(float_block);
                let lhs_value = ctx.get_small_value(types::F64, lhs);
                let rhs_value = ctx.get_small_value(types::F64, rhs);
                let result = ctx.builder.ins().fsub(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, &[]);

                // FIXME: finish this
                ctx.builder.switch_to_block(default_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.switch_to_block(neq_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.seal_block(default_block);
                ctx.builder.seal_block(int_block);
                ctx.builder.seal_block(float_check_block);
                ctx.builder.seal_block(float_block);
            }
            Div => {
                ctx.builder.switch_to_block(eq_block);
                let default_block = ctx.builder.create_block();
                let int_block = ctx.builder.create_block();
                let float_block = ctx.builder.create_block();
                let float_check_block = ctx.builder.create_block();

                let is_int = ctx
                    .builder
                    .ins()
                    .icmp_imm(IntCC::Equal, lhs_tag, Value::INT as i64);
                ctx.builder
                    .ins()
                    .brif(is_int, int_block, [], float_check_block, []);

                ctx.builder.switch_to_block(int_block);
                let lhs_value = ctx.get_small_value(types::I64, lhs);
                let rhs_value = ctx.get_small_value(types::I64, rhs);
                let result = ctx.builder.ins().sdiv(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, &[]);

                // FIXME: Non-float
                ctx.builder.switch_to_block(float_check_block);
                let is_float =
                    ctx.builder
                        .ins()
                        .icmp_imm(IntCC::Equal, lhs_tag, Value::FLOAT as i64);
                ctx.builder
                    .ins()
                    .brif(is_float, float_block, [], default_block, []);

                ctx.builder.switch_to_block(float_block);
                let lhs_value = ctx.get_small_value(types::F64, lhs);
                let rhs_value = ctx.get_small_value(types::F64, rhs);
                let result = ctx.builder.ins().fdiv(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, &[]);

                // FIXME: finish this
                ctx.builder.switch_to_block(default_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.switch_to_block(neq_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.seal_block(default_block);
                ctx.builder.seal_block(int_block);
                ctx.builder.seal_block(float_check_block);
                ctx.builder.seal_block(float_block);
            }
            And => {
                ctx.builder.switch_to_block(eq_block);
                let bool_block = ctx.builder.create_block();
                let non_bool_block = ctx.builder.create_block();

                let is_bool = ctx
                    .builder
                    .ins()
                    .icmp_imm(IntCC::Equal, lhs_tag, Value::BOOL as i64);
                ctx.builder
                    .ins()
                    .brif(is_bool, bool_block, [], non_bool_block, []);

                ctx.builder.switch_to_block(bool_block);
                let lhs_value = ctx.get_small_value(types::I64, lhs);
                let rhs_value = ctx.get_small_value(types::I64, rhs);
                let result = ctx.builder.ins().band(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, []);

                ctx.builder.switch_to_block(non_bool_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.switch_to_block(neq_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.seal_block(bool_block);
                ctx.builder.seal_block(non_bool_block);
            }
            Or => {
                ctx.builder.switch_to_block(eq_block);
                let bool_block = ctx.builder.create_block();
                let non_bool_block = ctx.builder.create_block();

                let is_bool = ctx
                    .builder
                    .ins()
                    .icmp_imm(IntCC::Equal, lhs_tag, Value::BOOL as i64);
                ctx.builder
                    .ins()
                    .brif(is_bool, bool_block, [], non_bool_block, []);

                ctx.builder.switch_to_block(bool_block);
                let lhs_value = ctx.get_small_value(types::I64, lhs);
                let rhs_value = ctx.get_small_value(types::I64, rhs);
                let result = ctx.builder.ins().bor(lhs_value, rhs_value);
                ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
                ctx.builder.ins().stack_store(result, lhs, 8);
                ctx.builder.ins().jump(exit_block, []);

                ctx.builder.switch_to_block(non_bool_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.switch_to_block(neq_block);
                ctx.builder.ins().trap(TrapCode::unwrap_user(1));

                ctx.builder.seal_block(bool_block);
                ctx.builder.seal_block(non_bool_block);
            }
            Eq => {
                ctx.builder.switch_to_block(eq_block);
                ctx.eq(lhs, rhs);
                ctx.builder.ins().jump(exit_block, []);
                ctx.builder.switch_to_block(neq_block);
                ctx.eq(lhs, rhs);
                ctx.builder.ins().jump(exit_block, []);
            }
            _ => todo!(),
        }

        ctx.builder.seal_block(exit_block);
        ctx.builder.seal_block(eq_block);
        ctx.builder.seal_block(neq_block);
        ctx.builder.switch_to_block(exit_block);
        ctx.free_slot(rhs);

        lhs
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for UnOp {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Attr {
    /// Compiles an attribute key to Cranelift IR.
    ///
    /// An attribute can be either a static string or a dynamic expression.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        use Attr::*;
        match self {
            Str(string) => ctx.create_string(string),
            Dynamic(ir) => ir.compile(ctx, engine, env),
        }
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Select {
    /// Compiles an attribute selection to Cranelift IR.
    ///
    /// This compiles the expression to select from, builds the attribute path,
    /// and calls the select helper function.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        let val = self.expr.compile(ctx, engine, env);
        let attrpath = ctx.alloc_array(self.attrpath.len());
        for (i, attr) in self.attrpath.iter().enumerate() {
            let arg = attr.compile(ctx, engine, env);
            let tag = ctx.builder.ins().stack_load(types::I64, arg, 0);
            let val0 = ctx.builder.ins().stack_load(types::I64, arg, 8);
            let val1 = ctx.builder.ins().stack_load(types::I64, arg, 16);
            let val2 = ctx.builder.ins().stack_load(types::I64, arg, 24);
            ctx.builder
                .ins()
                .store(MemFlags::new(), tag, attrpath, i as i32 * 32);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val0, attrpath, i as i32 * 32 + 8);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val1, attrpath, i as i32 * 32 + 16);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val2, attrpath, i as i32 * 32 + 24);
        }
        ctx.select(val, attrpath, self.attrpath.len(), engine, env);
        val
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for If {
    /// Compiles an if-expression to Cranelift IR.
    ///
    /// This generates code that evaluates the condition, checks that it's a boolean,
    /// and then jumps to the appropriate branch (true or false).
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        let cond = self.cond.compile(ctx, engine, env);
        let cond_type = ctx.builder.ins().stack_load(types::I64, cond, 0);
        let cond_value = ctx.builder.ins().stack_load(types::I64, cond, 8);

        let true_block = ctx.builder.create_block();
        let false_block = ctx.builder.create_block();
        let exit_block = ctx.builder.create_block();
        let error_block = ctx.builder.create_block();
        let judge_block = ctx.builder.create_block();
        let slot = ctx.alloca();

        let is_bool = ctx
            .builder
            .ins()
            .icmp_imm(IntCC::Equal, cond_type, Value::BOOL as i64);
        ctx.builder
            .ins()
            .brif(is_bool, judge_block, [], error_block, []);

        ctx.builder.switch_to_block(judge_block);
        ctx.builder
            .ins()
            .brif(cond_value, true_block, [], false_block, []);

        ctx.builder.switch_to_block(true_block);
        let ret = self.consq.compile(ctx, engine, env);
        let tag = ctx.builder.ins().stack_load(types::I64, ret, 0);
        let val0 = ctx.builder.ins().stack_load(types::I64, ret, 8);
        let val1 = ctx.builder.ins().stack_load(types::I64, ret, 16);
        let val2 = ctx.builder.ins().stack_load(types::I64, ret, 24);
        ctx.builder.ins().stack_store(tag, slot, 0);
        ctx.builder.ins().stack_store(val0, slot, 8);
        ctx.builder.ins().stack_store(val1, slot, 16);
        ctx.builder.ins().stack_store(val2, slot, 24);
        ctx.builder.ins().jump(exit_block, []);

        ctx.builder.switch_to_block(false_block);
        let ret = self.alter.compile(ctx, engine, env);
        let tag = ctx.builder.ins().stack_load(types::I64, ret, 0);
        let val0 = ctx.builder.ins().stack_load(types::I64, ret, 8);
        let val1 = ctx.builder.ins().stack_load(types::I64, ret, 16);
        let val2 = ctx.builder.ins().stack_load(types::I64, ret, 24);
        ctx.builder.ins().stack_store(tag, slot, 0);
        ctx.builder.ins().stack_store(val0, slot, 8);
        ctx.builder.ins().stack_store(val1, slot, 16);
        ctx.builder.ins().stack_store(val2, slot, 24);
        ctx.builder.ins().jump(exit_block, []);

        ctx.builder.switch_to_block(error_block);
        ctx.builder.ins().trap(TrapCode::unwrap_user(1));

        ctx.builder.switch_to_block(exit_block);
        slot
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Call {
    /// Compiles a function call to Cranelift IR.
    ///
    /// This compiles the function expression and all arguments, builds an argument array,
    /// and calls the call helper function.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        let func = self.func.compile(ctx, engine, env);
        let args = ctx.alloc_array(self.args.len());
        for (i, arg) in self.args.iter().enumerate() {
            let arg = arg.compile(ctx, engine, env);
            let tag = ctx.builder.ins().stack_load(types::I64, arg, 0);
            let val0 = ctx.builder.ins().stack_load(types::I64, arg, 8);
            let val1 = ctx.builder.ins().stack_load(types::I64, arg, 16);
            let val2 = ctx.builder.ins().stack_load(types::I64, arg, 24);
            ctx.builder
                .ins()
                .store(MemFlags::new(), tag, args, i as i32 * 32);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val0, args, i as i32 * 32 + 8);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val1, args, i as i32 * 32 + 16);
            ctx.builder
                .ins()
                .store(MemFlags::new(), val2, args, i as i32 * 32 + 24);
        }
        ctx.call(func, args, self.args.len(), engine, env);
        func
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for With {
    /// Compiles a `with` expression to Cranelift IR.
    ///
    /// This enters a new `with` scope with the compiled namespace, compiles the body expression,
    /// and then exits the `with` scope.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        let namespace = self.namespace.compile(ctx, engine, env);
        ctx.enter_with(env, namespace);
        let ret = self.expr.compile(ctx, engine, env);
        ctx.exit_with(env);
        ctx.free_slot(namespace);
        ret
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Assert {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for ConcatStrings {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Const {
    /// Compiles a constant value to Cranelift IR.
    ///
    /// This handles boolean, integer, float, and null constants by storing
    /// their values and type tags in a stack slot.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        use nixjit_value::Const::*;
        let slot = ctx.alloca();
        match self.val {
            Bool(x) => {
                let tag = ctx.builder.ins().iconst(types::I64, Value::BOOL as i64);
                let val = ctx.builder.ins().iconst(types::I64, x as i64);
                ctx.builder.ins().stack_store(tag, slot, 0);
                ctx.builder.ins().stack_store(val, slot, 8);
            }
            Int(x) => {
                let tag = ctx.builder.ins().iconst(types::I64, Value::INT as i64);
                let val = ctx.builder.ins().iconst(types::I64, x);
                ctx.builder.ins().stack_store(tag, slot, 0);
                ctx.builder.ins().stack_store(val, slot, 8);
            }
            Float(x) => {
                let tag = ctx.builder.ins().iconst(types::I64, Value::FLOAT as i64);
                let val = ctx.builder.ins().f64const(x);
                ctx.builder.ins().stack_store(tag, slot, 0);
                ctx.builder.ins().stack_store(val, slot, 8);
            }
            Null => {
                let tag = ctx.builder.ins().iconst(types::I64, Value::NULL as i64);
                ctx.builder.ins().stack_store(tag, slot, 0);
            }
        }
        slot
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Str {
    /// Compiles a string literal to Cranelift IR.
    ///
    /// This creates a string value from the string literal.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        ctx.create_string(&self.val)
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Var {
    /// Compiles a variable lookup to Cranelift IR.
    ///
    /// This looks up a variable by its symbol in the current environment.
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        ctx.lookup(env, &self.sym)
    }
}

impl<Ctx: JITContext> JITCompile<Ctx> for Path {
    fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
        todo!()
    }
}