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


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


@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)
            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)
            if resolved is not None:
                var_ty = resolved

        alloca: llvmlite.Value | t.CPtr = llvmlite.build_alloca(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)
                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)
            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)

        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