import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import viperlib
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.HandlesBody as HandlesBody


# ============================================================
# HandlesIf - if/elif/else 语句处理（Mixin 继承模式）
#
# 翻译 if 语句为 LLVM IR 控制流:
#   br i1 %cond, label %then, label %else
# then:
#   ... body ...
#   br label %end
# else:
#   ... orelse ...
#   br label %end
# end:
# ============================================================


@t.NoVTable
class IfHandle(HandlesBase.Mixin):
    """if/elif/else 语句处理器：继承 Mixin 获得 Trans 回指针"""

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

    # ============================================================
    # Handle - 处理 if 语句，返回新增变量数
    # ============================================================
    def Handle(self, node: ast.AST | t.CPtr) -> int:
        """翻译 if/elif/else 语句"""
        if node is None:
            return 0

        trans: HT.Translator | t.CPtr = self.Trans
        pool: memhub.MemBuddy | t.CPtr = trans.Pool
        builder: llvmlite.IRBuilder | t.CPtr = trans._cur_builder
        func: llvmlite.Function | t.CPtr = trans._cur_func

        if builder is None or func is None:
            return 0

        if_node: ast.If | t.CPtr = (ast.If | t.CPtr)(node)

        # 1. 求值条件表达式
        cond_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
            builder, pool, trans.Module, if_node.test,
            trans._funcs, trans._func_count, trans)

        # 2. 转换为 i1 条件
        # Compare/Not 表达式已返回 i1，直接使用；其他类型与 0 比较
        if cond_val is None:
            cond_val = llvmlite.const_int32(pool, 0)
        cond_bits: int = HandlesExpr.get_llvm_type_bits(cond_val.Ty)
        if cond_bits == 1:
            cond_i1: llvmlite.Value | t.CPtr = cond_val
        else:
            zero: llvmlite.Value | t.CPtr = llvmlite.const_int32(pool, 0)
            cond_i1 = llvmlite.build_icmp(
                builder, llvmlite.ICMP_NE, cond_val, zero)

        # 3. 创建基本块（使用 trans._label_counter 生成唯一标签名，不与 SSA 名共享）
        cnt: int = trans._label_counter
        trans._label_counter = cnt + 1

        name_buf: t.CChar | t.CPtr = pool.alloc(32)
        viperlib.snprintf(name_buf, 32, "if.then.%d", cnt)
        then_bb: llvmlite.BasicBlock | t.CPtr = llvmlite.create_block(pool, func, name_buf)

        viperlib.snprintf(name_buf, 32, "if.end.%d", cnt)
        merge_bb: llvmlite.BasicBlock | t.CPtr = llvmlite.create_block(pool, func, name_buf)

        # 检查是否有 else 分支
        has_else: int = 0
        orelse_list: list[ast.AST | t.CPtr] | t.CPtr = if_node.orelse
        if orelse_list is not None:
            if orelse_list.__len__() > 0:
                has_else = 1

        else_bb: llvmlite.BasicBlock | t.CPtr = None
        if has_else == 1:
            viperlib.snprintf(name_buf, 32, "if.else.%d", cnt)
            else_bb = llvmlite.create_block(pool, func, name_buf)

        # 4. 发射条件分支
        if has_else == 1:
            llvmlite.build_cond_br(builder, cond_i1, then_bb, else_bb)
        else:
            llvmlite.build_cond_br(builder, cond_i1, then_bb, merge_bb)

        # 5. 翻译 then body
        llvmlite.position_at_end(builder, then_bb)
        body: list[ast.AST | t.CPtr] | t.CPtr = if_node.children
        if body is not None:
            body_count: t.CSizeT = body.__len__()
            for bi in range(body_count):
                stmt: ast.AST | t.CPtr = body.get(bi)
                if stmt is not None:
                    HandlesBody.translate_stmt(trans, stmt)

        # then 块未终止则跳到 merge
        if llvmlite.builder_cur_block_is_terminated(builder) == 0:
            llvmlite.build_br(builder, merge_bb)

        # 6. 翻译 else body（若有）
        if has_else == 1:
            llvmlite.position_at_end(builder, else_bb)
            else_count: t.CSizeT = orelse_list.__len__()
            for ei in range(else_count):
                stmt: ast.AST | t.CPtr = orelse_list.get(ei)
                if stmt is not None:
                    HandlesBody.translate_stmt(trans, stmt)
            if llvmlite.builder_cur_block_is_terminated(builder) == 0:
                llvmlite.build_br(builder, merge_bb)

        # 7. 定位到 merge 块继续后续代码
        llvmlite.position_at_end(builder, merge_bb)
        return 0


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