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.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.HandlesExprCall as HandlesExprCall
import lib.core.Handles.HandlesNonlocal as HandlesNonlocal
import lib.core.Handles.HandlesType as HandlesType
import lib.core.Handles.HandlesStruct as HandlesStruct
import lib.core.StubMerger as StubMerger


# ============================================================
# HandlesAssign - Assign 语句处理（Mixin 继承模式）
#
# 对应 TransPyC 的 class AssignHandle(BaseHandle):
#   @t.NoVTable 继承 Mixin 获得 Trans 字段（展平嵌入，无 vtable）
#   通过 self.Trans 访问共享状态（Pool/Module/_cur_builder/SymTab/...）
#   通过 self.Trans.ExprH / self.Trans.IfH 等访问其他 Handle
# ============================================================


@t.NoVTable
class AssignHandle(HandlesBase.Mixin):
    """Assign 语句处理器：继承 Mixin 获得 Trans 回指针 + 共享方法"""
    _CurrentClass: str               # 模块私有状态

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

    # ============================================================
    # Handle - 处理 Assign 语句，返回新增变量数（0 或 1）
    #
    # 对应 TransPyC AssignHandle._HandleAssignLlvm
    # 共享状态从 self.Trans 获取，无需 11 个参数
    # ============================================================
    def Handle(self, node: ast.AST | t.CPtr) -> int:
        asgn: ast.Assign | t.CPtr = (ast.Assign | t.CPtr)(node)
        if asgn is None:
            VLogger.error("cast failed", "ASGN")
            return 0

        targets: list[ast.AST | t.CPtr] | t.CPtr = asgn.targets
        if targets is None:
            VLogger.error("targets is None", "ASGN")
            return 0

        # 从 self.Trans 取共享状态（替代 11 个参数）
        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

        # 翻译 RHS 值
        rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
            builder, pool, mod, asgn.value, None, 0, self.Trans)
        if rhs_val is None:
            # 增强错误信息：包含文件名 + lineno + AST 节点详情，便于定位
            fb: t.CChar | t.CPtr = VLogger.fmt_buf()
            if fb is not None and asgn.value is not None:
                sha1: str = self.Trans.ModuleSha1
                val_kind: int = asgn.value.kind()
                val_line: t.CInt = asgn.value.lineno
                # 通过 sha1 查找文件名（人类可读）
                rel_path: str = None
                if sha1 is not None:
                    rel_path = StubMerger.LookupSha1RelPath(sha1)
                # 根据 AST 节点类型提取详情
                if val_kind == ast.ASTKind.Attribute:
                    at_nv: ast.Attribute | t.CPtr = (ast.Attribute | t.CPtr)(asgn.value)
                    attr_nm: str = at_nv.attr
                    target_nm: str = "?"
                    if at_nv.value is not None:
                        if at_nv.value.kind() == ast.ASTKind.Name:
                            tn_nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(at_nv.value)
                            target_nm = tn_nm.id
                    if rel_path is not None:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: %s:%d '%s.%s' 属性访问失败",
                            rel_path, val_line, target_nm, attr_nm)
                    else:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: sha1=%s:%d '%s.%s' 属性访问失败",
                            sha1, val_line, target_nm, attr_nm)
                elif val_kind == ast.ASTKind.Name:
                    nm_nv: ast.Name | t.CPtr = (ast.Name | t.CPtr)(asgn.value)
                    if rel_path is not None:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: %s:%d 变量 '%s' 未找到",
                            rel_path, val_line, nm_nv.id)
                    else:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: sha1=%s:%d 变量 '%s' 未找到",
                            sha1, val_line, nm_nv.id)
                elif val_kind == ast.ASTKind.Call:
                    cl_nv: ast.Call | t.CPtr = (ast.Call | t.CPtr)(asgn.value)
                    func_nm: str = "?"
                    if cl_nv.func is not None:
                        if cl_nv.func.kind() == ast.ASTKind.Name:
                            fn_nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(cl_nv.func)
                            func_nm = fn_nm.id
                    if rel_path is not None:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: %s:%d 调用 '%s(...)' 返回 None",
                            rel_path, val_line, func_nm)
                    else:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: sha1=%s:%d 调用 '%s(...)' 返回 None",
                            sha1, val_line, func_nm)
                else:
                    if rel_path is not None:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: %s:%d [kind=%d]",
                            rel_path, val_line, val_kind)
                    else:
                        viperlib.snprintf(fb, 1024,
                            "赋值右侧为 None: sha1=%s:%d [kind=%d]",
                            sha1, val_line, val_kind)
                VLogger.error(fb, "ASGN")
            else:
                VLogger.error("赋值右侧为 None (asgn.value 为空)", "ASGN")
            return 0

        new_vars: int = 0
        tn: t.CSizeT = targets.__len__()
        for ti in range(tn):
            target: ast.AST | t.CPtr = targets.get(ti)
            if target is None:
                continue

            tk: int = target.kind()

            # Subscript 赋值: arr[i] = val / ptr[i] = val / list[i] = val
            if tk == ast.ASTKind.Subscript:
                # 检查是否是 list[T] 类型的 Subscript（泛型类不注册 struct）
                # list 的 subscript 赋值走 __setitem__ 内联路径
                list_obj: llvmlite.Value | t.CPtr = HandlesExpr.is_list_subscript(
                    target, self.Trans)
                if list_obj is not None:
                    # list[T] 类型: 内联生成 __setitem__ 逻辑
                    sub_node: ast.Subscript | t.CPtr = (ast.Subscript | t.CPtr)(target)
                    if sub_node is not None and sub_node.slice is not None:
                        list_idx_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                            builder, pool, mod, sub_node.slice, None, 0, self.Trans)
                        if list_idx_val is not None:
                            HandlesExpr.list_setitem_inline(
                                builder, pool, list_obj, list_idx_val, rhs_val)
                    continue
                # 普通 Subscript 赋值
                elem_ptr: llvmlite.Value | t.CPtr = HandlesExpr.get_subscript_ptr(
                    builder, pool, mod, target, self.Trans)
                if elem_ptr is not None:
                    store_val: llvmlite.Value | t.CPtr = rhs_val
                    if elem_ptr.Ty is not None:
                        elem_ty: llvmlite.LLVMType | t.CPtr = elem_ptr.Ty.Pointee
                        if elem_ty is not None:
                            store_val = HandlesExpr.coerce_to_type(
                                builder, rhs_val, elem_ty)
                    llvmlite.build_store(builder, store_val, elem_ptr)
                else:
                    # get_subscript_ptr 返回 None: 尝试 __setitem__ 运算符重载
                    # 适用于自定义类（如 hashtable[key]=val → hashtable.__setitem__(key, val)）
                    sub_asgn: ast.Subscript | t.CPtr = (ast.Subscript | t.CPtr)(target)
                    setitem_done: int = 0
                    if sub_asgn is not None and sub_asgn.value is not None:
                        if sub_asgn.value.kind() == ast.ASTKind.Name:
                            sub_nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(sub_asgn.value)
                            if sub_nm.id is not None:
                                sub_alloca: llvmlite.Value | t.CPtr = HandlesVar.lookup_var(
                                    self.Trans.SymTab, sub_nm.id)
                                if sub_alloca is not None and sub_alloca.Ty is not None:
                                    if HandlesExpr.is_ptr_type(sub_alloca.Ty) != 0:
                                        sub_pointee: llvmlite.LLVMType | t.CPtr = sub_alloca.Ty.Pointee
                                        if sub_pointee is not None:
                                            cls_nm_set: str = HandlesStruct.get_class_name_by_type(pool, sub_pointee)
                                            obj_val_set: llvmlite.Value | t.CPtr = llvmlite.build_load(
                                                builder, sub_pointee, sub_alloca)
                                            # 指针类型变量 (X|t.CPtr): sub_pointee 是 Ptr(Struct),
                                            # 需 load 获取 Ptr(Struct) 再检查 inner struct
                                            if cls_nm_set is None and obj_val_set is not None:
                                                if HandlesExpr.is_ptr_type(sub_pointee) != 0:
                                                    inner_ty_set: llvmlite.LLVMType | t.CPtr = sub_pointee.Pointee
                                                    if inner_ty_set is not None:
                                                        cls_nm_set = HandlesStruct.get_class_name_by_type(pool, inner_ty_set)
                                            if cls_nm_set is not None and obj_val_set is not None:
                                                key_val_set: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                                                    builder, pool, mod, sub_asgn.slice, None, 0, self.Trans)
                                                if key_val_set is not None:
                                                    arg_vals_set: t.CSizeT | t.CPtr = pool.alloc(16)
                                                    if arg_vals_set is not None:
                                                        arg_vals_set[0] = t.CSizeT(key_val_set)
                                                        arg_vals_set[1] = t.CSizeT(rhs_val)
                                                        HandlesExprCall._call_method_on_ptr(
                                                            pool, builder, mod, cls_nm_set, "__setitem__",
                                                            obj_val_set, arg_vals_set, 2, self.Trans)
                                                        setitem_done = 1
                        elif sub_asgn.value.kind() == ast.ASTKind.Attribute:
                            # self._ht[key] = val → self._ht.__setitem__(key, val)
                            # 通过 get_attribute_ptr 获取字段指针，再调用 __setitem__
                            field_ptr_set: llvmlite.Value | t.CPtr = HandlesExpr.get_attribute_ptr(
                                builder, pool, mod, sub_asgn.value, self.Trans)
                            if field_ptr_set is not None and field_ptr_set.Ty is not None:
                                if HandlesExpr.is_ptr_type(field_ptr_set.Ty) != 0:
                                    field_pointee_set: llvmlite.LLVMType | t.CPtr = field_ptr_set.Ty.Pointee
                                    if field_pointee_set is not None:
                                        cls_nm_attr: str = HandlesStruct.get_class_name_by_type(pool, field_pointee_set)
                                        obj_val_attr: llvmlite.Value | t.CPtr = field_ptr_set
                                        # field_ptr 是 Ptr(Ptr(Struct)) (X|t.CPtr 字段):
                                        # load 解引用获取 Ptr(Struct)
                                        if cls_nm_attr is None and HandlesExpr.is_ptr_type(field_pointee_set) != 0:
                                            inner_struct_attr: llvmlite.LLVMType | t.CPtr = field_pointee_set.Pointee
                                            if inner_struct_attr is not None:
                                                cls_nm_attr = HandlesStruct.get_class_name_by_type(pool, inner_struct_attr)
                                                if cls_nm_attr is not None:
                                                    obj_val_attr = llvmlite.build_load(
                                                        builder, field_pointee_set, field_ptr_set)
                                        if cls_nm_attr is not None and obj_val_attr is not None:
                                            key_val_attr: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                                                builder, pool, mod, sub_asgn.slice, None, 0, self.Trans)
                                            if key_val_attr is not None:
                                                arg_vals_attr: t.CSizeT | t.CPtr = pool.alloc(16)
                                                if arg_vals_attr is not None:
                                                    arg_vals_attr[0] = t.CSizeT(key_val_attr)
                                                    arg_vals_attr[1] = t.CSizeT(rhs_val)
                                                    HandlesExprCall._call_method_on_ptr(
                                                        pool, builder, mod, cls_nm_attr, "__setitem__",
                                                        obj_val_attr, arg_vals_attr, 2, self.Trans)
                                                    setitem_done = 1
                    if setitem_done == 0:
                        # 通用 fallback: 翻译 sub.value 并尝试 __setitem__ 或直接 GEP+store
                        # 处理 Name/Attribute 之外的节点以及 SymTab 查找失败的情况
                        if sub_asgn is not None and sub_asgn.value is not None:
                            obj_val_gen: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                                builder, pool, mod, sub_asgn.value, None, 0, self.Trans)
                            key_val_gen: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                                builder, pool, mod, sub_asgn.slice, None, 0, self.Trans)
                            if obj_val_gen is not None and key_val_gen is not None \
                                    and obj_val_gen.Ty is not None:
                                # 内联 _get_custom_struct_cls_nm 逻辑（避免依赖新函数）
                                cls_nm_gen: str = HandlesStruct.get_class_name_by_type(
                                    pool, obj_val_gen.Ty)
                                if cls_nm_gen is None and HandlesExpr.is_ptr_type(obj_val_gen.Ty) != 0:
                                    inner_gen: llvmlite.LLVMType | t.CPtr = obj_val_gen.Ty.Pointee
                                    if inner_gen is not None:
                                        cls_nm_gen = HandlesStruct.get_class_name_by_type(pool, inner_gen)
                                if cls_nm_gen is not None:
                                    # 自定义结构体: 调用 __setitem__
                                    arg_vals_gen: t.CSizeT | t.CPtr = pool.alloc(16)
                                    if arg_vals_gen is not None:
                                        arg_vals_gen[0] = t.CSizeT(key_val_gen)
                                        arg_vals_gen[1] = t.CSizeT(rhs_val)
                                        HandlesExprCall._call_method_on_ptr(
                                            pool, builder, mod, cls_nm_gen, "__setitem__",
                                            obj_val_gen, arg_vals_gen, 2, self.Trans)
                                        setitem_done = 1
                                else:
                                    # 非自定义结构体: 直接 GEP + store
                                    # 处理 get_subscript_ptr 因边界情况返回 None 的指针下标
                                    if HandlesExpr.is_ptr_type(obj_val_gen.Ty) != 0:
                                        elem_gen: llvmlite.LLVMType | t.CPtr = obj_val_gen.Ty.Pointee
                                        if elem_gen is not None:
                                            store_val_gen: llvmlite.Value | t.CPtr = rhs_val
                                            elem_ptr_gen: llvmlite.Value | t.CPtr = None
                                            match elem_gen:
                                                case llvmlite.LLVMType.Array(arr_elem_gen, _):
                                                    elem_ptr_gen = llvmlite.build_gep_array(
                                                        builder, elem_gen, arr_elem_gen,
                                                        obj_val_gen, key_val_gen)
                                                case _:
                                                    elem_ptr_gen = llvmlite.build_gep(
                                                        builder, elem_gen, obj_val_gen, key_val_gen)
                                            if elem_ptr_gen is not None and elem_ptr_gen.Ty is not None:
                                                pt_gen: llvmlite.LLVMType | t.CPtr = elem_ptr_gen.Ty.Pointee
                                                if pt_gen is not None:
                                                    store_val_gen = HandlesExpr.coerce_to_type(
                                                        builder, rhs_val, pt_gen)
                                                llvmlite.build_store(builder, store_val_gen, elem_ptr_gen)
                                                setitem_done = 1
                                    if setitem_done == 0:
                                        sub_vk: int = sub_asgn.value.kind()
                        if setitem_done == 0:
                            HandlesType.fatal_error(target, "subscript ptr is None")
                continue

            # Attribute 赋值: obj.field = val
            if tk == ast.ASTKind.Attribute:
                field_ptr: llvmlite.Value | t.CPtr = HandlesExpr.get_attribute_ptr(
                    builder, pool, mod, target, self.Trans)
                if field_ptr is not None:
                    # 获取字段类型，对 rhs_val 进行类型转换（如 i32 → i64）
                    store_val: llvmlite.Value | t.CPtr = rhs_val
                    if field_ptr.Ty is not None:
                        field_ty: llvmlite.LLVMType | t.CPtr = field_ptr.Ty.Pointee
                        if field_ty is not None:
                            store_val = HandlesExpr.coerce_to_type(
                                builder, rhs_val, field_ty)
                    llvmlite.build_store(builder, store_val, field_ptr)
                else:
                    # 构造详细错误信息
                    attr_node: ast.Attribute | t.CPtr = (ast.Attribute | t.CPtr)(target)
                    attr_name: str = "(unknown)"
                    obj_name: str = "(unknown)"
                    if attr_node is not None:
                        attr_name = attr_node.attr
                        if attr_node.value is not None and attr_node.value.kind() == ast.ASTKind.Name:
                            obj_nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(attr_node.value)
                            obj_name = obj_nm.id
                    err_buf: t.CChar | t.CPtr = pool.alloc(256)
                    if err_buf is not None:
                        viperlib.snprintf(err_buf, 256,
                            "attribute ptr is None: %s.%s",
                            obj_name, attr_name)
                        HandlesType.fatal_error(target, err_buf)
                    else:
                        HandlesType.fatal_error(target, "attribute ptr is None")
                continue

            # Name 赋值: var = val
            if tk == ast.ASTKind.Name:
                nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(target)
                if nm.id is not None:
                    # 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:
                            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)
                        continue

                    # nonlocal 变量：通过闭包 env 写入
                    if HT.is_nonlocal_name(self.Trans, nm.id) != 0:
                        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)
                        continue

                    # 普通局部变量
                    alloca: llvmlite.Value | t.CPtr = HandlesVar.get_or_create_sym(
                        self.Trans.SymTab, pool, builder, nm.id, rhs_val.Ty)
                    if alloca is not None:
                        # 按 alloca 类型对值进行转换（如 double → float）
                        store_val: llvmlite.Value | t.CPtr = rhs_val
                        if alloca.Ty is not None:
                            alloca_ty: llvmlite.LLVMType | t.CPtr = alloca.Ty.Pointee
                            if alloca_ty is not None:
                                store_val = HandlesExpr.coerce_to_type(
                                    builder, rhs_val, alloca_ty)
                        llvmlite.build_store(builder, store_val, alloca)
                        existing: llvmlite.Value | t.CPtr = HandlesVar.lookup_current(
                            self.Trans.SymTab, nm.id)
                        if existing is None:
                            new_vars += 1
                    else:
                        fb: t.CChar | t.CPtr = VLogger.fmt_buf()
                        if fb is not None:
                            viperlib.snprintf(fb, 1024, "alloca failed for %s", nm.id)
                            VLogger.error(fb, "ASGN")

        return new_vars


# ============================================================
# NewAssignHandle - 工厂函数：分配并初始化 AssignHandle
# ============================================================
def NewAssignHandle(pool: memhub.MemBuddy | t.CPtr,
                    trans: HT.Translator | t.CPtr) -> AssignHandle | t.CPtr:
    h: AssignHandle | t.CPtr = pool.alloc(AssignHandle.__sizeof__())
    if h is None:
        return None
    string.memset(h, 0, AssignHandle.__sizeof__())
    h.Trans = trans
    return h
