import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import lib.core.Handles.HandlesBase as HandlesBase
import lib.core.Handles.HandlesTranslator as HT
import lib.core.Handles.HandlesVar as HandlesVar
import lib.core.Handles.HandlesExpr as HandlesExpr
import lib.core.Handles.HandlesType as HandlesType
import lib.core.Handles.HandlesNonlocal as HandlesNonlocal
import lib.core.Handles.HandlesClassDef as HandlesClassDef
import lib.core.Handles.HandlesExprCall as HandlesExprCall
import lib.core.Handles.HandlesStruct as HandlesStruct


# ============================================================
# HandlesAnnAssign - AnnAssign 语句处理（Mixin 继承模式）
# ============================================================


# ============================================================
# is_cdefine_annotation - 检测注解是否为 t.CDefine（模块级函数）
#
# 支持两种形式:
#   1. Attribute(Name('t'), 'CDefine')
#   2. BinOp(... | Attribute(Name('t'), 'CDefine'))（联合注解）
# ============================================================
def is_cdefine_annotation(annot: ast.AST | t.CPtr) -> int:
    """检测注解是否为 t.CDefine，返回 1 表示是，0 表示否"""
    if annot is None:
        return 0
    k: int = annot.kind()
    if k == ast.ASTKind.Attribute:
        at: ast.Attribute | t.CPtr = (ast.Attribute | t.CPtr)(annot)
        if at.attr is not None and string.strcmp(at.attr, "CDefine") == 0:
            return 1
        return 0
    if k == ast.ASTKind.BinOp:
        bop: ast.BinOp | t.CPtr = (ast.BinOp | t.CPtr)(annot)
        if is_cdefine_annotation(bop.left) != 0:
            return 1
        if is_cdefine_annotation(bop.right) != 0:
            return 1
    return 0


# ============================================================
# extract_cdefine_int_value - 从 AnnAssign.value 提取整数常量（模块级函数）
#
# 支持的表达式形式:
#   1. Constant(INT)         — 如 0x0002, 42
#   2. BinOp(BitOr/BitAnd)   — 如 FOREGROUND_RED | FOREGROUND_GREEN
#   3. Name                  — 引用已注册的 CDefine 常量
#   4. Call                  — 如 t.CUnsignedLong(-11) → 取第一个参数
#   5. UnaryOp(USub/UAdd/Invert) — 如 -11
# ============================================================
def extract_cdefine_int_value(val_node: ast.AST | t.CPtr) -> int:
    """从值节点提取整数常量（支持 Constant/BinOp/Name/Call/UnaryOp）"""
    if val_node is None:
        return 0
    k: int = val_node.kind()

    # Case 1: Constant(INT) — 如 0x0002
    if k == ast.ASTKind.Constant:
        cn: ast.Constant | t.CPtr = (ast.Constant | t.CPtr)(val_node)
        if cn.const_kind != ast.CONST_INT:
            return 0
        return cn.int_val

    # Case 2: BinOp — 如 FOREGROUND_RED | FOREGROUND_GREEN
    if k == ast.ASTKind.BinOp:
        bop: ast.BinOp | t.CPtr = (ast.BinOp | t.CPtr)(val_node)
        left_val: int = extract_cdefine_int_value(bop.left)
        right_val: int = extract_cdefine_int_value(bop.right)
        if bop.op == ast.OpKind.BitOr:
            return left_val | right_val
        if bop.op == ast.OpKind.BitAnd:
            return left_val & right_val
        if bop.op == ast.OpKind.Add:
            return left_val + right_val
        if bop.op == ast.OpKind.Sub:
            return left_val - right_val
        return 0

    # Case 3: Name — 引用已注册的 CDefine 常量
    if k == ast.ASTKind.Name:
        nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(val_node)
        if nm.id is not None:
            looked_up: int = HandlesType.lookup_cdefine_constant(nm.id)
            if HandlesType.is_cdefine_found() != 0:
                return looked_up
        return 0

    # Case 4: Call — 如 t.CUnsignedLong(-11)
    if k == ast.ASTKind.Call:
        cl: ast.Call | t.CPtr = (ast.Call | t.CPtr)(val_node)
        if cl.args is not None and cl.args.__len__() > 0:
            first_arg: ast.AST | t.CPtr = cl.args.get(0)
            arg_val: int = extract_cdefine_int_value(first_arg)
            return arg_val
        return 0

    # Case 5: UnaryOp — 如 -11
    if k == ast.ASTKind.UnaryOp:
        uop: ast.UnaryOp | t.CPtr = (ast.UnaryOp | t.CPtr)(val_node)
        operand_val: int = extract_cdefine_int_value(uop.operand)
        if uop.op == ast.OpKind.USub:
            return -operand_val
        if uop.op == ast.OpKind.UAdd:
            return operand_val
        if uop.op == ast.OpKind.Invert:
            return ~operand_val
        return 0

    return 0


# ============================================================
# _annotation_contains_name - 递归检查注解中是否包含指定名称
#
# 支持 BinOp(BitOr) 递归，同时检查 Attribute.attr 和 Name.id。
# ============================================================
def _annotation_contains_name(annot: ast.AST | t.CPtr, name: str) -> int:
    """递归检查注解中是否包含指定名称，返回 1=包含 / 0=不包含"""
    if annot is None or name is None:
        return 0
    k: int = annot.kind()
    if k == ast.ASTKind.Attribute:
        at: ast.Attribute | t.CPtr = (ast.Attribute | t.CPtr)(annot)
        if at.attr is not None and string.strcmp(at.attr, name) == 0:
            return 1
        return 0
    if k == ast.ASTKind.Name:
        nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(annot)
        if nm.id is not None and string.strcmp(nm.id, name) == 0:
            return 1
        return 0
    if k == ast.ASTKind.BinOp:
        bop: ast.BinOp | t.CPtr = (ast.BinOp | t.CPtr)(annot)
        if _annotation_contains_name(bop.left, name) != 0:
            return 1
        if _annotation_contains_name(bop.right, name) != 0:
            return 1
    return 0


# ============================================================
# _is_ptr_element_annotation - 检查注解是否为 bytes|t.CPtr 或 str|t.CPtr
#
# 返回 1=是（下标按 8 字节步长, i8** 语义）, 0=否（下标按 1 字节步长, i8* 语义）
#
# 规则:
#   bytes|t.CPtr / str|t.CPtr       → i8**（8 字节步长）→ 返回 1
#   t.CChar|t.CPtr / t.CInt8T|t.CPtr → i8*（1 字节步长）→ 返回 0
#   纯 bytes / str                  → i8*（1 字节步长）→ 返回 0
#
# 注: 注解中的 "|" 是"和"的意思（组合类型），非右值中的位或运算
# ============================================================
def _is_ptr_element_annotation(annot: ast.AST | t.CPtr) -> int:
    """检查注解是否为 bytes|t.CPtr 或 str|t.CPtr 形式"""
    if annot is None:
        return 0
    k: int = annot.kind()
    if k != ast.ASTKind.BinOp:
        return 0
    bop: ast.BinOp | t.CPtr = (ast.BinOp | t.CPtr)(annot)
    if bop.op != ast.OpKind.BitOr:
        return 0
    if _annotation_contains_name(annot, "str") != 0:
        return 1
    if _annotation_contains_name(annot, "bytes") != 0:
        return 1
    return 0


# ============================================================
# _init_global_array_from_list — 从列表字面量初始化全局数组
#
# 为每个元素生成 GEP + store 指令，将值存入全局数组的对应位置。
# 适用于模块级 t.CArray[elem_ty, count] = [v0, v1, ...] 的初始化。
#
# Args:
#   builder: IRBuilder
#   pool: 编译器内存池
#   mod: LLVMModule
#   array_alloca: 全局数组变量的 Value 引用（类型为 [N x elem_ty]*）
#   list_node: ast.List 节点
#   trans: Translator 对象
#
# Returns:
#   0 成功，1 失败
# ============================================================
def _init_global_array_from_list(builder: llvmlite.IRBuilder | t.CPtr,
                                 pool: memhub.MemBuddy | t.CPtr,
                                 mod: llvmlite.LLVMModule | t.CPtr,
                                 array_alloca: llvmlite.Value | t.CPtr,
                                 list_node: ast.AST | t.CPtr,
                                 trans: HT.Translator | t.CPtr) -> int:
    """从列表字面量初始化全局数组，为每个元素生成 GEP + store"""
    if builder is None or array_alloca is None or list_node is None:
        return 1

    # 获取数组类型（Pointee of [N x elem_ty]*）
    arr_ty: llvmlite.LLVMType | t.CPtr = None
    if array_alloca.Ty is not None:
        arr_ty = array_alloca.Ty.Pointee
    if arr_ty is None:
        return 1

    # 匹配数组类型获取元素类型
    elem_ty: llvmlite.LLVMType | t.CPtr = None
    match arr_ty:
        case llvmlite.LLVMType.Array(et, _):
            elem_ty = et
    if elem_ty is None:
        return 1

    # 获取列表元素
    lst: ast.List | t.CPtr = (ast.List | t.CPtr)(list_node)
    if lst is None or lst.elts is None:
        return 1
    elts: list[ast.AST | t.CPtr] | t.CPtr = lst.elts
    elts_count: t.CSizeT = elts.__len__()

    # 为每个元素生成 GEP + store
    i: t.CSizeT = 0
    while i < elts_count:
        elem_node: ast.AST | t.CPtr = elts.get(i)
        elem_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
            builder, pool, mod, elem_node, None, 0, trans)
        if elem_val is not None:
            idx_val: llvmlite.Value | t.CPtr = llvmlite.const_int32(pool, i)
            elem_ptr: llvmlite.Value | t.CPtr = llvmlite.build_gep_array(
                builder, arr_ty, elem_ty, array_alloca, idx_val)
            if elem_ptr is not None:
                # 类型转换（如 i32 → i8 截断）
                store_val: llvmlite.Value | t.CPtr = elem_val
                if elem_ptr.Ty is not None and elem_ptr.Ty.Pointee is not None:
                    store_val = HandlesExpr.coerce_to_type(
                        builder, elem_val, elem_ptr.Ty.Pointee)
                llvmlite.build_store(builder, store_val, elem_ptr)
        i += 1

    return 0


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

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

    # ============================================================
    # PreScan - 预扫描 AnnAssign，为有类型注解的变量提前创建 alloca
    #
    # 返回新增的变量数
    # ============================================================
    def PreScan(self, node: ast.AST | t.CPtr) -> int:
        """预扫描 AnnAssign，提前创建 alloca 到 entry block 顶部"""
        if node is None:
            return 0
        k: int = node.kind()
        if k != ast.ASTKind.AnnAssign:
            return 0

        aa: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(node)
        if aa is None or aa.target is None:
            return 0
        target: ast.AST | t.CPtr = aa.target
        if target.kind() != ast.ASTKind.Name:
            return 0

        nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(target)
        if nm.id is None:
            return 0

        # CDefine 注解: 不创建 alloca，但注册常量到全局表
        # 供后续 t.CArray[elem_ty, NAME] 编译期解析使用
        if is_cdefine_annotation(aa.annotation) != 0:
            pool_ps: memhub.MemBuddy | t.CPtr = self.Trans.Pool
            val_ps: int = extract_cdefine_int_value(aa.value)
            HandlesType.register_cdefine_constant(pool_ps, nm.id, val_ps)
            return 0

        # global/nonlocal 变量不需要局部 alloca
        if HT.is_global_name(self.Trans, nm.id) != 0:
            return 0
        if HT.is_nonlocal_name(self.Trans, nm.id) != 0:
            return 0

        # 模块级全局变量已由 handle_module_level_var 注册到模块作用域，
        # 不需要创建局部 alloca（否则会导致局部变量遮蔽全局变量）
        if HandlesVar.lookup_module_var(self.Trans.SymTab, nm.id) is not None:
            return 0

        # 检查是否已存在
        existing: llvmlite.Value | t.CPtr = HandlesVar.lookup_current(
            self.Trans.SymTab, nm.id)
        if existing is not None:
            return 0

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

        # 从 annotation 推断类型（修复：原硬编码 i32 导致 str 等类型错误）
        var_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
        annot: ast.AST | t.CPtr = aa.annotation
        if annot is not None:
            resolved: llvmlite.LLVMType | t.CPtr = HandlesType.resolve_annotation_type(
                pool, annot, self.Trans._imported_modules, self.Trans._from_imports, self.Trans)
            if resolved is None:
                # 尝试特化泛型类注解（如 list[str]）
                # resolve_annotation_type 返回 None 可能是因为泛型类未特化
                if annot.kind() == ast.ASTKind.Subscript:
                    sub_annot: ast.Subscript | t.CPtr = (ast.Subscript | t.CPtr)(annot)
                    if sub_annot.value is not None and sub_annot.value.kind() == ast.ASTKind.Name:
                        sub_nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(sub_annot.value)
                        if sub_nm.id is not None:
                            template_cd: ast.ClassDef | t.CPtr = HandlesClassDef._find_generic_template(sub_nm.id)
                            if template_cd is not None:
                                type_args_ps: list[str] | t.CPtr = HandlesExprCall._extract_type_args_from_slice(pool, sub_annot.slice)
                                if type_args_ps is not None and type_args_ps.__len__() > 0:
                                    spec_name_ps: str = HandlesClassDef._specialize_generic_class(
                                        self.Trans, sub_nm.id, type_args_ps)
                                    if spec_name_ps is not None:
                                        resolved = HandlesType.resolve_annotation_type(
                                            pool, annot, self.Trans._imported_modules, self.Trans._from_imports, self.Trans)
            if resolved is not None:
                var_ty = resolved

        # 确保跨模块结构体的完整定义在当前模块中可用（供 alloca 分配空间）
        HandlesStruct.ensure_struct_def_in_module(pool, self.Trans.Module, var_ty)
        alloca: llvmlite.Value | t.CPtr = HandlesVar._alloca_at_entry(builder, var_ty)
        if alloca is not None:
            if HandlesVar.define_var(
                    self.Trans.SymTab, nm.id, alloca) == 0:
                # 存储原始类型注解的类名（方法调用检测时，Ptr(i8) 回退到类名查找结构体）
                if annot is not None:
                    cls_nm_ps: str = HandlesType.extract_class_name_from_annotation(
                        annot, self.Trans._imported_modules)
                    if cls_nm_ps is not None:
                        HandlesVar.set_var_annot_class_name(
                            self.Trans.SymTab, nm.id, cls_nm_ps)
                    # 检查是否为 bytes|t.CPtr 或 str|t.CPtr（i8** 语义，8 字节步长）
                    if _is_ptr_element_annotation(annot) != 0:
                        HandlesVar.set_var_ptr_element(
                            self.Trans.SymTab, nm.id)
                return 0
        return 0

    # ============================================================
    # Handle - 翻译 AnnAssign 语句
    #
    # 返回新增的变量数
    # ============================================================
    def Handle(self, node: ast.AST | t.CPtr) -> int:
        """翻译 AnnAssign(target=Name, annotation=..., value=expr)"""
        aa: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(node)
        if aa is None:
            return 0

        target: ast.AST | t.CPtr = aa.target
        if target is None or target.kind() != ast.ASTKind.Name:
            return 0

        nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(target)
        if nm.id is None:
            return 0

        # CDefine 注解: 编译期常量，不生成运行时代码
        # 注册到全局 CDefine 表供 t.CArray[elem_ty, NAME] 解析
        # 注: PreScan 已注册过，此处保证即使跳过 PreScan 也能正确注册
        if is_cdefine_annotation(aa.annotation) != 0:
            pool_cd: memhub.MemBuddy | t.CPtr = self.Trans.Pool
            val_cd: int = extract_cdefine_int_value(aa.value)
            HandlesType.register_cdefine_constant(pool_cd, nm.id, val_cd)
            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

        # 确定类型（从 annotation 推断）
        var_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
        annot: ast.AST | t.CPtr = aa.annotation
        if annot is not None:
            resolved: llvmlite.LLVMType | t.CPtr = HandlesType.resolve_annotation_type(
                pool, annot, self.Trans._imported_modules, self.Trans._from_imports, self.Trans)
            if resolved is not None:
                var_ty = resolved

        # global 变量：写入模块作用域中的全局变量
        if HT.is_global_name(self.Trans, nm.id) != 0:
            mod_alloca: llvmlite.Value | t.CPtr = HandlesVar.lookup_module_var(
                self.Trans.SymTab, nm.id)
            if mod_alloca is not None and aa.value is not None:
                rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                    builder, pool, mod, aa.value,
                    None, 0, self.Trans)
                if rhs_val is not None:
                    target_ty: llvmlite.LLVMType | t.CPtr = None
                    if mod_alloca.Ty is not None:
                        target_ty = mod_alloca.Ty.Pointee
                    if target_ty is not None:
                        rhs_val = HandlesExpr.coerce_to_type(builder, rhs_val, target_ty)
                    llvmlite.build_store(builder, rhs_val, mod_alloca)
            return 0

        # nonlocal 变量：通过闭包 env 写入
        if HT.is_nonlocal_name(self.Trans, nm.id) != 0:
            if aa.value is not None:
                rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                    builder, pool, mod, aa.value,
                    None, 0, self.Trans)
                if rhs_val is not None:
                    nl_ptr: llvmlite.Value | t.CPtr = HandlesNonlocal.get_nonlocal_var_ptr(
                        self.Trans, nm.id)
                    if nl_ptr is not None:
                        i32_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
                        rhs_coerced: llvmlite.Value | t.CPtr = HandlesExpr.coerce_to_type(
                            builder, rhs_val, i32_ty)
                        llvmlite.build_store(builder, rhs_coerced, nl_ptr)
            return 0

        # 模块级全局变量（由 handle_module_level_var 注册到模块作用域）
        # 非函数内 global 声明，但变量已在模块作用域中（如模块级 t.CArray 初始化）
        mod_glob: llvmlite.Value | t.CPtr = HandlesVar.lookup_module_var(
            self.Trans.SymTab, nm.id)
        if mod_glob is not None:
            if aa.value is not None:
                # 列表字面量 → 逐元素 GEP + store 初始化数组
                if aa.value.kind() == ast.ASTKind.List:
                    _init_global_array_from_list(
                        builder, pool, mod, mod_glob, aa.value, self.Trans)
                else:
                    rhs_val_mg: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                        builder, pool, mod, aa.value,
                        None, 0, self.Trans)
                    if rhs_val_mg is not None:
                        target_ty_mg: llvmlite.LLVMType | t.CPtr = None
                        if mod_glob.Ty is not None:
                            target_ty_mg = mod_glob.Ty.Pointee
                        if target_ty_mg is not None:
                            rhs_val_mg = HandlesExpr.coerce_to_type(builder, rhs_val_mg, target_ty_mg)
                        llvmlite.build_store(builder, rhs_val_mg, mod_glob)
            return 0

        # 普通局部变量
        # 创建 alloca
        alloca: llvmlite.Value | t.CPtr = HandlesVar.get_or_create_sym(
            self.Trans.SymTab, pool, builder, nm.id, var_ty)
        if alloca is None:
            return 0
        # 存储原始类型注解的类名（方法调用检测时，Ptr(i8) 回退到类名查找结构体）
        if annot is not None:
            cls_nm_hd: str = HandlesType.extract_class_name_from_annotation(
                annot, self.Trans._imported_modules)
            if cls_nm_hd is not None:
                HandlesVar.set_var_annot_class_name(
                    self.Trans.SymTab, nm.id, cls_nm_hd)
            # 检查是否为 bytes|t.CPtr 或 str|t.CPtr（i8** 语义，8 字节步长）
            if _is_ptr_element_annotation(annot) != 0:
                HandlesVar.set_var_ptr_element(
                    self.Trans.SymTab, nm.id)

        new_vars: int = 0
        existing: llvmlite.Value | t.CPtr = HandlesVar.lookup_current(
            self.Trans.SymTab, nm.id)
        if existing is None:
            new_vars = 1

        # 如果有初始值，store
        if aa.value is not None:
            rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                builder, pool, mod, aa.value,
                None, 0, self.Trans)
            if rhs_val is not None:
                # 类型转换：将 rhs_val 转换为 alloca 的 pointee 类型
                # 修复：整数字面量是 i32，但变量可能是 i8/i16/i64，需 trunc/sext
                target_ty: llvmlite.LLVMType | t.CPtr = None
                if alloca.Ty is not None:
                    target_ty = alloca.Ty.Pointee
                if target_ty is not None:
                    rhs_val = HandlesExpr.coerce_to_type(builder, rhs_val, target_ty)
                llvmlite.build_store(builder, rhs_val, alloca)

        return 0


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