import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import viperlib
import lib.core.Handles.HandlesTranslator as HT
import lib.core.Handles.HandlesVar as HandlesVar
import lib.core.Handles.HandlesExprCall as HandlesExprCall
import lib.core.Handles.HandlesImports as HandlesImports
import lib.core.Handles.HandlesBody as HandlesBody
import lib.core.Handles.HandlesFunctions as HandlesFunctions
import lib.core.Handles.HandlesType as HandlesType
import lib.core.Handles.HandlesClassDef as HandlesClassDef
import lib.core.Handles.HandlesAnnAssign as HandlesAnnAssign


# ============================================================
# HandlesMain - 模块级翻译入口 + wrapper main 创建
#
# 从 translator.py 拆分出来，负责:
#   1. create_wrapper_main() - 无用户 main 时创建包装 main
#   2. translate_children() - 遍历 AST 子节点并分派翻译
#
# trans 单参模式：所有共享状态从 trans 获取，无需 11 个参数
# 注意: str = bytes = t.CChar | t.CPtr = i8*
# ============================================================


# ============================================================
# _register_cexport_from_funcdef - Phase 1a 预注册 CExport/State 函数
#
# 解决翻译顺序依赖问题：Phase 1b 按字母序翻译，后翻译的模块的
# CExport 函数无法被先翻译的模块正确识别为裸名调用。
# Phase 1a 预注册所有 CExport/State 函数名到全局表。
# ============================================================
def _register_cexport_from_funcdef(trans: HT.Translator | t.CPtr,
                                   node: ast.AST | t.CPtr) -> int:
    """预注册 CExport/State 函数到全局表（仅注册，不生成 IR）"""
    fd: ast.FunctionDef | t.CPtr = (ast.FunctionDef | t.CPtr)(node)
    if fd is None or fd.name is None:
        return 0
    if trans is None or trans.ModuleSha1 is None:
        return 0

    # 检查返回类型是否有 CExport 或 State 标记
    has_export: int = 0
    has_state: int = 0
    if fd.returns is not None:
        has_export = HandlesType.has_decorator_marker(fd.returns, "CExport")
        has_state = HandlesType.has_decorator_marker(fd.returns, "State")

    if has_export != 0 or has_state != 0:
        HandlesExprCall.register_cexport_func(trans.ModuleSha1, fd.name)

    return 0


# ============================================================
# translate_children - 遍历 AST 子节点并分派翻译
#
# 对应 TransPyC translator._translate_children()
# 共享状态从 trans 获取，imported_modules/from_imports 更新到 trans
# ============================================================
def translate_children(trans: HT.Translator | t.CPtr,
                       tree: ast.AST | t.CPtr) -> int:
    """遍历 tree.children 并翻译每个子节点

    Args:
        trans: 翻译器（含所有共享状态）
        tree: AST 模块节点

    Returns:
        t.CInt: 新增的变量数
    """
    ch: list[ast.AST | t.CPtr] | t.CPtr = tree.children
    if ch is None:
        return 0

    cn_count: t.CSizeT = ch.__len__()
    added_total: int = 0

    for ci in range(cn_count):
        child: ast.AST | t.CPtr = ch.get(ci)
        if child is None: continue
        kd: int = child.kind()

        if kd == ast.ASTKind.Import:
            trans.ImportsH.HandleImport(child)
        elif kd == ast.ASTKind.ImportFrom:
            trans.ImportsH.HandleImportFromModule(child)
            trans.ImportsH.HandleImportFromNames(child)
        elif kd == ast.ASTKind.FunctionDef:
            # Phase 1a 声明模式：只注册 CExport/State 函数到全局表（解决翻译顺序依赖）
            # Phase 1b 全量翻译：正常翻译函数体
            if trans._declare_only == 0:
                added: int = HandlesFunctions.translate_function_def(trans, child)
                added_total += added
            elif trans._declare_only == 1:
                _register_cexport_from_funcdef(trans, child)
        elif kd == ast.ASTKind.ClassDef:
            # ClassDef 在模块级直接处理（不需要 builder）
            # _declare_only=2（import扫描模式）时跳过，只处理 import 依赖
            if trans._declare_only != 2:
                cd_node: ast.ClassDef | t.CPtr = (ast.ClassDef | t.CPtr)(child)
                cd_name: str = "?" if cd_node is None or cd_node.name is None else cd_node.name
                HandlesClassDef.translate_class_def(trans, child)
        elif trans._declare_only == 0 and trans._cur_builder is not None:
            # 有 builder → 委托 HandlesBody 分派
            added = HandlesBody.translate_stmt(trans, child)
            added_total += added
        elif kd == ast.ASTKind.AnnAssign and trans._declare_only != 2:
            # 模块级 AnnAssign
            # _declare_only=2（import扫描模式）时跳过
            # _declare_only=1（struct注册模式）时只处理 CDefine（在 handle_module_level_var 内部判断）
            # _declare_only=0（全量翻译）时处理所有模块级 AnnAssign
            added = handle_module_level_var(trans, child)
            added_total += added
        elif trans._declare_only == 0 and kd == ast.ASTKind.Assign:
            # 无 builder 的模块级 Assign → 创建全局变量（仅全量翻译模式）
            added = handle_module_level_var(trans, child)
            added_total += added

    return added_total


# ============================================================
# handle_module_level_var - 模块级变量声明 → 创建 LLVM 全局变量
#
# 当用户已定义 main（无 wrapper main builder）时，模块级
# AnnAssign/Assign 创建全局变量 @var_name 并注册到 SymTab 模块作用域
# ============================================================
def handle_module_level_var(trans: HT.Translator | t.CPtr,
                            node: ast.AST | t.CPtr) -> int:
    """处理模块级变量声明，创建全局变量"""
    pool: memhub.MemBuddy | t.CPtr = trans.Pool
    mod: llvmlite.LLVMModule | t.CPtr = trans.Module

    # CDefine 注解: 编译期常量，不创建全局变量
    # 注册到全局 CDefine 表供 t.CArray[elem_ty, NAME] 解析
    k: int = node.kind()
    if k == ast.ASTKind.AnnAssign:
        aa_cd: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(node)
        if aa_cd is not None and aa_cd.target is not None:
            if aa_cd.target.kind() == ast.ASTKind.Name:
                nm_cd: ast.Name | t.CPtr = (ast.Name | t.CPtr)(aa_cd.target)
                if nm_cd.id is not None:
                    if HandlesAnnAssign.is_cdefine_annotation(aa_cd.annotation) != 0:
                        val_cd: int = HandlesAnnAssign.extract_cdefine_int_value(aa_cd.value)
                        HandlesType.register_cdefine_constant(pool, nm_cd.id, val_cd)
                        return 0

    # _declare_only=1（struct注册模式）时只处理 CDefine，不创建全局变量
    # CDefine 已在上面处理并返回，到这里说明不是 CDefine，直接跳过
    if trans._declare_only == 1:
        return 0

    var_name: str = None
    var_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
    init_val: t.CInt64T = 0
    has_init: int = 0

    if k == ast.ASTKind.AnnAssign:
        aa: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(node)
        if aa is None or aa.target is None:
            return 0
        if aa.target.kind() != ast.ASTKind.Name:
            return 0
        nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(aa.target)
        var_name = nm.id
        # 解析类型
        if aa.annotation is not None:
            resolved: llvmlite.LLVMType | t.CPtr = HandlesType.resolve_annotation_type(
                pool, aa.annotation, trans._imported_modules, trans._from_imports)
            if resolved is not None:
                var_ty = resolved
        # 解析初始值
        if aa.value is not None and aa.value.kind() == ast.ASTKind.Constant:
            cn: ast.Constant | t.CPtr = (ast.Constant | t.CPtr)(aa.value)
            if cn.const_kind == ast.CONST_INT:
                init_val = cn.int_val
                has_init = 1
    elif k == ast.ASTKind.Assign:
        asgn: ast.Assign | t.CPtr = (ast.Assign | t.CPtr)(node)
        if asgn is None or asgn.targets is None:
            return 0
        targets: list[ast.AST | t.CPtr] | t.CPtr = asgn.targets
        if targets.__len__() < 1:
            return 0
        t0: ast.AST | t.CPtr = targets.get(0)
        if t0 is None or t0.kind() != ast.ASTKind.Name:
            return 0
        nm2: ast.Name | t.CPtr = (ast.Name | t.CPtr)(t0)
        var_name = nm2.id
        if asgn.value is not None and asgn.value.kind() == ast.ASTKind.Constant:
            cn2: ast.Constant | t.CPtr = (ast.Constant | t.CPtr)(asgn.value)
            if cn2.const_kind == ast.CONST_INT:
                init_val = cn2.int_val
                has_init = 1

    if var_name is None:
        return 0

    # 检查是否已注册
    existing: llvmlite.Value | t.CPtr = HandlesVar.lookup_module_var(
        trans.SymTab, var_name)
    if existing is not None:
        return 0

    # 创建全局变量 @var_name
    gv: llvmlite.GlobalVariable | t.CPtr = llvmlite.new_global_variable(pool, var_name, var_ty)
    if gv is None:
        return 0
    llvmlite.module_add_global(mod, gv)

    # 设置初始值（有初始值时清除 external linkage，因为 LLVM 22+ 不允许 external global 带初始值）
    gv.Linkage = None
    if has_init != 0:
        init_buf: t.CChar | t.CPtr = pool.alloc(48)
        if init_buf is not None:
            viperlib.snprintf(init_buf, 48, "%lld", init_val)
            gv.Initializer = init_buf
    else:
        gv.Initializer = "0"

    # 创建 Value 引用（@var_name, 类型为 var_ty*）
    var_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, var_ty)
    ref_name: t.CChar | t.CPtr = pool.alloc(64)
    if ref_name is not None:
        viperlib.snprintf(ref_name, 64, "@%s", var_name)
    gv_ref: llvmlite.Value | t.CPtr = llvmlite.SSAValue(pool, var_ptr_ty, ref_name)

    # 注册到模块作用域
    if HandlesVar.define_module_var(trans.SymTab, var_name, gv_ref) == 0:
        return 1
    return 0


# ============================================================
# create_wrapper_main - 创建包装 main 函数
#
# 当用户未定义 main 函数时调用。
# 创建 main() -> i32 函数 → 设置 trans._cur_func/_cur_builder →
# 预扫描 alloca → 翻译子节点 → ret 0
# ============================================================
def create_wrapper_main(trans: HT.Translator | t.CPtr,
                        tree: ast.AST | t.CPtr,
                        i32_ty: llvmlite.LLVMType | t.CPtr) -> int:
    """创建包装 main 函数并翻译所有顶层语句

    Args:
        trans: 翻译器（含所有共享状态）
        tree: AST 模块节点
        i32_ty: i32 LLVMType

    Returns:
        t.CInt: 新增的变量数
    """
    pool: memhub.MemBuddy | t.CPtr = trans.Pool
    mod: llvmlite.LLVMModule | t.CPtr = trans.Module

    main_func: llvmlite.Function | t.CPtr = llvmlite.create_function(
        pool, mod, "main", i32_ty)
    if main_func is None:
        stdio.printf("[TR] CreateFunction main returned NULL\n")
        return 0

    entry_blk: llvmlite.BasicBlock | t.CPtr = llvmlite.create_block(
        pool, main_func, "entry")
    if entry_blk is None:
        stdio.printf("[TR] CreateBlock returned NULL\n")
        return 0

    builder: llvmlite.IRBuilder | t.CPtr = llvmlite.new_builder(pool, main_func)
    if builder is None:
        stdio.printf("[TR] NewBuilder returned NULL\n")
        return 0
    llvmlite.position_at_end(builder, entry_blk)

    # 设置当前翻译上下文（供 Handle 通过 self.Trans._cur_builder 访问）
    trans._cur_func = main_func
    trans._cur_builder = builder

    # 预处理模块级变量：为 AnnAssign/Assign 创建 LLVM 全局变量
    # 必须在 pre_scan_allocas 之前执行，否则 PreScan 会创建局部 alloca，
    # 导致其他函数通过 SSA 编号引用 wrapper main 的局部变量（无效 IR）
    ch: list[ast.AST | t.CPtr] | t.CPtr = tree.children
    if ch is not None:
        cn_count: t.CSizeT = ch.__len__()
        for ci in range(cn_count):
            child: ast.AST | t.CPtr = ch.get(ci)
            if child is not None:
                kd: int = child.kind()
                if kd == ast.ASTKind.AnnAssign or kd == ast.ASTKind.Assign:
                    handle_module_level_var(trans, child)

    # 预扫描顶层语句：为函数内 AnnAssign 提前创建 alloca
    # （模块级变量已在上面注册到模块作用域，PreScan 的 lookup_current 会跳过它们）
    if ch is not None:
        for ci in range(cn_count):
            child: ast.AST | t.CPtr = ch.get(ci)
            if child is not None:
                HandlesBody.pre_scan_allocas(trans, child)

    # 翻译子节点
    added_total: int = translate_children(trans, tree)

    # 返回 0
    zero_val: llvmlite.Value | t.CPtr = llvmlite.const_int32(pool, 0)
    llvmlite.build_ret(builder, zero_val)

    return added_total