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


# ============================================================
# HandlesWhile - while 循环语句处理（Mixin 继承模式）
#
# 翻译 while 语句为 LLVM IR 控制流:
#   br label %cond
# cond:
#   %t = <test>
#   %c = icmp ne i32 %t, 0
#   br i1 %c, label %body, label %end
# body:
#   ... body ...
#   br label %cond
# end:
# ============================================================


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

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

    # ============================================================
    # Handle - 处理 while 语句，返回新增变量数
    # ============================================================
    def Handle(self, node: ast.AST | t.CPtr) -> int:
        """翻译 while 循环语句"""
        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

        while_node: ast.While | t.CPtr = (ast.While | t.CPtr)(node)

        # 1. 创建基本块: cond / body / end（使用 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, "while.cond.%d", cnt)
        cond_bb: llvmlite.BasicBlock | t.CPtr = llvmlite.create_block(pool, func, name_buf)

        viperlib.snprintf(name_buf, 32, "while.body.%d", cnt)
        body_bb: llvmlite.BasicBlock | t.CPtr = llvmlite.create_block(pool, func, name_buf)

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

        # 2. 跳转到 cond 块
        llvmlite.build_br(builder, cond_bb)

        # 3. cond 块: 求值条件，条件分支
        llvmlite.position_at_end(builder, cond_bb)
        cond_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
            builder, pool, trans.Module, while_node.test,
            trans._funcs, trans._func_count, trans)
        if cond_val is None:
            cond_val = llvmlite.const_int32(pool, 0)

        # Compare/Not 表达式已返回 i1，直接使用；其他类型与 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)
        llvmlite.build_cond_br(builder, cond_i1, body_bb, end_bb)

        # 4. body 块: 翻译循环体，跳回 cond
        llvmlite.position_at_end(builder, body_bb)

        # 保存旧循环上下文，设置 break/continue 目标
        old_break: llvmlite.BasicBlock | t.CPtr = trans._break_bb
        old_continue: llvmlite.BasicBlock | t.CPtr = trans._continue_bb
        trans._break_bb = end_bb
        trans._continue_bb = cond_bb

        body: list[ast.AST | t.CPtr] | t.CPtr = while_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)

        # 恢复旧循环上下文
        trans._break_bb = old_break
        trans._continue_bb = old_continue

        if llvmlite.builder_cur_block_is_terminated(builder) == 0:
            llvmlite.build_br(builder, cond_bb)

        # 5. 定位到 end 块
        llvmlite.position_at_end(builder, end_bb)
        return 0


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