import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
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 提取整数常量（模块级函数）
#
# 仅支持 ast.Constant(INT)，其他形式返回 0
# ============================================================
def extract_cdefine_int_value(val_node: ast.AST | t.CPtr) -> int:
    """从值节点提取整数常量（仅支持 Constant INT）"""
    if val_node is None:
        return 0
    if val_node.kind() != ast.ASTKind.Constant:
        return 0
    cn: ast.Constant | t.CPtr = (ast.Constant | t.CPtr)(val_node)
    if cn.const_kind != ast.CONST_INT:
        return 0
    return cn.int_val


# ============================================================
# _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


@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

        # 检查是否已存在
        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

        # 普通局部变量
        # 创建 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