import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import viperlib
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


# ============================================================
# 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:
            stdio.printf("[ASGN] cast failed\n")
            return 0

        targets: list[ast.AST | t.CPtr] | t.CPtr = asgn.targets
        if targets is None:
            stdio.printf("[ASGN] targets is None\n")
            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:
            stdio.printf("[ASGN] rhs_val is None\n")
            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)
                                            if cls_nm_set is not None:
                                                obj_val_set: llvmlite.Value | t.CPtr = llvmlite.build_load(
                                                    builder, sub_pointee, sub_alloca)
                                                if 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
                    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:
                        stdio.printf("[ASGN] alloca failed for %s\n", nm.id)

        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
