torsh-jit 0.1.3

//! Cranelift backend for JIT code generation
//!
//! This module provides code generation using the Cranelift compiler framework,
//! translating IR to native machine code.

#[cfg(feature = "cranelift-backend")]
use cranelift::prelude::*;
#[cfg(feature = "cranelift-backend")]
use cranelift_codegen::ir::Function;
#[cfg(feature = "cranelift-backend")]
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
#[cfg(feature = "cranelift-backend")]
use cranelift_module::{Linkage, Module};
#[cfg(feature = "cranelift-backend")]
use cranelift_object::{ObjectBuilder, ObjectModule};

use crate::ir::{IrModule, IrOpcode, IrType, IrValue, TypeKind};
use crate::{CompiledKernel, JitError, JitResult, KernelMetadata, TensorDesc};
use std::collections::HashMap;

/// Cranelift-based code generator
#[cfg(feature = "cranelift-backend")]
pub struct CraneliftCodeGen {
    /// Cranelift module for object generation
    module: ObjectModule,

    /// Function builder context
    #[allow(dead_code)]
    builder_context: FunctionBuilderContext,

    /// Compiled functions
    #[allow(dead_code)]
    compiled_functions: HashMap<String, *const u8>,
}

#[cfg(feature = "cranelift-backend")]
impl CraneliftCodeGen {
    /// Create a new Cranelift code generator
    pub fn new() -> JitResult<Self> {
        let mut flag_builder = settings::builder();
        flag_builder
            .set("use_colocated_libcalls", "false")
            .map_err(|e| JitError::CodeGenError(format!("Settings error: {}", e)))?;
        flag_builder
            .set("is_pic", "false")
            .map_err(|e| JitError::CodeGenError(format!("Settings error: {}", e)))?;

        let isa_builder = cranelift_native::builder()
            .map_err(|e| JitError::CodeGenError(format!("ISA builder error: {}", e)))?;

        let isa = isa_builder
            .finish(settings::Flags::new(flag_builder))
            .map_err(|e| JitError::CodeGenError(format!("ISA creation error: {}", e)))?;

        let object_builder = ObjectBuilder::new(
            isa,
            "jit_module".to_string(),
            cranelift_module::default_libcall_names(),
        )
        .map_err(|e| JitError::CodeGenError(format!("Object builder error: {}", e)))?;

        let module = ObjectModule::new(object_builder);

        Ok(Self {
            module,
            builder_context: FunctionBuilderContext::new(),
            compiled_functions: HashMap::new(),
        })
    }

    /// Generate code for an IR module
    pub fn generate(&mut self, ir_module: &IrModule) -> JitResult<Vec<CompiledKernel>> {
        let mut kernels = Vec::new();

        // Generate main function for the module
        let function_id = self.declare_function(&ir_module.name, ir_module)?;
        let compiled_code = self.compile_function(ir_module, function_id)?;

        // Create kernel metadata
        let metadata = KernelMetadata {
            inputs: ir_module
                .inputs
                .iter()
                .filter_map(|&input| self.ir_value_to_tensor_desc(ir_module, input))
                .collect(),
            outputs: ir_module
                .outputs
                .iter()
                .filter_map(|&output| self.ir_value_to_tensor_desc(ir_module, output))
                .collect(),
            shared_memory: 0,
            block_size: (1, 1, 1),
            grid_size: (1, 1, 1),
        };

        let kernel = CompiledKernel {
            id: format!("cranelift_{}", ir_module.name),
            source_nodes: Vec::new(), // Would need to track this through lowering
            code: compiled_code,
            metadata,
        };

        kernels.push(kernel);
        Ok(kernels)
    }

    /// Declare a function in the module
    fn declare_function(
        &mut self,
        name: &str,
        ir_module: &IrModule,
    ) -> JitResult<cranelift_module::FuncId> {
        let mut signature = self.module.make_signature();

        // Add parameters for each input based on IR module
        for _ in &ir_module.inputs {
            signature.params.push(AbiParam::new(types::F64));
        }

        // Add return value if there are outputs
        if !ir_module.outputs.is_empty() {
            signature.returns.push(AbiParam::new(types::F64));
        }

        signature.call_conv = self.module.isa().default_call_conv();

        let func_id = self
            .module
            .declare_function(name, Linkage::Export, &signature)
            .map_err(|e| JitError::CodeGenError(format!("Function declaration error: {}", e)))?;

        Ok(func_id)
    }

    /// Compile a function from IR
    fn compile_function(
        &mut self,
        ir_module: &IrModule,
        func_id: cranelift_module::FuncId,
    ) -> JitResult<Vec<u8>> {
        let mut func = Function::new();
        let mut ctx = cranelift_codegen::Context::new();

        // Get the function signature
        func.signature = self
            .module
            .declarations()
            .get_function_decl(func_id)
            .signature
            .clone();

        // Build the function body
        {
            let mut local_builder_context = FunctionBuilderContext::new();
            let mut builder = FunctionBuilder::new(&mut func, &mut local_builder_context);

            // Create entry block
            let entry_block = builder.create_block();
            builder.append_block_params_for_function_params(entry_block);
            builder.switch_to_block(entry_block);
            builder.seal_block(entry_block);

            // Generate code for IR
            let result_value = Self::generate_ir_body_static(ir_module, &mut builder)?;

            // Return the computed value if there are outputs
            if !ir_module.outputs.is_empty() {
                let return_val = result_value.unwrap_or_else(|| {
                    // If no result, return a default f64 value
                    builder.ins().f64const(0.0)
                });
                builder.ins().return_(&[return_val]);
            } else {
                // Void function - no return value
                builder.ins().return_(&[]);
            }
            builder.finalize();
        }

        // Compile the function.
        // Per cranelift-module docs: "After calling define_function the given Context
        // will contain the compiled function."  We extract the machine-code bytes from
        // ctx.compiled_code() immediately after the call.
        ctx.func = func;
        self.module
            .define_function(func_id, &mut ctx)
            .map_err(|e| JitError::CodeGenError(format!("Function definition error: {}", e)))?;

        let code_bytes = ctx
            .compiled_code()
            .ok_or_else(|| {
                JitError::CodeGenError(
                    "Cranelift produced no compiled code after define_function — \
                     the function body may be malformed"
                        .to_string(),
                )
            })?
            .code_buffer()
            .to_vec();

        if code_bytes.is_empty() {
            return Err(JitError::CodeGenError(
                "Cranelift emitted zero bytes for the compiled function — \
                 the IR body may be empty or degenerate"
                    .to_string(),
            ));
        }

        Ok(code_bytes)
    }

    /// Static version of generate_ir_body to avoid borrowing issues
    fn generate_ir_body_static(
        ir_module: &IrModule,
        builder: &mut FunctionBuilder,
    ) -> JitResult<Option<Value>> {
        let mut value_map: HashMap<IrValue, Value> = HashMap::new();

        // Add function inputs to value_map
        let entry_block_id = builder.current_block().expect("Entry block should exist");
        let block_params = builder.block_params(entry_block_id).to_vec();

        for (i, &ir_value) in ir_module.inputs.iter().enumerate() {
            if let Some(&param_value) = block_params.get(i) {
                value_map.insert(ir_value, param_value);
            }
        }

        // Add any input parameters that might be missing (fallback)
        for (i, &param_value) in block_params.iter().enumerate() {
            let ir_value = IrValue(i as u32);
            if !value_map.contains_key(&ir_value) {
                value_map.insert(ir_value, param_value);
            }
        }

        // Process entry block
        if let Some(entry_block) = ir_module.blocks.get(&ir_module.entry_block) {
            for instruction in &entry_block.instructions {
                let cranelift_values =
                    Self::generate_instruction_static(instruction, &value_map, builder, ir_module)?;

                // Map result value if any
                if let (Some(result), Some(cranelift_value)) =
                    (instruction.result, cranelift_values.first())
                {
                    value_map.insert(result, *cranelift_value);
                }
            }
        }

        // Return the value of the first output (if any)
        let result_value = if !ir_module.outputs.is_empty() {
            let output_ir_value = ir_module.outputs[0];
            value_map.get(&output_ir_value).copied()
        } else {
            None
        };

        Ok(result_value)
    }

    /// Generate Cranelift IR for the IR module body
    #[allow(dead_code)]
    fn generate_ir_body(
        &self,
        ir_module: &IrModule,
        builder: &mut FunctionBuilder,
    ) -> JitResult<()> {
        let mut value_map: HashMap<IrValue, Value> = HashMap::new();

        // Add function inputs to value_map
        let entry_block_id = builder.current_block().expect("Entry block should exist");
        let block_params = builder.block_params(entry_block_id).to_vec();
        for (i, &ir_value) in ir_module.inputs.iter().enumerate() {
            if let Some(&param_value) = block_params.get(i) {
                value_map.insert(ir_value, param_value);
            }
        }

        // Process entry block
        if let Some(entry_block) = ir_module.blocks.get(&ir_module.entry_block) {
            for instruction in &entry_block.instructions {
                let cranelift_values =
                    self.generate_instruction(instruction, &value_map, builder, ir_module)?;

                // Map result value if any
                if let (Some(result), Some(cranelift_value)) =
                    (instruction.result, cranelift_values.first())
                {
                    value_map.insert(result, *cranelift_value);
                }
            }
        }

        Ok(())
    }

    /// Static version of generate_instruction to avoid borrowing issues
    fn generate_instruction_static(
        instruction: &crate::ir::Instruction,
        value_map: &HashMap<IrValue, Value>,
        builder: &mut FunctionBuilder,
        ir_module: &IrModule,
    ) -> JitResult<Vec<Value>> {
        use cranelift_codegen::ir::condcodes::FloatCC;

        // Get operand values
        let operands: Result<Vec<Value>, JitError> = instruction
            .operands
            .iter()
            .map(|&operand| {
                value_map
                    .get(&operand)
                    .copied()
                    .or_else(|| Self::generate_constant_value_static(operand, builder, ir_module))
                    .ok_or_else(|| {
                        JitError::CodeGenError(format!("Operand {:?} not found", operand))
                    })
            })
            .collect();

        let operands = operands?;

        // Generate instruction based on opcode
        let result = match &instruction.opcode {
            // Arithmetic operations
            IrOpcode::Add => {
                if operands.len() == 2 {
                    Some(builder.ins().fadd(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Add requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Sub => {
                if operands.len() == 2 {
                    Some(builder.ins().fsub(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Sub requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Mul => {
                if operands.len() == 2 {
                    Some(builder.ins().fmul(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Mul requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Div => {
                if operands.len() == 2 {
                    Some(builder.ins().fdiv(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Div requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Rem => {
                if operands.len() == 2 {
                    // Remainder: a - floor(a/b) * b
                    let quotient = builder.ins().fdiv(operands[0], operands[1]);
                    let floor_quotient = builder.ins().floor(quotient);
                    let product = builder.ins().fmul(floor_quotient, operands[1]);
                    Some(builder.ins().fsub(operands[0], product))
                } else {
                    return Err(JitError::CodeGenError(
                        "Rem requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Neg => {
                if operands.len() == 1 {
                    Some(builder.ins().fneg(operands[0]))
                } else {
                    return Err(JitError::CodeGenError("Neg requires 1 operand".to_string()));
                }
            }

            IrOpcode::Abs => {
                if operands.len() == 1 {
                    Some(builder.ins().fabs(operands[0]))
                } else {
                    return Err(JitError::CodeGenError("Abs requires 1 operand".to_string()));
                }
            }

            // Mathematical functions
            IrOpcode::Sqrt => {
                if operands.len() == 1 {
                    Some(builder.ins().sqrt(operands[0]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Sqrt requires 1 operand".to_string(),
                    ));
                }
            }

            IrOpcode::Exp | IrOpcode::Log | IrOpcode::Sin | IrOpcode::Cos | IrOpcode::Tanh => {
                // These would require libm calls - placeholder implementation
                if operands.len() == 1 {
                    Some(operands[0]) // Pass-through for now
                } else {
                    return Err(JitError::CodeGenError(format!(
                        "{:?} requires 1 operand",
                        instruction.opcode
                    )));
                }
            }

            IrOpcode::Sigmoid => {
                if operands.len() == 1 {
                    // sigmoid(x) = 1 / (1 + exp(-x)) - simplified
                    let one = builder.ins().f64const(1.0);
                    let neg_x = builder.ins().fneg(operands[0]);
                    let exp_approx = builder.ins().fabs(neg_x); // Placeholder
                    let denominator = builder.ins().fadd(one, exp_approx);
                    Some(builder.ins().fdiv(one, denominator))
                } else {
                    return Err(JitError::CodeGenError(
                        "Sigmoid requires 1 operand".to_string(),
                    ));
                }
            }

            // Activation functions
            IrOpcode::Relu => {
                if operands.len() == 1 {
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().fmax(operands[0], zero))
                } else {
                    return Err(JitError::CodeGenError(
                        "ReLU requires 1 operand".to_string(),
                    ));
                }
            }

            IrOpcode::Gelu => {
                if operands.len() == 1 {
                    // Simplified GELU: 0.5 * x * (1 + tanh(...))
                    let half = builder.ins().f64const(0.5);
                    let one = builder.ins().f64const(1.0);
                    let x_half = builder.ins().fmul(half, operands[0]);
                    Some(builder.ins().fmul(x_half, one)) // Simplified
                } else {
                    return Err(JitError::CodeGenError(
                        "GELU requires 1 operand".to_string(),
                    ));
                }
            }

            // Comparison operations
            IrOpcode::Eq => {
                if operands.len() == 2 {
                    let cmp = builder.ins().fcmp(FloatCC::Equal, operands[0], operands[1]);
                    let one = builder.ins().f64const(1.0);
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().select(cmp, one, zero))
                } else {
                    return Err(JitError::CodeGenError("Eq requires 2 operands".to_string()));
                }
            }

            IrOpcode::Ne => {
                if operands.len() == 2 {
                    let cmp = builder
                        .ins()
                        .fcmp(FloatCC::NotEqual, operands[0], operands[1]);
                    let one = builder.ins().f64const(1.0);
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().select(cmp, one, zero))
                } else {
                    return Err(JitError::CodeGenError("Ne requires 2 operands".to_string()));
                }
            }

            IrOpcode::Lt => {
                if operands.len() == 2 {
                    let cmp = builder
                        .ins()
                        .fcmp(FloatCC::LessThan, operands[0], operands[1]);
                    let one = builder.ins().f64const(1.0);
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().select(cmp, one, zero))
                } else {
                    return Err(JitError::CodeGenError("Lt requires 2 operands".to_string()));
                }
            }

            IrOpcode::Le => {
                if operands.len() == 2 {
                    let cmp =
                        builder
                            .ins()
                            .fcmp(FloatCC::LessThanOrEqual, operands[0], operands[1]);
                    let one = builder.ins().f64const(1.0);
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().select(cmp, one, zero))
                } else {
                    return Err(JitError::CodeGenError("Le requires 2 operands".to_string()));
                }
            }

            IrOpcode::Gt => {
                if operands.len() == 2 {
                    let cmp = builder
                        .ins()
                        .fcmp(FloatCC::GreaterThan, operands[0], operands[1]);
                    let one = builder.ins().f64const(1.0);
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().select(cmp, one, zero))
                } else {
                    return Err(JitError::CodeGenError("Gt requires 2 operands".to_string()));
                }
            }

            IrOpcode::Ge => {
                if operands.len() == 2 {
                    let cmp =
                        builder
                            .ins()
                            .fcmp(FloatCC::GreaterThanOrEqual, operands[0], operands[1]);
                    let one = builder.ins().f64const(1.0);
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().select(cmp, one, zero))
                } else {
                    return Err(JitError::CodeGenError("Ge requires 2 operands".to_string()));
                }
            }

            // Memory operations
            IrOpcode::Load => {
                if operands.len() == 1 {
                    let flags = cranelift_codegen::ir::MemFlagsData::trusted();
                    Some(builder.ins().load(types::F64, flags, operands[0], 0))
                } else {
                    return Err(JitError::CodeGenError(
                        "Load requires 1 operand".to_string(),
                    ));
                }
            }

            IrOpcode::Store => {
                if operands.len() == 2 {
                    let flags = cranelift_codegen::ir::MemFlagsData::trusted();
                    builder.ins().store(flags, operands[1], operands[0], 0);
                    None // Store doesn't return a value
                } else {
                    return Err(JitError::CodeGenError(
                        "Store requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Const => {
                // Constants should be handled separately
                None
            }

            IrOpcode::Call => {
                // Function calls - would need more implementation
                return Err(JitError::CodeGenError(
                    "Function calls not yet implemented".to_string(),
                ));
            }

            IrOpcode::Intrinsic(ref name) => {
                // Handle intrinsic operations (e.g., fused operations)
                Self::generate_intrinsic(name, &operands, builder, ir_module)?
            }

            _ => {
                return Err(JitError::CodeGenError(format!(
                    "Unsupported opcode: {:?}",
                    instruction.opcode
                )));
            }
        };

        Ok(if let Some(value) = result {
            vec![value]
        } else {
            vec![]
        })
    }

    /// Generate Cranelift instruction for IR instruction
    #[allow(dead_code)]
    fn generate_instruction(
        &self,
        instruction: &crate::ir::Instruction,
        value_map: &HashMap<IrValue, Value>,
        builder: &mut FunctionBuilder,
        ir_module: &IrModule,
    ) -> JitResult<Vec<Value>> {
        // Get operand values
        let operands: Result<Vec<Value>, JitError> = instruction
            .operands
            .iter()
            .map(|&operand| {
                value_map
                    .get(&operand)
                    .copied()
                    .or_else(|| self.generate_constant_value(operand, builder, ir_module))
                    .ok_or_else(|| {
                        JitError::CodeGenError(format!("Operand {:?} not found", operand))
                    })
            })
            .collect();

        let operands = operands?;

        // Generate instruction based on opcode
        let result = match &instruction.opcode {
            IrOpcode::Add => {
                if operands.len() == 2 {
                    Some(builder.ins().fadd(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Add requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Sub => {
                if operands.len() == 2 {
                    Some(builder.ins().fsub(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Sub requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Mul => {
                if operands.len() == 2 {
                    Some(builder.ins().fmul(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Mul requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Div => {
                if operands.len() == 2 {
                    Some(builder.ins().fdiv(operands[0], operands[1]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Div requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Neg => {
                if operands.len() == 1 {
                    Some(builder.ins().fneg(operands[0]))
                } else {
                    return Err(JitError::CodeGenError("Neg requires 1 operand".to_string()));
                }
            }

            IrOpcode::Abs => {
                if operands.len() == 1 {
                    Some(builder.ins().fabs(operands[0]))
                } else {
                    return Err(JitError::CodeGenError("Abs requires 1 operand".to_string()));
                }
            }

            IrOpcode::Sqrt => {
                if operands.len() == 1 {
                    Some(builder.ins().sqrt(operands[0]))
                } else {
                    return Err(JitError::CodeGenError(
                        "Sqrt requires 1 operand".to_string(),
                    ));
                }
            }

            IrOpcode::Relu => {
                if operands.len() == 1 {
                    // ReLU: max(0, x)
                    let zero = builder.ins().f64const(0.0);
                    Some(builder.ins().fmax(operands[0], zero))
                } else {
                    return Err(JitError::CodeGenError(
                        "ReLU requires 1 operand".to_string(),
                    ));
                }
            }

            IrOpcode::Load => {
                // Load from memory - simplified
                if operands.len() == 1 {
                    let flags = cranelift_codegen::ir::MemFlagsData::trusted();
                    Some(builder.ins().load(types::F64, flags, operands[0], 0))
                } else {
                    return Err(JitError::CodeGenError(
                        "Load requires 1 operand".to_string(),
                    ));
                }
            }

            IrOpcode::Store => {
                // Store to memory - simplified
                if operands.len() == 2 {
                    let flags = cranelift_codegen::ir::MemFlagsData::trusted();
                    builder.ins().store(flags, operands[1], operands[0], 0);
                    None // Store doesn't return a value
                } else {
                    return Err(JitError::CodeGenError(
                        "Store requires 2 operands".to_string(),
                    ));
                }
            }

            IrOpcode::Const => {
                // Constants should be handled separately
                None
            }

            IrOpcode::Call => {
                // Function calls - would need more implementation
                return Err(JitError::CodeGenError(
                    "Function calls not yet implemented".to_string(),
                ));
            }

            IrOpcode::Intrinsic(ref name) => {
                // Handle intrinsic operations (e.g., fused operations)
                Self::generate_intrinsic(name, &operands, builder, ir_module)?
            }

            _ => {
                return Err(JitError::CodeGenError(format!(
                    "Unsupported opcode: {:?}",
                    instruction.opcode
                )));
            }
        };

        Ok(if let Some(value) = result {
            vec![value]
        } else {
            vec![]
        })
    }

    /// Static version of generate_constant_value
    fn generate_constant_value_static(
        ir_value: IrValue,
        builder: &mut FunctionBuilder,
        ir_module: &IrModule,
    ) -> Option<Value> {
        if let Some(value_def) = ir_module.get_value(ir_value) {
            if let crate::ir::ValueKind::Constant { ref data } = value_def.kind {
                match data {
                    crate::ir::ConstantData::Float(f) => Some(builder.ins().f64const(*f)),
                    crate::ir::ConstantData::Int(i) => Some(builder.ins().iconst(types::I64, *i)),
                    crate::ir::ConstantData::Bool(b) => {
                        Some(builder.ins().iconst(types::I8, if *b { 1 } else { 0 }))
                    }
                    _ => None, // Other constant types not supported yet
                }
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Generate a constant value
    #[allow(dead_code)]
    fn generate_constant_value(
        &self,
        ir_value: IrValue,
        builder: &mut FunctionBuilder,
        ir_module: &IrModule,
    ) -> Option<Value> {
        if let Some(value_def) = ir_module.get_value(ir_value) {
            if let crate::ir::ValueKind::Constant { ref data } = value_def.kind {
                match data {
                    crate::ir::ConstantData::Float(f) => Some(builder.ins().f64const(*f)),
                    crate::ir::ConstantData::Int(i) => Some(builder.ins().iconst(types::I64, *i)),
                    crate::ir::ConstantData::Bool(b) => {
                        Some(builder.ins().iconst(types::I8, if *b { 1 } else { 0 }))
                    }
                    _ => None, // Other constant types not supported yet
                }
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Convert IR type to Cranelift type
    #[allow(dead_code)]
    fn ir_type_to_cranelift(&self, ir_type: IrType, ir_module: &IrModule) -> Option<Type> {
        if let Some(type_def) = ir_module.get_type(ir_type) {
            match &type_def.kind {
                TypeKind::Bool => Some(types::I8),
                TypeKind::I8 => Some(types::I8),
                TypeKind::I16 => Some(types::I16),
                TypeKind::I32 => Some(types::I32),
                TypeKind::I64 => Some(types::I64),
                TypeKind::U8 => Some(types::I8), // Cranelift doesn't distinguish signed/unsigned
                TypeKind::U16 => Some(types::I16),
                TypeKind::U32 => Some(types::I32),
                TypeKind::U64 => Some(types::I64),
                TypeKind::F16 => Some(types::F32), // Promote to F32
                TypeKind::F32 => Some(types::F32),
                TypeKind::F64 => Some(types::F64),
                TypeKind::Pointer { .. } => Some(types::I64), // Pointer as 64-bit int
                _ => None,                                    // Complex types not supported yet
            }
        } else {
            None
        }
    }

    /// Convert IR value to tensor descriptor
    fn ir_value_to_tensor_desc(
        &self,
        ir_module: &IrModule,
        ir_value: IrValue,
    ) -> Option<TensorDesc> {
        if let Some(value_def) = ir_module.get_value(ir_value) {
            if let Some(type_def) = ir_module.get_type(value_def.ty) {
                match &type_def.kind {
                    TypeKind::Tensor { shape, .. } => {
                        Some(TensorDesc {
                            dtype: torsh_core::DType::F32, // Simplified
                            shape: shape.clone(),
                            strides: self.compute_strides(shape),
                            offset: 0,
                        })
                    }
                    _ => None,
                }
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Compute strides for a shape
    fn compute_strides(&self, shape: &[usize]) -> Vec<usize> {
        let mut strides = vec![1; shape.len()];
        for i in (0..shape.len().saturating_sub(1)).rev() {
            strides[i] = strides[i + 1] * shape[i + 1];
        }
        strides
    }

    /// Finalize the module and get the compiled object
    pub fn finalize(self) -> JitResult<Vec<u8>> {
        let object = self
            .module
            .finish()
            .object
            .write()
            .map_err(|e| JitError::CodeGenError(format!("Object finalization error: {}", e)))?;

        Ok(object)
    }

    /// Generate code for intrinsic operations (e.g., fused operations)
    fn generate_intrinsic(
        name: &str,
        operands: &[Value],
        builder: &mut FunctionBuilder,
        _ir_module: &IrModule,
    ) -> JitResult<Option<Value>> {
        match name {
            "fused_1" | "linear_relu" | "conv_relu" => {
                // For now, implement fused operations as separate operations
                // This is a simplified version - in production, you'd want proper fusion
                if !operands.is_empty() {
                    // Use the first (or only) operand as the input
                    let input = operands[0];

                    // Apply ReLU activation: max(0, x)
                    let zero = builder.ins().f64const(0.0);
                    let result = builder.ins().fmax(input, zero);

                    Ok(Some(result))
                } else {
                    Err(JitError::CodeGenError(format!(
                        "Intrinsic {} requires at least 1 operand",
                        name
                    )))
                }
            }
            _ => {
                // For unknown intrinsics, just return the first operand
                // This is a fallback to prevent compilation failure
                if !operands.is_empty() {
                    Ok(Some(operands[0]))
                } else {
                    Ok(None)
                }
            }
        }
    }
}

#[cfg(not(feature = "cranelift-backend"))]
pub struct CraneliftCodeGen;

#[cfg(not(feature = "cranelift-backend"))]
impl CraneliftCodeGen {
    pub fn new() -> JitResult<Self> {
        Err(JitError::CodeGenError(
            "Cranelift backend not enabled".to_string(),
        ))
    }

    pub fn generate(&mut self, _ir_module: &IrModule) -> JitResult<Vec<CompiledKernel>> {
        Err(JitError::CodeGenError(
            "Cranelift backend not enabled".to_string(),
        ))
    }

    pub fn finalize(self) -> JitResult<Vec<u8>> {
        Err(JitError::CodeGenError(
            "Cranelift backend not enabled".to_string(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cranelift_codegen_creation() {
        #[cfg(feature = "cranelift-backend")]
        {
            let _codegen = CraneliftCodeGen::new();
            // Test passes if it can be created
        }

        #[cfg(not(feature = "cranelift-backend"))]
        {
            let result = CraneliftCodeGen::new();
            assert!(result.is_err());
        }
    }

    /// Verify that compile_function produces non-empty machine code bytes for a
    /// trivially valid IR module (void function, no inputs, no outputs, single
    /// empty block).  An empty byte vector would indicate the previous silent
    /// fabrication is back.
    #[test]
    #[cfg(feature = "cranelift-backend")]
    fn test_compile_function_emits_non_empty_bytes() {
        use crate::ir::IrModule;

        let mut ir_module = IrModule::new("test_fn".to_string());
        // Add an empty entry block with no inputs/outputs — the simplest valid IR.
        let block_id = ir_module.add_block();
        ir_module.entry_block = block_id;

        let mut codegen = CraneliftCodeGen::new().expect("CraneliftCodeGen::new() should succeed");

        let kernels = codegen
            .generate(&ir_module)
            .expect("generate() should succeed for a void function");

        assert!(
            !kernels.is_empty(),
            "generate() must produce at least one kernel"
        );
        assert!(
            !kernels[0].code.is_empty(),
            "compiled kernel code must be non-empty — empty bytes mean no machine code was emitted"
        );
    }
}