import t, c
from stdint import *
import ast
import llvmlite
import memhub
import stdio
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.HandlesNonlocal as HandlesNonlocal


# ============================================================
# HandlesAugAssign - 增强赋值语句处理（Mixin 继承模式）
#
# 处理 += -= *= /= %= &= |= ^= <<= >>=
# 流程: load target → apply binop → store result
#
# 支持 local/global/nonlocal 三种变量作用域
# ============================================================


@t.NoVTable
class AugAssignHandle(HandlesBase.Mixin):
    """增强赋值处理器 (+=, -=, *=, etc.)：继承 Mixin 获得 Trans 回指针"""

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

    # ============================================================
    # Handle - 处理 AugAssign 语句，返回新增变量数（始终为 0）
    # ============================================================
    def Handle(self, node: ast.AST | t.CPtr) -> int:
        """翻译增强赋值语句 (x += 1, y -= 2, etc.)"""
        if node is None:
            return 0

        aug: ast.AugAssign | t.CPtr = (ast.AugAssign | t.CPtr)(node)
        if aug is None:
            return 0

        target: ast.AST | t.CPtr = aug.target
        if target is None:
            return 0

        tk: int = target.kind()

        # Attribute 目标: self.field += 1
        # 流程: get_attribute_ptr → load → binop → store
        if tk == ast.ASTKind.Attribute:
            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

            # 1. 获取字段指针
            field_ptr: llvmlite.Value | t.CPtr = HandlesExpr.get_attribute_ptr(
                builder, pool, mod, target, self.Trans)
            if field_ptr is None:
                stdio.printf("[AUGASGN] attribute ptr is None\n")
                return 0

            # 2. 确定字段类型并加载当前值
            target_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
            if field_ptr.Ty is not None:
                target_ty = field_ptr.Ty.Pointee
            cur_val: llvmlite.Value | t.CPtr = llvmlite.build_load(
                builder, target_ty, field_ptr)
            if cur_val is None:
                stdio.printf("[AUGASGN] cannot load attribute\n")
                return 0

            # 3. 翻译 RHS 值
            rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
                builder, pool, mod, aug.value, None, 0, self.Trans)
            if rhs_val is None:
                stdio.printf("[AUGASGN] rhs is None\n")
                return 0

            # 4. 应用二元运算
            result: llvmlite.Value | t.CPtr = _apply_aug_op(
                pool, builder, aug.op, cur_val, rhs_val)
            if result is None:
                stdio.printf("[AUGASGN] binop failed for attr op=%d\n", aug.op)
                return 0

            # 5. 类型对齐并存储
            result = HandlesExpr.coerce_to_type(builder, result, target_ty)
            if result is None:
                return 0
            llvmlite.build_store(builder, result, field_ptr)
            return 0

        if tk != ast.ASTKind.Name:
            stdio.printf("[AUGASGN] only Name/Attribute target supported\n")
            return 0

        nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(target)
        if nm is None or nm.id is None:
            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
        i32_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)

        # 1. 确定变量作用域类型: 0=local, 1=global, 2=nonlocal
        scope_type: int = 0
        if HT.is_global_name(self.Trans, nm.id) != 0:
            scope_type = 1
        elif HT.is_nonlocal_name(self.Trans, nm.id) != 0:
            scope_type = 2

        # 2. 加载当前值
        cur_val: llvmlite.Value | t.CPtr = None
        target_alloca: llvmlite.Value | t.CPtr = None
        target_ty: llvmlite.LLVMType | t.CPtr = i32_ty

        if scope_type == 1:
            # global 变量
            target_alloca = HandlesVar.lookup_module_var(
                self.Trans.SymTab, nm.id)
            if target_alloca is not None:
                if target_alloca.Ty is not None:
                    target_ty = target_alloca.Ty.Pointee
                cur_val = llvmlite.build_load(builder, target_ty, target_alloca)
        elif scope_type == 2:
            # nonlocal 变量（通过闭包 env）
            cur_val = HandlesNonlocal.load_nonlocal_var(self.Trans, nm.id)
            if cur_val is not None:
                target_ty = cur_val.Ty
        else:
            # 普通局部变量
            target_alloca = HandlesVar.lookup_var(self.Trans.SymTab, nm.id)
            if target_alloca is not None:
                if target_alloca.Ty is not None:
                    target_ty = target_alloca.Ty.Pointee
                cur_val = llvmlite.build_load(builder, target_ty, target_alloca)

        if cur_val is None:
            stdio.printf("[AUGASGN] cannot load target %s\n", nm.id)
            return 0

        # 3. 翻译 RHS 值
        rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value(
            builder, pool, mod, aug.value, None, 0, self.Trans)
        if rhs_val is None:
            stdio.printf("[AUGASGN] rhs is None\n")
            return 0

        # 4. 类型对齐 + 应用二元运算
        # 注意: AugAssign 不走运算符重载（语义上需要 __iadd__ 而非 __add__）
        result: llvmlite.Value | t.CPtr = _apply_aug_op(
            pool, builder, aug.op, cur_val, rhs_val)
        if result is None:
            stdio.printf("[AUGASGN] binop failed for op=%d\n", aug.op)
            return 0

        # 5. 存储结果
        result = HandlesExpr.coerce_to_type(builder, result, target_ty)
        if result is None:
            return 0

        if scope_type == 1:
            # global 变量
            if target_alloca is not None:
                llvmlite.build_store(builder, result, target_alloca)
        elif scope_type == 2:
            # nonlocal 变量
            nl_ptr: llvmlite.Value | t.CPtr = HandlesNonlocal.get_nonlocal_var_ptr(
                self.Trans, nm.id)
            if nl_ptr is not None:
                llvmlite.build_store(builder, result, nl_ptr)
        else:
            # 普通局部变量
            if target_alloca is not None:
                llvmlite.build_store(builder, result, target_alloca)

        return 0


# ============================================================
# _apply_aug_op - 应用增强赋值的二元运算
#
# 支持指针算术: ptr += int / ptr -= int
# 整数运算自动类型提升
# ============================================================
def _apply_aug_op(pool: memhub.MemBuddy | t.CPtr,
                  builder: llvmlite.IRBuilder | t.CPtr,
                  op: int,
                  lhs: llvmlite.Value | t.CPtr,
                  rhs: llvmlite.Value | t.CPtr) -> llvmlite.Value | t.CPtr:
    """应用增强赋值的二元运算（指针算术 + 整数运算）"""
    lhs_bits: int = HandlesExpr.get_llvm_type_bits(lhs.Ty)
    rhs_bits: int = HandlesExpr.get_llvm_type_bits(rhs.Ty)

    # 指针算术: ptr += int / ptr -= int
    if lhs_bits == 0 and rhs_bits != 0:
        if op == ast.OpKind.Add or op == ast.OpKind.Sub:
            i64_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int64(pool)
            ptr_as_int: llvmlite.Value | t.CPtr = llvmlite.build_ptrtoint(builder, lhs, i64_ty)
            int_val: llvmlite.Value | t.CPtr = HandlesExpr.coerce_to_type(builder, rhs, i64_ty)
            if ptr_as_int is None or int_val is None:
                return None
            if op == ast.OpKind.Add:
                result: llvmlite.Value | t.CPtr = llvmlite.build_add(builder, ptr_as_int, int_val)
            else:
                result = llvmlite.build_sub(builder, ptr_as_int, int_val)
            if result is None:
                return None
            return llvmlite.build_inttoptr(builder, result, lhs.Ty)
        return None

    # 整数运算：类型提升
    if lhs_bits > rhs_bits and rhs_bits > 0:
        rhs = HandlesExpr.coerce_to_type(builder, rhs, lhs.Ty)
    elif rhs_bits > lhs_bits and lhs_bits > 0:
        lhs = HandlesExpr.coerce_to_type(builder, lhs, rhs.Ty)

    if op == ast.OpKind.Add:
        return llvmlite.build_add(builder, lhs, rhs)
    elif op == ast.OpKind.Sub:
        return llvmlite.build_sub(builder, lhs, rhs)
    elif op == ast.OpKind.Mult:
        return llvmlite.build_mul(builder, lhs, rhs)
    elif op == ast.OpKind.Div:
        return llvmlite.build_sdiv(builder, lhs, rhs)
    elif op == ast.OpKind.Mod:
        return llvmlite.build_srem(builder, lhs, rhs)
    elif op == ast.OpKind.BitAnd:
        return llvmlite.build_and(builder, lhs, rhs)
    elif op == ast.OpKind.BitOr:
        return llvmlite.build_or(builder, lhs, rhs)
    elif op == ast.OpKind.BitXor:
        return llvmlite.build_xor(builder, lhs, rhs)
    elif op == ast.OpKind.LShift:
        return llvmlite.build_shl(builder, lhs, rhs)
    elif op == ast.OpKind.RShift:
        return llvmlite.build_ashr(builder, lhs, rhs)
    return None


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