import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import lib.core.Handles.HandlesBase as HandlesBase
import lib.core.Handles.HandlesTranslator as HT
import lib.core.Handles.HandlesExpr as HandlesExpr
import lib.core.Handles.HandlesVar as HandlesVar


# ============================================================
# HandlesReturn - Return 语句处理（Mixin 继承模式）
# ============================================================


@t.NoVTable
class ReturnHandle(HandlesBase.Mixin):
    """Return 语句处理器：继承 Mixin 获得 Trans 回指针"""

    def __init__(self, trans: HT.Translator | t.CPtr):
        self.Trans = trans

    # ============================================================
    # Handle - 处理 Return 语句，返回 0
    #
    # 翻译返回值（若有）并生成 ret 指令
    # ============================================================
    def Handle(self, node: ast.AST | t.CPtr) -> int:
        """翻译 Return 语句"""
        rt: ast.Return | t.CPtr = (ast.Return | t.CPtr)(node)
        if rt is None:
            return 0

        pool: memhub.MemBuddy | t.CPtr = self.Trans.Pool
        builder: llvmlite.IRBuilder | t.CPtr = self.Trans._cur_builder
        mod: llvmlite.LLVMModule | t.CPtr = self.Trans.Module

        val: llvmlite.Value | t.CPtr = None
        if rt.value is not None:
            # return self: 直接返回指针（self 是 SSA 参数 Ptr(struct_ty)）
            # translate_value 会 load 得到结构体值，但 return self 需要指针本身
            if rt.value.kind() == ast.ASTKind.Name:
                ret_nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(rt.value)
                if ret_nm is not None and ret_nm.id is not None:
                    if string.strcmp(ret_nm.id, "self") == 0:
                        self_ptr: llvmlite.Value | t.CPtr = HandlesVar.lookup_var(
                            self.Trans.SymTab, "self")
                        if self_ptr is not None:
                            val = self_ptr
            if val is None:
                val = HandlesExpr.translate_value(
                    builder, pool, mod, rt.value,
                    self.Trans._funcs, self.Trans._func_count, self.Trans)

        # val 为 None 时（裸 return）：检查当前函数返回类型
        # void 函数（如 __init__/__before_init__）生成 ret void，否则 ret i32 0
        if val is None:
            cur_func: llvmlite.Function | t.CPtr = self.Trans._cur_func
            is_void_ret: int = 0
            if cur_func is not None:
                ret_ty: llvmlite.LLVMType | t.CPtr = llvmlite.function_get_ret_ty(cur_func)
                if ret_ty is not None:
                    match ret_ty:
                        case llvmlite.LLVMType.Void():
                            is_void_ret = 1
            if is_void_ret != 0:
                llvmlite.build_ret_void(builder)
                return 0
            val = llvmlite.const_int32(pool, 0)

        # 类型转换：确保 val 类型与函数返回类型匹配
        # 处理 i1（bool 比较结果）→ i8（t.CBool）等情况
        # i1 → 更宽整数用 zext（bool 语义：1 保持 1，而非 sext 的 0xFF）
        cur_func_rt: llvmlite.Function | t.CPtr = self.Trans._cur_func
        if cur_func_rt is not None and val is not None and val.Ty is not None:
            ret_ty_rt: llvmlite.LLVMType | t.CPtr = llvmlite.function_get_ret_ty(cur_func_rt)
            if ret_ty_rt is not None:
                val_bits: int = HandlesExpr.get_llvm_type_bits(val.Ty)
                ret_bits: int = HandlesExpr.get_llvm_type_bits(ret_ty_rt)
                if val_bits != 0 and ret_bits != 0 and val_bits != ret_bits:
                    if val_bits == 1 and val_bits < ret_bits:
                        val = llvmlite.build_zext(builder, val, ret_ty_rt)
                    else:
                        val = HandlesExpr.coerce_to_type(builder, val, ret_ty_rt)
        llvmlite.build_ret(builder, val)
        return 0


# ============================================================
# NewReturnHandle - 工厂函数
# ============================================================
def NewReturnHandle(pool: memhub.MemBuddy | t.CPtr,
                    trans: HT.Translator | t.CPtr) -> ReturnHandle | t.CPtr:
    h: ReturnHandle | t.CPtr = pool.alloc(ReturnHandle.__sizeof__())
    if h is None:
        return None
    string.memset(h, 0, ReturnHandle.__sizeof__())
    h.Trans = trans
    return h
