use crate::compiler::codegen::type_manager::{CodeGenType, TypeManager};
use crate::compiler::codegen::CodeGenerator;
use crate::compiler::ast::Type;
use crate::compiler::builtin_loader::{BuiltinLoader, BuiltinTypeData};
use inkwell::values::BasicValueEnum;





pub const SOURCE: &str = include_str!("llm_types.tl");

pub fn load_llm_data() -> BuiltinTypeData {
    BuiltinLoader::load_module_data(SOURCE, "LLM")
        .expect("Failed to load LLM type data")
}


pub fn register_llm_types(manager: &mut TypeManager) {
    let string_type = Type::String("String".to_string());
    
    // Register Tokenizer
    let mut tokenizer = CodeGenType::new("Tokenizer");
    // Tokenizer::new(path: String) -> Tokenizer
    tokenizer.register_evaluated_static_method(
        "new", 
        compile_tokenizer_new,
        vec![string_type.clone()],
        Type::Struct("Tokenizer".to_string(), vec![])
    );
    // Tokenizer.encode(prompt: String) -> Tensor<i64>
    tokenizer.register_evaluated_instance_method(
        "encode", 
        compile_tokenizer_encode,
        vec![string_type.clone()],
        Type::Tensor(Box::new(Type::I64), 1) // 1D tensor of IDs
    );
    // Tokenizer.decode(ids: Tensor<i64>) -> String
    tokenizer.register_evaluated_instance_method(
        "decode", 
        compile_tokenizer_decode,
        vec![Type::Tensor(Box::new(Type::I64), 1)],
        string_type.clone()
    );
    // Tokenizer.encode_chat(user_msg: String) -> Tensor<i64>
    // Llama 3 チャットテンプレートに準拠したトークン列を生成
    tokenizer.register_evaluated_instance_method(
        "encode_chat", 
        compile_tokenizer_encode_chat,
        vec![string_type.clone()],
        Type::Tensor(Box::new(Type::I64), 1)
    );
    manager.register_type(tokenizer);

    // Register KVCache
    let mut kv_cache = CodeGenType::new("KVCache");
    // KVCache::new(layers: i64) -> KVCache
    kv_cache.register_evaluated_static_method(
        "new", 
        compile_kv_cache_new,
        vec![Type::I64],
        Type::Struct("KVCache".to_string(), vec![])
    );
    // KVCache.free() -> Void
    kv_cache.register_evaluated_instance_method(
        "free", 
        compile_kv_cache_free,
        vec![],
        Type::Void
    );
    // KVCache.get_k(layer: i64) -> Tensor<f32>
    kv_cache.register_evaluated_instance_method(
        "get_k", 
        compile_kv_cache_get_k,
        vec![Type::I64],
        Type::Tensor(Box::new(Type::F32), 2) // Assuming 2D? (batch, dim) or (seq, dim)
    );
    // KVCache.get_v(layer: i64) -> Tensor<f32>
    kv_cache.register_evaluated_instance_method(
        "get_v", 
        compile_kv_cache_get_v,
        vec![Type::I64],
        Type::Tensor(Box::new(Type::F32), 2)
    );
    // KVCache.update(layer: i64, k: Tensor, v: Tensor) -> Void
    // k and v ranks depend on implementation, usually 2D or 3D. 
    // Let's assume generic Tensor for now as runtime handles it, but signature expects specific type.
    // Based on usage in compile_kv_cache_update, it checks if it is Tensor.
    // I'll define it as taking Tensor<f32> with rank 0 (generic) or explicit if known. 
    // Usually these are 2D/3D tensors. 
    // Let's use generic Tensor<f32>.
    let tensor_type = Type::Tensor(Box::new(Type::F32), 0); 
    kv_cache.register_evaluated_instance_method(
        "update", 
        compile_kv_cache_update,
        vec![Type::I64, tensor_type.clone(), tensor_type.clone()],
        Type::Void
    );

    manager.register_type(kv_cache);

    // VarBuilder (for loading model weights from safetensors)
    let mut var_builder = CodeGenType::new("VarBuilder");
    // VarBuilder::new(path: String) -> VarBuilder
    var_builder.register_static_signature("new", vec![string_type.clone()], Type::Struct("VarBuilder".into(), vec![]));
    // vb.get(name: String) -> Tensor
    var_builder.register_instance_signature("get", vec![string_type.clone()], tensor_type.clone());
    // vb.get_or_zeros(name: String, shape: Vec<i64>) -> Tensor
    var_builder.register_instance_signature("get_or_zeros", vec![string_type.clone(), Type::Struct("Vec".into(), vec![Type::I64])], tensor_type.clone());
    // vb.child(prefix: String) -> VarBuilder
    var_builder.register_instance_signature("child", vec![string_type.clone()], Type::Struct("VarBuilder".into(), vec![]));
    // vb.grad(tensor: Tensor) -> Tensor
    var_builder.register_instance_signature("grad", vec![tensor_type.clone()], tensor_type.clone());
    // vb.update() -> Void
    var_builder.register_instance_signature("update", vec![], Type::Void);
    // vb.save(path: String) -> Void
    var_builder.register_instance_signature("save", vec![string_type.clone()], Type::Void);
    manager.register_type(var_builder);
}

// ... existing static methods ...

// ... existing instance methods ...



fn compile_kv_cache_update<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    _instance_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 3 { return Err("KVCache::update requires 3 arguments".into()); }
    
    // Convert instance (ptr) to ptr
    if !instance_val.is_pointer_value() { return Err("Instance update requires pointer".into()); }
    
    let (layer_val, _) = args[0];
    let (k_val, k_ty) = &args[1];
    let (v_val, v_ty) = &args[2];

    // Helper to extract tensor pointer
    let mut get_tensor_ptr = |val: BasicValueEnum<'ctx>, ty: &Type| -> Result<inkwell::values::BasicMetadataValueEnum<'ctx>, String> {
        let ptr_val = if let Type::Struct(name, _) = ty {
            if name == "Tensor" {
                 let h = codegen.load_struct_i64_field(val, ty, "handle")?;
                 codegen.builder.build_int_to_ptr(h.into_int_value(), codegen.context.ptr_type(inkwell::AddressSpace::default()), "handle_ptr").map_err(|e| e.to_string())?
            } else {
                 return Err("Expected Tensor argument".into());
            }
        } else if let Type::Tensor(_, _) = ty {
             if val.is_pointer_value() {
                 val.into_pointer_value()
             } else {
                 return Err("Expected Tensor pointer".into());
             }
        } else {
             return Err(format!("Expected Tensor argument, got {:?}", ty));
        };
        Ok(ptr_val.into())
    };

    let k_ptr = get_tensor_ptr(*k_val, k_ty)?;
    let v_ptr = get_tensor_ptr(*v_val, v_ty)?;

    let fn_val = codegen.module.get_function("tl_kvcache_update").ok_or("tl_kvcache_update not found")?;
    codegen.builder.build_call(fn_val, &[instance_val.into(), layer_val.into(), k_ptr, v_ptr], "").map_err(|e| e.to_string())?;
    
    Ok((codegen.context.i64_type().const_int(0, false).into(), Type::Void))
}
// Tokenizer Static Methods
fn compile_tokenizer_new<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    _target_type: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 {
        return Err("Tokenizer::new requires 1 argument".into());
    }
    let (path_val, path_ty) = &args[0];
    
    let path_ptr = match path_ty {
        Type::String(_) => {
             // Extract ptr from String struct
             codegen.load_struct_i64_field(*path_val, path_ty, "ptr")?
                .into_int_value()
        },
        _ => return Err(format!("Tokenizer::new expects String argument, got {:?}", path_ty)),
    };

    let path_ptr_val = codegen.builder.build_int_to_ptr(path_ptr, codegen.context.ptr_type(inkwell::AddressSpace::default()), "path_str_ptr").map_err(|e| e.to_string())?;

    let fn_val = codegen.module.get_function("tl_tokenizer_new").ok_or("tl_tokenizer_new not found")?;
    let call = codegen.builder.build_call(fn_val, &[inkwell::values::BasicMetadataValueEnum::from(path_ptr_val)], "tok_new").map_err(|e| e.to_string())?;
    let handle = match call.try_as_basic_value() {
        inkwell::values::ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from Tokenizer::new".into()),
    };

    let struct_type = *codegen.struct_types.get("Tokenizer").ok_or("Struct type Tokenizer not found")?;
    let struct_def = codegen.struct_defs.get("Tokenizer").ok_or("Struct definition Tokenizer not found")?;
    let size = struct_type.size_of().ok_or("Cannot determine size of Tokenizer")?;
    
    let size_int = size;
    let size_i64 = if size_int.get_type() == codegen.context.i32_type() {
        codegen.builder.build_int_z_extend(size_int, codegen.context.i64_type(), "size_i64").unwrap()
    } else {
        size_int
    };

    let malloc_fn = codegen.module.get_function("malloc").ok_or("malloc not found (declare in builtins)")?;
    let call = codegen.builder.build_call(malloc_fn, &[size_i64.into()], "tokenizer_malloc").map_err(|e| e.to_string())?;
    let raw_ptr = match call.try_as_basic_value() {
        inkwell::values::ValueKind::Basic(v) => v.into_pointer_value(),
        _ => return Err("malloc returned invalid value".into()),
    };
    
    if let Some(register_fn) = codegen.module.get_function("tl_mem_register_struct") {
        let cast_ptr = codegen.builder.build_pointer_cast(raw_ptr, codegen.context.ptr_type(inkwell::AddressSpace::default()), "cast_ptr").unwrap();
        codegen.builder.build_call(register_fn, &[cast_ptr.into()], "").map_err(|e| e.to_string())?;
    }
    
    let struct_ptr = codegen.builder.build_pointer_cast(raw_ptr, codegen.context.ptr_type(inkwell::AddressSpace::default()), "tokenizer_ptr").map_err(|e| e.to_string())?;
    
    let field_idx = struct_def.fields.iter().position(|(n, _)| n == "handle").ok_or("Field handle not found in Tokenizer")?;
    let field_ptr = codegen.builder.build_struct_gep(struct_type, struct_ptr, field_idx as u32, "tokenizerhandle").map_err(|e| e.to_string())?;
    codegen.builder.build_store(field_ptr, handle).map_err(|e| e.to_string())?;
    
    Ok((struct_ptr.into(), Type::Struct("Tokenizer".to_string(), vec![])))
}

// Tokenizer Instance Methods
fn compile_tokenizer_encode<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    instance_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("Tokenizer::encode requires 1 argument".into()); }
    let handle = codegen.load_struct_i64_field(instance_val, &instance_ty, "handle")?;
    let (prompt_val, prompt_ty) = &args[0];
    
    let prompt_ptr = match prompt_ty {
        Type::String(_) => {
             // Extract ptr from String struct
             codegen.load_struct_i64_field(*prompt_val, prompt_ty, "ptr")?
                .into_int_value()
        },
        _ => return Err("Tokenizer::encode expects String argument".into()),
    };
    let prompt_ptr_val = codegen.builder.build_int_to_ptr(prompt_ptr, codegen.context.ptr_type(inkwell::AddressSpace::default()), "prompt_str_ptr").map_err(|e| e.to_string())?;
    let fn_val = codegen.module.get_function("tl_tokenizer_encode").ok_or("tl_tokenizer_encode not found")?;
    let call = codegen.builder.build_call(fn_val, &[handle.into(), inkwell::values::BasicMetadataValueEnum::from(prompt_ptr_val)], "tok_encode").map_err(|e| e.to_string())?;
    // check_tensor_result(val, msg) is method of codegen.
    let res = codegen.check_tensor_result(call, "tok_encode_error")?;
    Ok((res, Type::Tensor(Box::new(Type::I64), 0)))
}

fn compile_tokenizer_encode_chat<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    instance_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("Tokenizer::encode_chat requires 1 argument".into()); }
    let handle = codegen.load_struct_i64_field(instance_val, &instance_ty, "handle")?;
    let (prompt_val, prompt_ty) = &args[0];
    
    let prompt_ptr = match prompt_ty {
        Type::String(_) => {
             codegen.load_struct_i64_field(*prompt_val, prompt_ty, "ptr")?
                .into_int_value()
        },
        _ => return Err("Tokenizer::encode_chat expects String argument".into()),
    };
    let prompt_ptr_val = codegen.builder.build_int_to_ptr(prompt_ptr, codegen.context.ptr_type(inkwell::AddressSpace::default()), "prompt_str_ptr").map_err(|e| e.to_string())?;
    let fn_val = codegen.module.get_function("tl_tokenizer_encode_chat").ok_or("tl_tokenizer_encode_chat not found")?;
    let call = codegen.builder.build_call(fn_val, &[handle.into(), inkwell::values::BasicMetadataValueEnum::from(prompt_ptr_val)], "tok_encode_chat").map_err(|e| e.to_string())?;
    let res = codegen.check_tensor_result(call, "tok_encode_chat_error")?;
    Ok((res, Type::Tensor(Box::new(Type::I64), 0)))
}

fn compile_tokenizer_decode<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    instance_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("Tokenizer::decode requires 1 argument".into()); }
    let handle = codegen.load_struct_i64_field(instance_val, &instance_ty, "handle")?;
    let (ids_val, _) = args[0];
    let fn_val = codegen.module.get_function("tl_tokenizer_decode").ok_or("tl_tokenizer_decode not found")?;
    let call = codegen.builder.build_call(fn_val, &[handle.into(), ids_val.into()], "tok_decode").map_err(|e| e.to_string())?;
    let char_ptr = match call.try_as_basic_value() {
        inkwell::values::ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from Tokenizer::decode".into()),
    };
    
    // Wrap char* in String struct using tl_string_new
    let string_new_fn = codegen.module.get_function("tl_string_new").ok_or("tl_string_new not found")?;
    let str_call = codegen.builder.build_call(string_new_fn, &[char_ptr.into()], "tok_decode_str").map_err(|e| e.to_string())?;
    let res = match str_call.try_as_basic_value() {
        inkwell::values::ValueKind::Basic(v) => v,
        _ => return Err("Invalid return from tl_string_new".into()),
    };

    Ok((res, Type::String("String".to_string())))
}

// KVCache Static Methods
fn compile_kv_cache_new<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
    _target_type: Option<&Type>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 {
        return Err("KVCache::new requires 1 argument".into());
    }
    let layers_val = args[0].0;

    let fn_val = codegen.module.get_function("tl_kvcache_new").ok_or("tl_kvcache_new not found")?;
    
    // Alloc stack slot for return value SRET
    // CodeGenerator usually manages alloca in entry block but here we do it ad-hoc?
    // Using simple alloca in current block is risky if inside loop?
    // But `new` is usually top level.
    // Better: use codegen.builder
    
    let i64_type = codegen.context.i64_type();
    let struct_ty = codegen.context.struct_type(&[i64_type.into()], false);
    
    let alloca = codegen.builder.build_alloca(struct_ty, "kv_new_sret").unwrap();
    
    codegen.builder.build_call(fn_val, &[alloca.into(), layers_val.into()], "").map_err(|e| e.to_string())?;

    // Return the POINTER to the struct (standard for Struct return types in TL Codegen)
    Ok((alloca.into(), Type::Struct("KVCache".to_string(), vec![])))
}

// KVCache Instance Methods
fn compile_kv_cache_free<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    _instance_ty: Type,
    _args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    // KVCache.free(self) -> self is ptr
    // tl_kvcache_free(ptr)
    if !instance_val.is_pointer_value() {
         return Err("Instance not pointer in free".into());
    }
    let fn_val = codegen.module.get_function("tl_kvcache_free").ok_or("tl_kvcache_free not found")?;
    codegen.builder.build_call(fn_val, &[instance_val.into()], "").map_err(|e| e.to_string())?;
    Ok((codegen.context.i64_type().const_int(0, false).into(), Type::Void))
}


fn compile_kv_cache_get_k<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    _instance_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("KVCache::get_k requires 1 argument".into()); }
    let (layer_val, _) = args[0];
    
    // tl_kvcache_get_k(ptr, layer)
    let fn_val = codegen.module.get_function("tl_kvcache_get_k").ok_or("tl_kvcache_get_k not found")?;
    let call = codegen.builder.build_call(fn_val, &[instance_val.into(), layer_val.into()], "kv_get_k").map_err(|e| e.to_string())?;
    let res = codegen.check_tensor_result(call, "kv_get_k_error")?;
    Ok((res, Type::Tensor(Box::new(Type::F32), 2))) // Assuming 2D tensor
}


fn compile_kv_cache_get_v<'ctx>(
    codegen: &mut CodeGenerator<'ctx>,
    instance_val: BasicValueEnum<'ctx>,
    _instance_ty: Type,
    args: Vec<(BasicValueEnum<'ctx>, Type)>,
) -> Result<(BasicValueEnum<'ctx>, Type), String> {
    if args.len() != 1 { return Err("KVCache::get_v requires 1 argument".into()); }
    let (layer_val, _) = args[0];
    
    // tl_kvcache_get_v(ptr, layer)
    let fn_val = codegen.module.get_function("tl_kvcache_get_v").ok_or("tl_kvcache_get_v not found")?;
    let call = codegen.builder.build_call(fn_val, &[instance_val.into(), layer_val.into()], "kv_get_v").map_err(|e| e.to_string())?;
    let res = codegen.check_tensor_result(call, "kv_get_v_error")?;
    Ok((res, Type::Tensor(Box::new(Type::F32), 2))) // Assuming 2D tensor
}

