import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import viperlib
import lib.core.VLogger as VLogger
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()

        # _declare_only == 0 时，_translate_module_level 已预处理导入语句，
        # 此处跳过避免重复处理（from_imports 条目重复）
        if trans._declare_only == 0 and (kd == ast.ASTKind.Import or kd == ast.ASTKind.ImportFrom):
            continue
        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:
                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


# ============================================================
# _build_array_initializer_text - 从 AST List 生成 LLVM IR 数组初始化器
#
# 生成格式: [N x elem_ty] [elem_ty val0, elem_ty val1, ...]
# 用于模块级 t.CArray 列表字面量初始化（不依赖 builder，直接生成常量初始化器）
# ============================================================
def _build_array_initializer_text(pool: memhub.MemBuddy | t.CPtr,
                                  var_ty: llvmlite.LLVMType | t.CPtr,
                                  list_node: ast.List | t.CPtr) -> str:
    """从 AST List 节点生成 LLVM IR 数组初始化器文本，None 失败"""
    if var_ty is None or list_node is None:
        return None

    # 重新进行类型转换，确保 TPV 编译器正确识别 ast.List 类型
    # （避免 init_list_node 初始值为 None 时类型推断退化为 t.CPtr）
    ln: ast.List | t.CPtr = (ast.List | t.CPtr)(list_node)
    if ln is None:
        return None

    # 从 var_ty 提取数组元素类型和数量
    elem_ty: llvmlite.LLVMType | t.CPtr = None
    arr_count: t.CInt = 0
    match var_ty:
        case llvmlite.LLVMType.Array(et, cnt):
            elem_ty = et
            arr_count = cnt

    if elem_ty is None or arr_count <= 0:
        return None

    # 确定元素类型的 IR 表示
    elem_ir_ty: str = "i8"
    match elem_ty:
        case llvmlite.LLVMType.Int(bits):
            if bits == 8:
                elem_ir_ty = "i8"
            elif bits == 16:
                elem_ir_ty = "i16"
            elif bits == 32:
                elem_ir_ty = "i32"
            elif bits == 64:
                elem_ir_ty = "i64"

    # 获取列表元素（通过 ln.elts 而非 list_node.elts，确保类型正确识别）
    elts: list[ast.AST | t.CPtr] | t.CPtr = ln.elts
    if elts is None:
        return None
    elts_count: t.CSizeT = elts.__len__()

    # 计算缓冲区大小: 每个元素最多 "i32 -9223372036854775808, " 约 26 字符
    buf_size: t.CSizeT = 64 + elts_count * 32
    buf: t.CChar | t.CPtr = pool.alloc(buf_size)
    if buf is None:
        return None

    # 写入前缀: "["（类型前缀由 _print_global 输出，初始化器只需元素列表）
    written: t.CInt = viperlib.snprintf(buf, buf_size, "[")
    pos: t.CSizeT = t.CSizeT(written)

    # 逐个元素写入
    i: t.CSizeT = 0
    while i < elts_count:
        elem_node: ast.AST | t.CPtr = elts.get(i)
        # 直接从 Constant 节点提取 int_val（避免 extract_cdefine_int_value 的 t.CInt 截断）
        elem_val: t.CInt64T = 0
        if elem_node is not None:
            ek: int = elem_node.kind()
            if ek == ast.ASTKind.Constant:
                ec: ast.Constant | t.CPtr = (ast.Constant | t.CPtr)(elem_node)
                if ec is not None and ec.const_kind == ast.CONST_INT:
                    elem_val = ec.int_val
            elif ek == ast.ASTKind.UnaryOp:
                # 负数: UnaryOp(USub, Constant)
                elem_val = HandlesAnnAssign.extract_cdefine_int_value(elem_node)
        if i > 0:
            written = viperlib.snprintf(buf + pos, buf_size - pos, ", ")
            pos += t.CSizeT(written)
        written = viperlib.snprintf(buf + pos, buf_size - pos, "%s %lld", elem_ir_ty, elem_val)
        pos += t.CSizeT(written)
        i += 1

    # 写入后缀 "]"
    viperlib.snprintf(buf + pos, buf_size - pos, "]")
    return buf


# ============================================================
# _build_string_ptr_initializer - 为字符串字面量初始化指针类型全局变量
#
# 创建内部字符串常量全局 @.str.{var_name} = private constant [len+1 x i8] c"...\00"
# 返回 GEP 初始化器文本: getelementptr inbounds ([len+1 x i8], [len+1 x i8]* @.str.{var_name}, i32 0, i32 0)
# 用于 b64_tab: t.CArray[t.CChar, None] = "ABC..." 等字符串初始化指针类型
# ============================================================
def _build_string_ptr_initializer(pool: memhub.MemBuddy | t.CPtr,
                                  mod: llvmlite.LLVMModule | t.CPtr,
                                  var_name: str,
                                  str_val: str) -> str:
    """为字符串字面量创建内部常量全局并返回 GEP 初始化器文本，None 失败"""
    if str_val is None or var_name is None:
        return None

    # 构造内部字符串常量名称: .str.{var_name}
    str_name_buf: t.CChar | t.CPtr = pool.alloc(64)
    if str_name_buf is None:
        return None
    viperlib.snprintf(str_name_buf, 64, ".str.%s", var_name)

    # 调用 llvmlite.create_global_string 创建字符串常量全局
    str_gv: llvmlite.GlobalVariable | t.CPtr = llvmlite.create_global_string(
        pool, mod, str_name_buf, str_val)
    if str_gv is None:
        return None

    # 生成 GEP 初始化器文本
    str_len: t.CSizeT = string.strlen(str_val)
    init_buf: t.CChar | t.CPtr = pool.alloc(128)
    if init_buf is None:
        return None
    viperlib.snprintf(init_buf, 128,
        "getelementptr inbounds ([%d x i8], [%d x i8]* @%s, i32 0, i32 0)",
        str_len + 1, str_len + 1, str_name_buf)
    return init_buf


# ============================================================
# 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
    init_kind: int = 0  # 0=none, 1=int, 2=list, 3=str
    init_str_val: str = None
    init_list_node: ast.List | t.CPtr = None

    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, trans)
            if resolved is not None:
                var_ty = resolved
        # 解析初始值（支持整数、列表字面量、字符串字面量）
        if aa.value is not None:
            val_kind: int = aa.value.kind()
            if val_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
                    init_kind = 1
                elif cn.const_kind == ast.CONST_STR:
                    init_str_val = cn.str_val
                    init_kind = 3
            elif val_kind == ast.ASTKind.List:
                init_list_node = (ast.List | t.CPtr)(aa.value)
                init_kind = 2
    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:
            val_kind2: int = asgn.value.kind()
            if val_kind2 == 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
                    init_kind = 1
                elif cn2.const_kind == ast.CONST_STR:
                    init_str_val = cn2.str_val
                    init_kind = 3
            elif val_kind2 == ast.ASTKind.List:
                init_list_node = (ast.List | t.CPtr)(asgn.value)
                init_kind = 2

    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
    # 先确保 var_ty 中的跨模块结构体类型在模块中有 opaque 声明
    # 否则 @var_name = global %"sha1.ClassName"* zeroinitializer 会报
    # "use of undefined type named 'sha1.ClassName'" 错误
    # （如 _mbuddy: memhub.MemManager | t.CPtr → Ptr(Struct("aeb3...MemManager"))）
    llvmlite.module_ensure_opaque_for_type(mod, pool, var_ty)
    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 init_kind == 1:
        # 整数初始值
        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
    elif init_kind == 2:
        # 列表字面量 → 数组初始化器 [N x elem_ty] [elem_ty val0, ...]
        arr_init: str = _build_array_initializer_text(pool, var_ty, init_list_node)
        if arr_init is not None:
            gv.Initializer = arr_init
        else:
            gv.Initializer = "zeroinitializer"
    elif init_kind == 3:
        # 字符串字面量 → 创建字符串常量全局 + GEP 引用
        str_init: str = _build_string_ptr_initializer(pool, mod, var_name, init_str_val)
        if str_init is not None:
            gv.Initializer = str_init
        else:
            gv.Initializer = "zeroinitializer"
    else:
        # 使用 zeroinitializer 而非 "0"：指针类型必须用 null/zeroinitializer，
        # 整数/聚合类型也兼容 zeroinitializer，避免 "integer constant must have integer type"
        gv.Initializer = "zeroinitializer"

    # 创建 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:
        # 设置 AnnotClassName（联合类型如 Logger|t.CPtr 简化为 Ptr(i8) 时，
        # 属性访问需要通过 AnnotClassName 回退到类名查找结构体）
        if k == ast.ASTKind.AnnAssign:
            aa_ml: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(node)
            if aa_ml is not None and aa_ml.annotation is not None:
                cls_nm_ml: str = HandlesType.extract_class_name_from_annotation(
                    aa_ml.annotation, trans._imported_modules)
                if cls_nm_ml is not None:
                    HandlesVar.set_var_annot_class_name(
                        trans.SymTab, var_name, cls_nm_ml)
        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:
        VLogger.error("CreateFunction main returned NULL", "TR")
        return 0

    entry_blk: llvmlite.BasicBlock | t.CPtr = llvmlite.create_block(
        pool, main_func, "entry")
    if entry_blk is None:
        VLogger.error("CreateBlock returned NULL", "TR")
        return 0

    builder: llvmlite.IRBuilder | t.CPtr = llvmlite.new_builder(pool, main_func)
    if builder is None:
        VLogger.error("NewBuilder returned NULL", "TR")
        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