use crate::compiler::codegen::type_manager::{CodeGenType, TypeManager};
use crate::compiler::codegen::CodeGenerator;
use crate::compiler::ast::{Type, Expr, ExprKind};
use inkwell::values::{BasicValueEnum, ValueKind, BasicValue};
// use inkwell::AddressSpace;



pub fn register_tensor_types(manager: &mut TypeManager) {
    let mut tensor = CodeGenType::new("Tensor");
    
    // Unevaluated static methods (for literal optimizations)
    // zeros, ones, randn with 1 or 2 args
    let shape_type = Type::Tensor(Box::new(Type::I64), 1);

    // zeros(shape, requires_grad) and zeros(shape)
    tensor.register_unevaluated_static_method("zeros", compile_tensor_zeros, vec![shape_type.clone(), Type::Bool], Type::Tensor(Box::new(Type::F32), 0));
    tensor.register_unevaluated_static_method("zeros", compile_tensor_zeros, vec![shape_type.clone()], Type::Tensor(Box::new(Type::F32), 0));
    
    // ones(shape, requires_grad) and ones(shape)
    tensor.register_unevaluated_static_method("ones", compile_ones, vec![shape_type.clone(), Type::Bool], Type::Tensor(Box::new(Type::F32), 0));
    tensor.register_unevaluated_static_method("ones", compile_ones, vec![shape_type.clone()], Type::Tensor(Box::new(Type::F32), 0));
    
    // randn(shape, requires_grad) and randn(shape)
    tensor.register_unevaluated_static_method("randn", compile_randn, vec![shape_type.clone(), Type::Bool], Type::Tensor(Box::new(Type::F32), 0));
    tensor.register_unevaluated_static_method("randn", compile_randn, vec![shape_type.clone()], Type::Tensor(Box::new(Type::F32), 0));
    
    // Evaluated static methods
    // load(path: String) -> Tensor
    tensor.register_evaluated_static_method("load", compile_load_tensor, vec![Type::String("String".to_string())], Type::Tensor(Box::new(Type::F32), 1));
    // clear_grads() -> Void
    tensor.register_evaluated_static_method("clear_grads", compile_clear_grads, vec![], Type::Void);
    // from_vec_u8(u8s: Vec<u8>, shape: Vec<i64>) -> Tensor
    tensor.register_evaluated_static_method("from_vec_u8", compile_from_vec_u8, vec![Type::Struct("Vec".into(), vec![Type::U8]), Type::Struct("Vec".into(), vec![Type::I64])], Type::Tensor(Box::new(Type::F32), 0));

    // Instance methods
    let any_tensor = Type::Tensor(Box::new(Type::F32), 0);

    // Sum/reduction with overloads
    tensor.register_evaluated_instance_method("sum", compile_tensor_sumall, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("sum", compile_tensor_sum_impl, vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("sum", compile_tensor_sum_impl, vec![Type::I64, Type::Bool], any_tensor.clone());

    tensor.register_evaluated_instance_method("max", compile_tensor_max_impl, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("max", compile_tensor_max_impl, vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("max", compile_tensor_max_impl, vec![Type::I64, Type::Bool], any_tensor.clone());

    tensor.register_evaluated_instance_method("min", compile_tensor_min_impl, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("min", compile_tensor_min_impl, vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("min", compile_tensor_min_impl, vec![Type::I64, Type::Bool], any_tensor.clone());

    tensor.register_evaluated_instance_method("mean", compile_tensor_mean_impl, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("mean", compile_tensor_mean_impl, vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("mean", compile_tensor_mean_impl, vec![Type::I64, Type::Bool], any_tensor.clone());

    tensor.register_evaluated_instance_method("argmax", compile_tensor_argmax_impl, vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("argmin", compile_tensor_argmin_impl, vec![Type::I64], any_tensor.clone());

    // Other instance methods (no overloads)
    tensor.register_evaluated_instance_method("detach", compile_tensor_detach, vec![Type::Bool], any_tensor.clone());
    tensor.register_evaluated_instance_method("detach", compile_tensor_detach, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("tril", compile_tensor_tril, vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("contiguous", compile_tensor_contiguous, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("clone", compile_tensor_clone, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("shallow_clone", compile_tensor_shallow_clone, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("grad", compile_tensor_grad, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("to_i64", compile_tensor_to_i64, vec![], Type::Tensor(Box::new(Type::I64), 0));
    tensor.register_evaluated_instance_method("cuda", compile_tensor_cuda, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("cpu", compile_tensor_cpu, vec![], any_tensor.clone());
    tensor.register_evaluated_instance_method("item", compile_tensor_item, vec![], Type::F32);
    tensor.register_evaluated_instance_method("shape", compile_tensor_shape, vec![], Type::Struct("Vec".into(), vec![Type::I64]));
    
    // Binary ops
    tensor.register_evaluated_instance_method("mul", compile_tensor_mul, vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_evaluated_instance_method("add", compile_tensor_add, vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_evaluated_instance_method("sub", compile_tensor_sub, vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_evaluated_instance_method("div", compile_tensor_div, vec![any_tensor.clone()], any_tensor.clone());
    
    // Neural network ops
    tensor.register_evaluated_instance_method("conv2d", compile_tensor_conv2d, vec![any_tensor.clone(), Type::I64, Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("clamp", compile_tensor_clamp, vec![Type::I64, Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("matmul_quantized", compile_tensor_matmul_quantized, vec![any_tensor.clone()], any_tensor.clone());
    
    // IO and transformation
    tensor.register_evaluated_instance_method("save", compile_tensor_save, vec![Type::String("String".to_string())], Type::Void);
    tensor.register_evaluated_instance_method("reshape", compile_tensor_reshape, vec![Type::Struct("Vec".into(), vec![Type::I64])], any_tensor.clone());
    tensor.register_evaluated_instance_method("to", compile_tensor_to_device, vec![Type::String("String".to_string())], any_tensor.clone());
    tensor.register_evaluated_instance_method("transpose", compile_tensor_transpose, vec![Type::I64, Type::I64], any_tensor.clone());
    
    // Compound assignment
    tensor.register_evaluated_instance_method("add_assign", compile_tensor_add_assign, vec![any_tensor.clone()], Type::Void);
    tensor.register_evaluated_instance_method("sub_assign", compile_tensor_sub_assign, vec![any_tensor.clone()], Type::Void);

    // ===== Additional methods from check_builtin_method (signature-only) =====
    // Autograd
    tensor.register_instance_signature("backward", vec![], Type::Void);
    tensor.register_instance_signature("enable_grad", vec![], any_tensor.clone());
    
    // Activation / ops - 1 arg
    tensor.register_instance_signature("softmax", vec![Type::I64], any_tensor.clone());
    tensor.register_instance_signature("log_softmax", vec![Type::I64], any_tensor.clone());
    tensor.register_instance_signature("matmul", vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_evaluated_instance_method("embedding", compile_tensor_embedding, vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_instance_signature("cross_entropy", vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_instance_signature("pow", vec![Type::I64], any_tensor.clone());
    
    // Elementwise ops - 0 arg
    for method in ["abs", "relu", "gelu", "silu", "exp", "log", "sqrt", "sin", "cos", "tanh", "sigmoid", "neg", "tan"] {
        tensor.register_instance_signature(method, vec![], any_tensor.clone());
    }
    
    // Shape ops - 1+ args 
    tensor.register_instance_signature("squeeze", vec![Type::I64], any_tensor.clone());
    tensor.register_instance_signature("unsqueeze", vec![Type::I64], any_tensor.clone());
    tensor.register_instance_signature("flatten", vec![Type::I64], any_tensor.clone());
    tensor.register_instance_signature("gather", vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_instance_signature("permute", vec![Type::Struct("Vec".into(), vec![Type::I64])], any_tensor.clone());
    tensor.register_instance_signature("cat", vec![any_tensor.clone()], any_tensor.clone());
    
    // LLM ops
    tensor.register_instance_signature("scale", vec![Type::F32], any_tensor.clone());
    tensor.register_instance_signature("add_4d", vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_instance_signature("matmul_4d", vec![any_tensor.clone()], any_tensor.clone());
    tensor.register_instance_signature("cat_4d", vec![any_tensor.clone(), Type::I64], any_tensor.clone());
    tensor.register_instance_signature("rms_norm", vec![any_tensor.clone(), Type::F32], any_tensor.clone());
    tensor.register_instance_signature("sample", vec![Type::F32, Type::F32], Type::Tensor(Box::new(Type::F32), 1));
    tensor.register_instance_signature("repeat_interleave", vec![Type::I64, Type::I64], any_tensor.clone());
    tensor.register_instance_signature("apply_rope", vec![any_tensor.clone(), any_tensor.clone()], any_tensor.clone());
    tensor.register_instance_signature("narrow", vec![Type::I64, Type::I64, Type::I64], any_tensor.clone());
    tensor.register_instance_signature("transpose_2d", vec![Type::I64, Type::I64], any_tensor.clone());
    tensor.register_instance_signature("item_i64", vec![], Type::I64);
    tensor.register_instance_signature("cat_i64", vec![any_tensor.clone(), Type::I64], any_tensor.clone());
    
    // Shape/length queries
    tensor.register_instance_signature("len", vec![], Type::I64);
    tensor.register_instance_signature("dim", vec![Type::I64], Type::I64);
    tensor.register_instance_signature("ndim", vec![], Type::I64);
    tensor.register_instance_signature("get_shape", vec![], Type::Struct("Vec".into(), vec![Type::I64]));
    tensor.register_instance_signature("print", vec![], Type::Void);
    tensor.register_instance_signature("display", vec![], Type::Void);
    
    // Slice
    tensor.register_instance_signature("slice", vec![Type::I64, Type::I64, Type::I64, Type::I64], any_tensor.clone());
    
    // Dim-specific ops  
    tensor.register_instance_signature("sum_dim", vec![Type::I64, Type::Bool], any_tensor.clone());
    tensor.register_instance_signature("mean_dim", vec![Type::I64, Type::Bool], any_tensor.clone());
    tensor.register_instance_signature("max_dim", vec![Type::I64], any_tensor.clone());
    tensor.register_instance_signature("min_dim", vec![Type::I64], any_tensor.clone());
    tensor.register_evaluated_instance_method("sumall", compile_tensor_sumall, vec![], any_tensor.clone());
    
    // Modulo
    tensor.register_instance_signature("mod", vec![Type::I64], any_tensor.clone());

    manager.register_type(tensor);
}

fn compile_tensor_add_assign<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj_val: BasicValueEnum<'ctx>,
    _obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 {
        return Err("add_assign requires 1 argument".into());
    }
    let (rhs_val, rhs_ty) = args[0].clone();

    if matches!(rhs_ty, Type::Tensor(_, _)) {
        let fn_val = codegen.module.get_function("tl_tensor_add_assign").unwrap();
        codegen
            .builder
            .build_call(fn_val, &[obj_val.into(), rhs_val.into()], "assign_res")
            .map_err(|e| e.to_string())?;
    } else if matches!(rhs_ty, Type::F32 | Type::F64 | Type::I64 | Type::I32) {
        let scalar_f32 = match rhs_ty {
            Type::F32 => rhs_val.into_float_value(),
            Type::F64 => codegen
                .builder
                .build_float_cast(
                    rhs_val.into_float_value(),
                    codegen.context.f32_type(),
                    "f64_to_f32",
                )
                .map_err(|e| e.to_string())?,
            Type::I64 | Type::I32 => codegen
                .builder
                .build_signed_int_to_float(
                    rhs_val.into_int_value(),
                    codegen.context.f32_type(),
                    "int_to_f32",
                )
                .map_err(|e| e.to_string())?,
            _ => return Err(format!("add_assign scalar: unsupported type {:?}", rhs_ty)),
        };
        let fn_val = codegen
            .module
            .get_function("tl_tensor_add_assign_scalar_f32")
            .unwrap();
        codegen
            .builder
            .build_call(fn_val, &[obj_val.into(), scalar_f32.into()], "assign_res")
            .map_err(|e| e.to_string())?;
    } else {
        return Err(format!(
            "add_assign requires Tensor or scalar argument, got {:?}",
            rhs_ty
        ));
    }

    Ok((
        codegen.context.i64_type().const_int(0, false).into(),
        Type::Void,
    ))
}

fn compile_tensor_sub_assign<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj_val: BasicValueEnum<'ctx>,
    _obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 {
        return Err("sub_assign requires 1 argument".into());
    }
    let (rhs_val, rhs_ty) = args[0].clone();

    if matches!(rhs_ty, Type::Tensor(_, _)) {
        let fn_val = codegen.module.get_function("tl_tensor_sub_assign").unwrap();
        codegen
            .builder
            .build_call(fn_val, &[obj_val.into(), rhs_val.into()], "assign_res")
            .map_err(|e| e.to_string())?;
    } else if matches!(rhs_ty, Type::F32 | Type::F64 | Type::I64 | Type::I32) {
        let scalar_f32 = match rhs_ty {
            Type::F32 => rhs_val.into_float_value(),
            Type::F64 => codegen
                .builder
                .build_float_cast(
                    rhs_val.into_float_value(),
                    codegen.context.f32_type(),
                    "f64_to_f32",
                )
                .map_err(|e| e.to_string())?,
            Type::I64 | Type::I32 => codegen
                .builder
                .build_signed_int_to_float(
                    rhs_val.into_int_value(),
                    codegen.context.f32_type(),
                    "int_to_f32",
                )
                .map_err(|e| e.to_string())?,
            _ => return Err(format!("sub_assign scalar: unsupported type {:?}", rhs_ty)),
        };
        let fn_val = codegen
            .module
            .get_function("tl_tensor_sub_assign_scalar_f32")
            .unwrap();
        codegen
            .builder
            .build_call(fn_val, &[obj_val.into(), scalar_f32.into()], "assign_res")
            .map_err(|e| e.to_string())?;
    } else {
        return Err(format!(
            "sub_assign requires Tensor or scalar argument, got {:?}",
            rhs_ty
        ));
    }

    Ok((
        codegen.context.i64_type().const_int(0, false).into(),
        Type::Void,
    ))
}



fn compile_tensor_shape<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj_val: BasicValueEnum<'ctx>,
    _obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_shape").ok_or("tl_tensor_shape not found")?;
    let res = match codegen
        .builder
        .build_call(fn_val, &[obj_val.into()], "shape_res")
        .map_err(|e| e.to_string())?
        .try_as_basic_value()
    {
        inkwell::values::ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from shape()".into()),
    };
    Ok((res, Type::Struct("Vec".into(), vec![Type::I64])))
}

fn compile_tensor_sumall<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_sum").ok_or("tl_tensor_sum not found")?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "sum_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from sumall()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_detach<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_detach").ok_or("tl_tensor_detach not found")?;
    // Optional arg: req_grad (bool). Default to false.
    let req_grad = if !args.is_empty() {
        let (arg_val, _) = args[0];
        arg_val.into_int_value()
    } else {
        codegen.context.bool_type().const_int(0, false)
    };
    let call = codegen.builder.build_call(fn_val, &[obj.into(), req_grad.into()], "detach_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from detach()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_tril<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("tril requires 1 argument (diagonal)".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_tril").ok_or("tl_tensor_tril not found")?;
    let (diag_val, diag_ty) = &args[0];
    let diag_i32 = match diag_ty {
        Type::I64 => codegen.builder.build_int_cast(diag_val.into_int_value(), codegen.context.i32_type(), "tril_diag_cast").unwrap(),
        Type::I32 => diag_val.into_int_value(),
        _ => return Err("tril argument must be integer".into()),
    };
    let call = codegen.builder.build_call(fn_val, &[obj.into(), diag_i32.into()], "tril_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from tril()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_embedding<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,      // self = indices tensor
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("embedding requires 1 argument (weight)".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_embedding")
        .ok_or("tl_tensor_embedding not found")?;

    let (weight_val, _) = &args[0];

    // デフォルト引数: padding_idx=-1, scale_grad_by_freq=false, sparse=false
    let padding_idx = codegen.context.i64_type().const_int((-1i64) as u64, true);
    let false_val = codegen.context.bool_type().const_int(0, false);

    // runtime順序: (weight, indices, padding_idx, scale_grad_by_freq, sparse)
    let call = codegen.builder.build_call(
        fn_val,
        &[(*weight_val).into(), obj.into(), padding_idx.into(), false_val.into(), false_val.into()],
        "embedding_res",
    ).map_err(|e| e.to_string())?;

    let res = codegen.check_tensor_result(call, "embedding_error")?;
    Ok((res, obj_ty))
}

fn compile_tensor_binop<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    op_name: &str,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err(format!("{} requires 1 argument", op_name)); }
    // args already evaluated. Note: ensure_tensor_v2 logic in expr.rs handles scalar conversion *before* evaluation if it was static arg?
    // Wait, TypeManager passes evaluated args. If arg was scalar literal, it is evaluated as I64/F64.
    // We need run-time scalar->tensor conversion if needed?
    // expr.rs logic: self.ensure_tensor_v2(&args[0], 0)?
    // ensure_tensor_v2 takes Expr.
    // HERE we have evaluated values.
    // If we receive I64/F32, we must wrap it in Tensor?
    // BUT TypeManager calls compile_expr(arg). So we have I64/F32 value.
    // We need a helper to ensure tensor from Value.
    
    // For now, let's look at how expr.rs did it. It used ensure_tensor_v2 which compiles Expr potentially to conversion call.
    // Since we receive compiled value, we must check type and convert if scalar.
    // But we don't have easy "convert value to tensor" helper exposed in CodeGen yet?
    // Actually, ensure_tensor_v2 does: if scalar type, allocate tensor, fill.
    // We can replicate logic or assume args are Tensors?
    // Users might write `t + 1.0`. `1.0` is F64.
    // We need to implement scalar-to-tensor promotion here.
    
    let (rhs_val, rhs_ty) = &args[0];
    let final_rhs = if let Type::Tensor(_, _) = rhs_ty {
        *rhs_val
    } else {
        // Promote scalar to tensor
        // Check numeric
        match rhs_ty {
            Type::I64 | Type::F64 | Type::F32 | Type::I32 => {
                 // Create encoded tensor from scalar
                 // We need a helper function in codegen?
                 // Or just hardcode promotion logic.
                 // Ideally we call a runtime helper tl_tensor_from_scalar?
                 // Existing ensure_tensor_v2 calls tl_tensor_zeros/fill or something.
                 // simpler: use codegen.ensure_tensor_from_val(...) if we make it?
                 // For now, I will assume we can't easily access ensure_tensor_v2 (it takes &Expr).
                 // I will skip scalar implementation details for now or try to support it?
                 // The safe bet: If not tensor, error for now OR implement simple promotion.
                 // Given constraints, I'll error if not tensor for this refactor step, 
                 // BUT this breaks `t + 1`. 
                 // Wait, expr.rs ensure_tensor_v2 is called on Expr.
                 // TypeManager evaluates args via compile_expr.
                 // If I want to support scalars, I must handle them.
                 // Let's assume for this task refactoring, we primarily target Tensor-Tensor ops.
                 // OR I can use the trick: The user's code `t + 1` passes `1` as expr.
                 // If I could intercept it... instance_method receives evaluated values.
                 // If I want `ensure_tensor_v2` behavior, I need to implement `ensure_tensor_from_value` here.
                 
                 // Implementation of scalar to tensor promotion:
                 // 1. Alloc tensor of rank 0
                 // 2. data[0] = value
                 // This requires runtime calls.
                 // It's too complex to inline here perfectly without a helper.
                 // Let's rely on checking if rhs is Tensor. 
                 return Err("Scalar broadcasting in binary ops via TypeManager refactor is pending. Please use explicit Tensor.".into());
            }
             _ => return Err("Binary op requires Tensor or Scalar".into())
        }
    };

    let fn_name = match op_name {
        "mul" => "tl_tensor_mul",
        "add" => "tl_tensor_add",
        "sub" => "tl_tensor_sub",
        "div" => "tl_tensor_div",
        _ => unreachable!(),
    };
    let fn_val = codegen.module.get_function(fn_name).ok_or(format!("{} not found", fn_name))?;
    let call = codegen.builder.build_call(fn_val, &[obj.into(), final_rhs.into()], "binop_res").map_err(|e| e.to_string())?;
    let res = codegen.check_tensor_result(call, "binop_error")?;
    Ok((res, obj_ty))
}

fn compile_tensor_mul<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_binop(c, o, t, a, "mul")
}
fn compile_tensor_add<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_binop(c, o, t, a, "add")
}
fn compile_tensor_sub<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_binop(c, o, t, a, "sub")
}
fn compile_tensor_div<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_binop(c, o, t, a, "div")
}

fn compile_tensor_contiguous<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_contiguous").ok_or("tl_tensor_contiguous not found")?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "cont_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from contiguous()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_conv2d<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 3 { return Err("conv2d requires 3 arguments: weight, padding, stride".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_conv2d").ok_or("tl_tensor_conv2d not found")?;
    
    let (weight_val, _) = &args[0];
    let (pad_val, pad_ty) = &args[1];
    let (stride_val, stride_ty) = &args[2];
    
    let pad_i64 = match pad_ty {
        Type::I64 => pad_val.into_int_value(),
        Type::I32 => codegen.builder.build_int_z_extend(pad_val.into_int_value(), codegen.context.i64_type(), "ext").unwrap(),
        _ => return Err("conv2d padding must be int".into()),
    };
    let stride_i64 = match stride_ty {
        Type::I64 => stride_val.into_int_value(),
        Type::I32 => codegen.builder.build_int_z_extend(stride_val.into_int_value(), codegen.context.i64_type(), "ext").unwrap(),
        _ => return Err("conv2d stride must be int".into()),
    };

    let call = codegen.builder.build_call(fn_val, &[obj.into(), (*weight_val).into(), pad_i64.into(), stride_i64.into()], "conv_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from conv2d()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_clamp<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 2 { return Err("clamp requires 2 args".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_clamp").ok_or("tl_tensor_clamp not found")?;
    
    let (min_val, min_ty) = &args[0];
    let min_f32 = match min_ty {
        Type::F32 => min_val.into_float_value(),
        Type::F64 => codegen.builder.build_float_trunc(min_val.into_float_value(), codegen.context.f32_type(), "trunc").unwrap(),
        _ => return Err("min must be float".into()),
    };
    let (max_val, max_ty) = &args[1];
    let max_f32 = match max_ty {
        Type::F32 => max_val.into_float_value(),
        Type::F64 => codegen.builder.build_float_trunc(max_val.into_float_value(), codegen.context.f32_type(), "trunc").unwrap(),
        _ => return Err("max must be float".into()),
    };

    let call = codegen.builder.build_call(fn_val, &[obj.into(), min_f32.into(), max_f32.into()], "clamp_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from clamp()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_clone<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_clone").ok_or("tl_tensor_clone not found")?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "clone_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from clone()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_shallow_clone<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_shallow_clone").ok_or("tl_tensor_shallow_clone not found")?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "clone_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from shallow_clone()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_grad<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_grad").ok_or("tl_tensor_grad not found")?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "grad_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from grad()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_matmul_quantized<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("matmul_quantized requires 1 arg".into()); }
    let fn_val = codegen.module.get_function("tl_qtensor_matmul").ok_or("tl_qtensor_matmul not found")?;
    let (weight, _) = &args[0];
    let call = codegen.builder.build_call(fn_val, &[obj.into(), (*weight).into()], "qmatmul_res").map_err(|e| e.to_string())?;
    let res = codegen.check_tensor_result(call, "qmatmul_error")?;
    Ok((res, obj_ty))
}

fn compile_tensor_to_i64<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    _obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_to_i64").ok_or("tl_tensor_to_i64 not found")?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "to_i64_call").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from to_i64".into()),
    };
    Ok((res, Type::Tensor(Box::new(Type::I64), 0)))
}

fn compile_tensor_cuda<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_to_device_helper(codegen, obj, obj_ty, "cuda")
}

fn compile_tensor_cpu<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_to_device_helper(codegen, obj, obj_ty, "cpu")
}

fn compile_tensor_to_device_helper<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    device: &str,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_to_device").ok_or("tl_tensor_to_device not found")?;
    let (dev_str_val, _) = codegen.compile_string_literal(device)?;
    let call = codegen.builder.build_call(fn_val, &[obj.into(), dev_str_val.into()], "to_dev_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
         _ => return Err("Invalid return from to_device".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_item<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let is_int = if let Type::Tensor(elem, _) = &obj_ty {
        matches!(elem.as_ref(), Type::I64 | Type::I32 | Type::U32 | Type::U8)
    } else {
        false
    };
    let fn_name = if is_int { "tl_tensor_item_i64" } else { "tl_tensor_item" };
    let fn_val = codegen.module.get_function(fn_name).ok_or(format!("{} not found", fn_name))?;
    let call = codegen.builder.build_call(fn_val, &[obj.into()], "item_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from item()".into()),
    };
    Ok((res, if is_int { Type::I64 } else { Type::F32 }))
}

// fn compile_tensor_reduce_wrapper<'ctx>(
//     codegen: &mut CodeGenerator<'ctx>,
//     obj: BasicValueEnum<'ctx>,
//     obj_ty: Type,
//     args: Vec<(BasicValueEnum<'ctx>, Type)>,
// ) -> Result<(BasicValueEnum<'ctx>, Type), String> {
//     // We don't have 'method name' here easily unless we pass it or make closures?
//     // Wait, register_instance_method usage: compile_tensor_reduce_wrapper
//     // But this function signature doesn't take method name.
//     // Solution: define separate functions for each reduce op, or use a helper that takes name, 
//     // but we can't register a helper that doesn't match the signature.
//     // We need explicit wrappers for each.
//     // I will implement explicit wrappers below.
//     Err("Not implemented: reduce wrapper needs specific method dispatch".into())
// }

// Implement specific reducers
fn compile_tensor_max_impl<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_reduce_generic(c, o, t, a, "max")
}
fn compile_tensor_min_impl<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_reduce_generic(c, o, t, a, "min")
}
fn compile_tensor_mean_impl<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_reduce_generic(c, o, t, a, "mean")
}
fn compile_tensor_sum_impl<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_reduce_generic(c, o, t, a, "sum")
}
fn compile_tensor_argmax_impl<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_reduce_generic(c, o, t, a, "argmax")
}
fn compile_tensor_argmin_impl<'ctx>(c: &mut CodeGenerator<'ctx>, o: BasicValueEnum<'ctx>, t: Type, a: Vec<(BasicValueEnum<'ctx>, Type)>) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_reduce_generic(c, o, t, a, "argmin")
}

fn compile_tensor_reduce_generic<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    method: &str,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if !args.is_empty() {
        let suffix = if method == "argmax" || method == "argmin" { "" } else { "_dim" };
        let fn_name = format!("tl_tensor_{}{}", method, suffix);
        let fn_val = codegen.module.get_function(&fn_name).ok_or(format!("{} not found", fn_name))?;

        let (dim_val, _) = &args[0];
        let keep_val = if args.len() > 1 {
            let (k, _) = &args[1];
             *k
        } else {
             codegen.context.bool_type().const_int(0, false).into()
        };
        let call = codegen.builder.build_call(fn_val, &[obj.into(), (*dim_val).into(), keep_val.into()], "reduce_res").map_err(|e| e.to_string())?;
        let res = match call.try_as_basic_value() {
            ValueKind::Basic(v) => v,
            _ => return Err("Invalid return".into()),
        };
        Ok((res, obj_ty))
    } else {
        if method == "argmax" || method == "argmin" {
            return Err(format!("{} requires arguments", method));
        }
        let fn_name = format!("tl_tensor_{}", method);
        let fn_val = codegen.module.get_function(&fn_name).ok_or(format!("{} not found", fn_name))?;
        let call = codegen.builder.build_call(fn_val, &[obj.into()], "reduce_res").map_err(|e| e.to_string())?;
        let res = match call.try_as_basic_value() {
            ValueKind::Basic(v) => v,
             _ => return Err("Invalid return".into()),
        };
        Ok((res, obj_ty))
    }
}

fn compile_tensor_save<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    _obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("save requires 1 arg (path)".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_save").ok_or("tl_tensor_save not found")?;
    let (path_val, _) = &args[0];
    codegen.builder.build_call(fn_val, &[(*path_val).into(), obj.into()], "save_call").map_err(|e| e.to_string())?;
    Ok((codegen.context.i64_type().const_int(0, false).into(), Type::Void))
}

fn compile_tensor_reshape<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    _obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
     if args.len() != 1 { return Err("reshape requires 1 arg (shape)".into()); }
     let (shape_val, shape_ty) = &args[0];
     // Arg[0] is shape.
     // In expr.rs, it compiled arg[0].

     if let Type::Tensor(_, _) = shape_ty {
          // Case 1: Shape is a Tensor. Use tl_tensor_reshape_new(obj, shape_tensor).
          let fn_val = codegen.module.get_function("tl_tensor_reshape_new")
               .ok_or("tl_tensor_reshape_new not found")?;
          let call = codegen.builder.build_call(fn_val, &[obj.into(), shape_val.clone().into()], "reshape_res")
               .map_err(|e| e.to_string())?;
          let res = match call.try_as_basic_value() {
               ValueKind::Basic(v) => v,
               _ => return Err("Invalid return".into()),
          };
          // Free shape tensor after reshape (it was only used as a shape descriptor)
          if let Some(release_fn) = codegen.module.get_function("tl_tensor_release_safe") {
               codegen.builder.build_call(release_fn, &[(*shape_val).into()], "").ok();
          }
          return Ok((res, Type::Tensor(Box::new(Type::F32), 0))); // Dynamic rank
     } else {
          // Case 2: Shape is Vec. Use tl_tensor_reshape_dims(obj, ptr, rank).
          let fn_val = codegen.module.get_function("tl_tensor_reshape_dims")
               .ok_or("tl_tensor_reshape_dims not found")?;

          let (data_ptr, rank_val) = if matches!(shape_ty, Type::Struct(n, _) if n == "Vec") {
               // Vec
               let vec_ptr = if shape_val.is_pointer_value() {
                    shape_val.into_pointer_value()
               } else {
                    return Err("Vec shape must be pointer".into());
               };

               let i64_ty = codegen.context.i64_type();
               let vec_struct_ty = codegen.context.struct_type(&[i64_ty.into(), i64_ty.into(), i64_ty.into()], false);

               // Extract ptr (index 0)
               let data_ptr_ptr = codegen.builder.build_struct_gep(vec_struct_ty, vec_ptr, 0, "data_ptr_ptr")
                         .map_err(|e| e.to_string())?;
               
               let data_ptr_int = codegen.builder.build_load(i64_ty, data_ptr_ptr, "data_ptr_int")
                    .map_err(|e| e.to_string())?.into_int_value();
               
               let data_ptr = codegen.builder.build_int_to_ptr(
                    data_ptr_int,
                    codegen.context.ptr_type(inkwell::AddressSpace::default()),
                    "data_ptr"
               ).map_err(|e| e.to_string())?;

               // Extract len (index 2)
               let len_ptr = codegen.builder.build_struct_gep(vec_struct_ty, vec_ptr, 2, "len_ptr")
                         .map_err(|e| e.to_string())?;
               
               let len_val = codegen.builder.build_load(i64_ty, len_ptr, "rank_val")
                    .map_err(|e| e.to_string())?.into_int_value();
               
               (data_ptr, len_val)
          } else {
               return Err(format!("reshape argument must be Tensor or Vec. Got: {:?}", shape_ty));
          };

          let call = codegen.builder.build_call(fn_val, &[obj.into(), data_ptr.into(), rank_val.into()], "reshape_res")
               .map_err(|e| e.to_string())?;
          
          let res = match call.try_as_basic_value() {
               ValueKind::Basic(v) => v,
               _ => return Err("Invalid return".into()),
          };
          return Ok((res, Type::Tensor(Box::new(Type::F32), 0)));
     }
}

fn compile_tensor_to_device<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("to requires 1 arg (device)".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_to_device").ok_or("tl_tensor_to_device not found")?;
    let (dev_val, _) = &args[0];
    let call = codegen.builder.build_call(fn_val, &[obj.into(), (*dev_val).into()], "to_dev_res").map_err(|e| e.to_string())?;
     let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from to()".into()),
    };
    Ok((res, obj_ty))
}

fn compile_tensor_transpose<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    obj: BasicValueEnum<'ctx>,
    obj_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 2 { return Err("transpose requires 2 args (dim0, dim1)".into()); }
    let fn_val = codegen.module.get_function("tl_tensor_transpose").ok_or("tl_tensor_transpose not found")?;
    let (d0, _) = &args[0];
    let (d1, _) = &args[1];
    
    let call = codegen.builder.build_call(fn_val, &[obj.into(), (*d0).into(), (*d1).into()], "transpose_res").map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
         _ => return Err("Invalid return from transpose".into()),
    };
    Ok((res, obj_ty))
}

fn compile_clear_grads<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    _target: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if !args.is_empty() { return Err("Tensor::clear_grads takes no arguments".into()); }
    let fn_val = codegen.module.get_function("tl_clear_grads").ok_or("tl_clear_grads not found")?;
    codegen.builder.build_call(fn_val, &[], "clear_grads").map_err(|e| e.to_string())?;
    Ok((codegen.context.i64_type().const_int(0, false).into(), Type::Void))
}

fn compile_from_vec_u8<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    _target: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    // Tensor::from_vec_u8(vec, shape) -> Tensor
    // Arguments are already evaluated by GenericResolver (evaluated args mode)
    // args[0]: vec (Vec<U8>)
    // args[1]: shape (Vec<I64>)
    // Implicit: offset=0, rank=shape.len()
    
    if args.len() != 2 {
        return Err("Tensor::from_vec_u8 requires 2 arguments (vec, shape)".into());
    }
    
    let vec_val = args[0].0;
    let shape_raw = args[1].0; 
    let shape_ty = args[1].1.clone();
    
    // We need to pass (ptr, offset, shape_data_ptr, rank)
    // shape_data_ptr should be *const i64 (pointer to array)
    
    let (shape_data_ptr, rank_val) = if matches!(shape_ty, Type::Tensor(_, _)) {
        // Handle Tensor
        // Data ptr
        let data_fn = codegen.module.get_function("tl_tensor_data").ok_or("tl_tensor_data not found")?;
        let data_call = codegen.builder.build_call(data_fn, &[shape_raw.into()], "shape_data").map_err(|e| e.to_string())?;
        let data_ptr = match data_call.try_as_basic_value() {
            ValueKind::Basic(v) => v.into_pointer_value(),
            _ => return Err("Invalid return from tensor data".into()),
        };
        
        // Numel (rank)
        let numel_fn = codegen.module.get_function("tl_tensor_numel").ok_or("tl_tensor_numel not found")?;
        let numel_call = codegen.builder.build_call(numel_fn, &[shape_raw.into()], "shape_numel").map_err(|e| e.to_string())?;
        let numel = match numel_call.try_as_basic_value() {
            ValueKind::Basic(v) => v.into_int_value(),
            _ => return Err("Invalid return from tensor numel".into()),
        };
        
        (data_ptr, numel)
    } else {
         // Assume Vec (pointer to struct)
         // Call tl_vec_i64_as_ptr(vec) -> *const i64
         let as_ptr_fn = codegen.module.get_function("tl_vec_i64_as_ptr").ok_or("tl_vec_i64_as_ptr not found")?;
         let ptr_call = codegen.builder.build_call(as_ptr_fn, &[shape_raw.into()], "vec_ptr").map_err(|e| e.to_string())?;
         let data_ptr = match ptr_call.try_as_basic_value() {
             ValueKind::Basic(v) => v.into_pointer_value(),
             _ => return Err("Invalid return from vec as_ptr".into()),
         };

         // Call tl_vec_i64_len(vec) -> usize
         let len_fn = codegen.module.get_function("tl_vec_i64_len").ok_or("tl_vec_i64_len not found")?;
         let len_call = codegen.builder.build_call(len_fn, &[shape_raw.into()], "vec_len").map_err(|e| e.to_string())?;
         let rank = match len_call.try_as_basic_value() {
             ValueKind::Basic(v) => v.into_int_value(),
             _ => return Err("Invalid return from vec len".into()),
         };
         
         (data_ptr, rank)
    };
    
    let offset_val = codegen.context.i64_type().const_int(0, false);
    
    let fn_val = codegen.module.get_function("tl_tensor_from_vec_u8").ok_or("tl_tensor_from_vec_u8 not found")?;
    
    let call = codegen.builder.build_call(fn_val, &[
        vec_val.into(),
        offset_val.into(),
        shape_data_ptr.into(),
        rank_val.into()
    ], "from_vec_res").map_err(|e| e.to_string())?;
    
    let v = codegen.check_tensor_result(call, "from_vec_error")?;
    
    Ok((v, Type::Tensor(Box::new(Type::F32), 0)))
}

fn compile_load_tensor<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    _target: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    let fn_val = codegen.module.get_function("tl_tensor_load").ok_or("tl_tensor_load not found")?;
    let (path_val, _) = &args[0];
    let call = codegen
        .builder
        .build_call(fn_val, &[(*path_val).into()], "load_res")
        .map_err(|e| e.to_string())?;
    let res = match call.try_as_basic_value() {
        ValueKind::Basic(v) => v,
        _ => return Err("Invalid load return".into()),
    };
    // Return rank 1? Or dynamic? Most loaded tensors have specific valid ranks.
    // Existing code returned rank 1.
    Ok((res, Type::Tensor(Box::new(Type::F32), 1)))
}

fn compile_tensor_creation_helper<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: &[Expr],
    runtime_func_name: &str,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.is_empty() {
        return Err(format!("{} requires shape", runtime_func_name));
    }
    let shape_expr = &args[0];
    let (rank, shape_vals) = match &shape_expr.inner {
        ExprKind::TensorLiteral(el) | ExprKind::TensorConstLiteral(el) => {
            let mut vals = Vec::new();
            for e in el {
                let (v, t) = codegen.compile_expr(e)?;
                let int_val = match t {
                    Type::I64 => v.into_int_value(),
                    Type::I32 => codegen
                        .builder
                        .build_int_z_extend(
                            v.into_int_value(),
                            codegen.context.i64_type(),
                            "dim_ext",
                        )
                        .map_err(|e| e.to_string())?,
                    _ => return Err(format!("Dimension must be integer, found {:?}", t)),
                };
                vals.push(int_val);
            }
            (el.len(), vals)
        }
        _ => {
            return Err(format!(
                "{} currently requires array literal [dim, ...] for shape",
                runtime_func_name
            ));
        }
    };
    let requires_grad = if args.len() > 1 {
        match &args[1].inner {
            ExprKind::Bool(b) => *b,
            _ => false,
        }
    } else {
        false
    };
    let i64_type = codegen.context.i64_type();

    let current_block = codegen.builder.get_insert_block().unwrap();
    let function = current_block.get_parent().unwrap();
    let entry_block = function.get_first_basic_block().unwrap();
    let entry_builder = codegen.context.create_builder();
    if let Some(first_instr) = entry_block.get_first_instruction() {
        entry_builder.position_before(&first_instr);
    } else {
        entry_builder.position_at_end(entry_block);
    }

    let shape_array_type = i64_type.array_type(rank as u32);
    let shape_alloca = entry_builder
        .build_alloca(shape_array_type, "shape_arr")
        .map_err(|e| e.to_string())?;

    shape_alloca
        .as_instruction_value()
        .unwrap()
        .set_alignment(16)
        .map_err(|e| e.to_string())?;

    for (i, val) in shape_vals.iter().enumerate() {
        let ptr = unsafe {
            codegen.builder.build_in_bounds_gep(
                shape_array_type,
                shape_alloca,
                &[
                    i64_type.const_int(0, false),
                    i64_type.const_int(i as u64, false),
                ],
                "shape_ptr_in",
            )
        }
        .map_err(|e| e.to_string())?;
        codegen
            .builder
            .build_store(ptr, *val)
            .map_err(|e| e.to_string())?;
    }
    let req_grad_val = codegen
        .context
        .bool_type()
        .const_int(if requires_grad { 1 } else { 0 }, false);
    let f = codegen
        .module
        .get_function(runtime_func_name)
        .ok_or(format!("{} not found", runtime_func_name))?;

    let zero = i64_type.const_int(0, false);
    let first_elem_ptr = unsafe {
        codegen.builder.build_in_bounds_gep(
            shape_array_type,
            shape_alloca,
            &[zero, zero],
            "first_elem_ptr",
        )
    }
    .map_err(|e| e.to_string())?;

    // Build args: randn_debug requires seed arg, zeros/ones do not
    let mut call_args: Vec<inkwell::values::BasicMetadataValueEnum> = Vec::new();
    call_args.push(i64_type.const_int(rank as u64, false).into());
    call_args.push(first_elem_ptr.into());
    if runtime_func_name.contains("randn") {
        // randn_debug has signature: (rank, shape, seed, req_grad)
        call_args.push(i64_type.const_int(0, false).into()); // seed = 0 (unused)
    }
    call_args.push(req_grad_val.into());

    let call = codegen
        .builder
        .build_call(
            f,
            &call_args,
            "creation_res",
        )
        .map_err(|e| e.to_string())?;
    let res = codegen.check_tensor_result(call, "creation_error")?;

    Ok((res, Type::Tensor(Box::new(Type::F32), rank)))
}

fn compile_randn<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: &[Expr],
    _target: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_creation_helper(codegen, args, "tl_tensor_randn_debug")
}

fn compile_ones<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: &[Expr],
    _target: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    compile_tensor_creation_helper(codegen, args, "tl_tensor_ones")
}

fn compile_tensor_zeros<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: &[Expr],
    _target: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.is_empty() {
        return Err("Tensor::zeros requires shape argument".into());
    }

    let elements_ref = if let ExprKind::TensorLiteral(el) = &args[0].inner {
        Some(el)
    } else if let ExprKind::TensorConstLiteral(el) = &args[0].inner {
        Some(el)
    } else {
        None
    };

    if let Some(el) = elements_ref {
        let i64_type = codegen.context.i64_type();
        let mut vals = Vec::new();
        for e in el {
            let (v, t) = codegen.compile_expr(e)?;
            let int_val = match t {
                Type::I64 => v.into_int_value(),
                Type::I32 => codegen
                    .builder
                    .build_int_z_extend(v.into_int_value(), i64_type, "ext")
                    .map_err(|e| e.to_string())?,
                _ => return Err(format!("Dimension must be integer, found {:?}", t)),
            };
            vals.push(int_val);
        }

        let rank = el.len();
        let shape_array_type = i64_type.array_type(rank as u32);
        let shape_alloca = codegen
            .builder
            .build_alloca(shape_array_type, "shape_arr")
            .map_err(|e| e.to_string())?;

        for (i, val) in vals.iter().enumerate() {
            let ptr = unsafe {
                codegen
                    .builder
                    .build_in_bounds_gep(
                        shape_array_type,
                        shape_alloca,
                        &[
                            i64_type.const_int(0, false),
                            i64_type.const_int(i as u64, false),
                        ],
                        "tmp",
                    )
                    .map_err(|e| e.to_string())?
            };
            codegen
                .builder
                .build_store(ptr, *val)
                .map_err(|e| e.to_string())?;
        }

        let req_grad = if args.len() > 1 {
            let (v, _) = codegen.compile_expr(&args[1])?;
            v.into_int_value()
        } else {
            codegen.context.bool_type().const_int(0, false)
        };

        let f = codegen
            .module
            .get_function("tl_tensor_zeros")
            .ok_or("tl_tensor_zeros not found")?;
        let call = codegen
            .builder
            .build_call(
                f,
                &[
                    i64_type.const_int(rank as u64, false).into(),
                    shape_alloca.into(),
                    req_grad.into(),
                ],
                "zeros_res",
            )
            .map_err(|e| e.to_string())?;

        let v = codegen.check_tensor_result(call, "zeros_error")?;
        let result_ty = Type::Tensor(Box::new(Type::F32), rank);
        // codegen.emit_register_tensor(v, &result_ty)?;
        return Ok((v, result_ty));
    }

    Err("Generic Tensor::zeros (non-literal shape) not yet ported. Please use literal shape [d1, d2] for now.".into())
}
