from __future__ import annotations
from typing import TYPE_CHECKING
import vast as ast
import sys
import llvmlite.ir as ir
if TYPE_CHECKING:
    from lib.core.translator import Translator
from lib.core.VLogger import get_logger as _vlog
from lib.core.Handles.HandlesBase import BaseHandle


class AugAssignHandle(BaseHandle):
    _AUGOP_MAP: dict[type, str] = {
        ast.Add: '__iadd__', ast.Sub: '__isub__', ast.Mult: '__imul__',
        ast.Div: '__itruediv__', ast.FloorDiv: '__ifloordiv__', ast.Mod: '__imod__', ast.Pow: '__ipow__',
        ast.BitAnd: '__iand__', ast.BitOr: '__ior__', ast.BitXor: '__ixor__',
        ast.LShift: '__ilshift__', ast.RShift: '__irshift__',
    }
    _AUGOP_FALLBACK: dict[type, str] = {
        ast.Add: '__add__', ast.Sub: '__sub__', ast.Mult: '__mul__',
        ast.Div: '__truediv__', ast.FloorDiv: '__floordiv__', ast.Mod: '__mod__', ast.Pow: '__pow__',
        ast.BitAnd: '__and__', ast.BitOr: '__or__', ast.BitXor: '__xor__',
        ast.LShift: '__lshift__', ast.RShift: '__rshift__',
    }

    def _resolve_global_var_type(self, Gen: "Translator.LlvmGen", VarName: str, fallback: ir.Type) -> ir.Type:
        """解析 global 变量的 LLVM 类型：优先用模块级 AnnAssign 的注解类型。

        问题场景：函数内 `_g_alloc_count += 1`（global 声明），模块级注解为
        `_g_alloc_count: CSizeT`（i64）。若不查模块级注解类型，此处会用 RHS 字面量
        类型（i32）自动建 i32 全局，与模块内其他读取/写入（i64）类型不匹配，
        产生 "Operands must be the same type, got (i64, i32)" 类错误。
        模块级 AnnAssign 已通过 _EmitGlobalAnnAssignLlvm 创建 LLVM 全局，
        但注解无初始值（如 `_g_alloc_max: CSizeT = 0`）时 initializer 为空、
        Gen.module.globals 中无该符号，故需在此兜底解析注解类型。
        """
        module_ann: dict = getattr(Gen, '_module_ann_types', None)
        if module_ann and VarName in module_ann:
            ann_type: ir.Type = module_ann[VarName]
            if ann_type is not None:
                return ann_type
        if isinstance(fallback, ir.IntType):
            return fallback
        return ir.IntType(64)

    def _is_aug_unsigned(self, Node: ast.AugAssign, OldVal: ir.Value | None, Value: ir.Value) -> bool:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        if isinstance(Node.target, ast.Name):
            if Gen._check_node_unsigned(Node.target):
                return True
        elif isinstance(Node.target, ast.Attribute):
            if Gen._check_node_unsigned(Node.target):
                return True
        elif isinstance(Node.target, ast.Subscript):
            if Gen._check_node_unsigned(Node.target):
                return True
        if Node.value is not None and Gen._check_node_unsigned(Node.value):
            return True
        return False

    def _HandleAugAssignLlvm(self, Node: ast.AugAssign) -> None:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        if isinstance(Node.target, ast.Name):
            VarName: str = Node.target.id
            if (VarName not in Gen.variables
                    and VarName not in Gen.global_vars
                    and VarName not in Gen._reg_values
                    and VarName not in Gen._direct_values):
                src_info: str = Gen._get_node_info()
                _vlog().error(f"变量 '{VarName}' 未声明，不能用于增量赋值（需用 global 声明或先声明局部变量）{src_info}")
                sys.exit(1)
            if Gen._IsVarConst(VarName):
                _vlog().error(Gen._ConstAssignMessage("增量赋值", VarName))
                sys.exit(1)
            ClassName: str | None = Gen.var_struct_class.get(VarName)
            iop_name: str | None = self._AUGOP_MAP.get(type(Node.op))
            if ClassName and iop_name:
                iFuncName: str = f'{ClassName}.{iop_name}'
                FuncName: str = f'{ClassName}.{self._AUGOP_FALLBACK.get(type(Node.op), "")}'
                if iFuncName in Gen.functions or FuncName in Gen.functions:
                    OldVal: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(ast.Name(id=VarName, ctx=ast.Load()))
                    Value: ir.Value | None = self.HandleExprLlvm(Node.value)
                    if OldVal and Value:
                        result: ir.Value | None
                        if iFuncName in Gen.functions:
                            result = self.Trans.ExprHandler._try_operator_overLoad(ClassName, iop_name, OldVal, Gen, Value)
                        else:
                            result = self.Trans.ExprHandler._try_operator_overLoad(ClassName, self._AUGOP_FALLBACK[type(Node.op)], OldVal, Gen, Value)
                        if result is not None:
                            if VarName in Gen.variables and Gen.variables[VarName] is not None:
                                VarPtr: ir.Value = Gen.variables[VarName]
                                if VarPtr.type.pointee == result.type:
                                    Gen._store(result, VarPtr)
                                elif isinstance(VarPtr.type.pointee, ir.PointerType) and isinstance(result.type, ir.PointerType):
                                    CastedVal: ir.Value = Gen.builder.bitcast(result, VarPtr.type.pointee, name=f"cast_{VarName}")
                                    Gen._store(CastedVal, VarPtr)
                                else:
                                    NewVar: ir.AllocaInstr = Gen._allocaEntry(result.type, name=VarName)
                                    Gen._store(result, NewVar)
                                    Gen.variables[VarName] = NewVar
                            else:
                                Gen._reg_values[VarName] = result
                                Gen.variables[VarName] = None
                            return
        Value: ir.Value | None = self.HandleExprLlvm(Node.value)
        if not Value:
            return
        is_unsigned: bool = self._is_aug_unsigned(Node, None, Value)
        if isinstance(Node.target, ast.Name):
            VarName: str = Node.target.id
            Op: str = self.Trans.ExprHandler.GetOpSymbol(Node.op)
            if VarName in Gen.global_vars:
                # 治本修复：global 声明的变量若未被模块级 AnnAssign 创建为 LLVM 全局
                # （注解无初始值的全局变量），此处自动创建零初始化全局变量
                GVar: ir.GlobalVariable
                if VarName not in Gen.module.globals:
                    _gv_type: ir.Type = self._resolve_global_var_type(Gen, VarName, Value.type)
                    GVar = ir.GlobalVariable(Gen.module, _gv_type, name=VarName)
                    GVar.initializer = ir.Constant(_gv_type, 0)
                else:
                    GVar = Gen.module.globals[VarName]
                Gen.variables[VarName] = GVar
                OldVal: ir.Value = Gen._load(GVar, name=VarName)
                NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned, is_aug=True)
                Gen._store(NewVal, GVar)
                return
            if VarName in Gen._reg_values:
                OldVal: ir.Value = Gen._reg_values[VarName]
                var: ir.AllocaInstr = Gen._allocaEntry(OldVal.type, name=VarName)
                Gen._store(OldVal, var)
                Gen.variables[VarName] = var
                del Gen._reg_values[VarName]
            if VarName in Gen._direct_values:
                OldVal: ir.Value = Gen._direct_values.pop(VarName)
                var: ir.AllocaInstr = Gen._allocaEntry(OldVal.type, name=VarName)
                Gen._store(OldVal, var)
                Gen.variables[VarName] = var
            if Gen.builder.block.is_terminated:
                return
            if VarName in Gen.variables and Gen.variables[VarName] is not None:
                OldVal: ir.Value = Gen._load(Gen.variables[VarName], name=VarName)
                store_ptr: ir.Value = Gen.variables[VarName]
                if isinstance(OldVal.type, ir.PointerType) and isinstance(Value.type, ir.IntType):
                    if Value.type.width != 64:
                        Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset")
                    NewVal: ir.Value
                    if Op == '-' or Op == '+':
                        if Op == '-':
                            NegValue: ir.Value = Gen.builder.neg(Value, name="neg_ptr_offset")
                            NewVal = Gen.builder.gep(OldVal, [NegValue], name="ptr_sub")
                        else:
                            NewVal = Gen.builder.gep(OldVal, [Value], name="ptr_add")
                    else:
                        ptr_val: ir.Value = OldVal
                        OldVal = Gen._load(OldVal, name="deref_ptr_for_op")
                        orig_type: ir.Type = OldVal.type
                        if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType):
                            if OldVal.type.width < Value.type.width:
                                OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left")
                            elif OldVal.type.width > Value.type.width:
                                Value = Gen.builder.zext(Value, OldVal.type, name="zext_right")
                        NewVal = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned, is_aug=True)
                        if isinstance(NewVal.type, ir.IntType) and isinstance(orig_type, ir.IntType):
                            if NewVal.type.width > orig_type.width:
                                NewVal = Gen.builder.trunc(NewVal, orig_type, name="trunc_result")
                        store_ptr = ptr_val
                else:
                    if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType):
                        if OldVal.type.width < Value.type.width:
                            OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left")
                        elif OldVal.type.width > Value.type.width:
                            Value = Gen.builder.zext(Value, OldVal.type, name="zext_right")
                    elif isinstance(OldVal.type, ir.PointerType) and isinstance(Value.type, ir.IntType):
                        if Op in ('+', '-'):
                            if Value.type.width < 64:
                                Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset_aug")
                            NewVal: ir.Value
                            if Op == '-':
                                NegValue: ir.Value = Gen.builder.neg(Value, name="neg_ptr_offset_aug")
                                NewVal = Gen.builder.gep(OldVal, [NegValue], name="ptr_sub_aug")
                            else:
                                NewVal = Gen.builder.gep(OldVal, [Value], name="ptr_add_aug")
                            Gen._store(NewVal, store_ptr)
                            return
                    NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned, is_aug=True)
                if isinstance(store_ptr.type, ir.PointerType):
                    TargetType: ir.Type = store_ptr.type.pointee
                    if NewVal.type != TargetType:
                        Coerced: ir.Value | None = Gen._coerce_value(NewVal, TargetType)
                        if Coerced is not None:
                            NewVal = Coerced
                Gen._store(NewVal, store_ptr)
            elif VarName == 'pos':
                var: ir.AllocaInstr = Gen._allocaEntry(Value.type, name=VarName)
                Gen.variables[VarName] = var
                OldVal: ir.Value = Gen._load(var, name=VarName)
                NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned, is_aug=True)
                Gen._store(NewVal, var)
        elif isinstance(Node.target, ast.Attribute):
            AttrName: str = Node.target.attr
            OldVal: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(Node.target)
            if OldVal:
                def _deref_if_ptr(val: ir.Value) -> ir.Value:
                    if isinstance(val.type, ir.PointerType):
                        pointee: ir.Type = val.type.pointee
                        if isinstance(pointee, (ir.IntType, ir.FloatType, ir.DoubleType)):
                            return Gen._load(val, name="deref_aug")
                    return val
                OldVal = _deref_if_ptr(OldVal)
                Op: str = self.Trans.ExprHandler.GetOpSymbol(Node.op)
                if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType):
                    if OldVal.type.width < Value.type.width:
                        OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left_attr")
                    elif OldVal.type.width > Value.type.width:
                        Value = Gen.builder.zext(Value, OldVal.type, name="zext_right_attr")
                NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned, is_aug=True)
                self.Trans.AssignHandler._HandleAttributeStoreLlvm(Node.target, NewVal)
        elif isinstance(Node.target, ast.Subscript):
            SubTarget: ast.Subscript = Node.target
            ElemPtr: ir.Value | None = self.Trans.ExprAttrHandle.HandleSubscriptPtrLlvm(SubTarget)
            if ElemPtr:
                if isinstance(ElemPtr.type, ir.PointerType):
                    OldVal: ir.Value = Gen._load(ElemPtr, name="old_subscript_val")
                    Op: str = self.Trans.ExprHandler.GetOpSymbol(Node.op)
                    if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType):
                        if OldVal.type.width < Value.type.width:
                            OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left_sub")
                        elif OldVal.type.width > Value.type.width:
                            Value = Gen.builder.zext(Value, OldVal.type, name="zext_right_sub")
                    elif isinstance(OldVal.type, ir.PointerType) and isinstance(Value.type, ir.IntType):
                        if Op in ('+', '-'):
                            if Value.type.width < 64:
                                Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset_sub")
                            NewVal: ir.Value
                            if Op == '-':
                                NegValue: ir.Value = Gen.builder.neg(Value, name="neg_ptr_offset_sub")
                                NewVal = Gen.builder.gep(OldVal, [NegValue], name="ptr_sub_sub")
                            else:
                                NewVal = Gen.builder.gep(OldVal, [Value], name="ptr_add_sub")
                            Gen._store(NewVal, ElemPtr)
                            return
                    NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned, is_aug=True)
                    if isinstance(NewVal.type, ir.IntType) and isinstance(ElemPtr.type.pointee, ir.IntType):
                        if NewVal.type.width > ElemPtr.type.pointee.width:
                            NewVal = Gen.builder.trunc(NewVal, ElemPtr.type.pointee, name="trunc_sub_result")
                    Gen._store(NewVal, ElemPtr)
        elif isinstance(Node.target, ast.UnaryOp):
            # 解引用增强赋值 `*p += v`：目标节点是 UnaryOp(Deref)
            # （SetStoreContext 不改 UnaryOp 的 ctx），与 `*p = v` 同一形态。
            # 流程：操作数求值取指针 → load 当前值 → 二元运算 → store 回同址。
            if not isinstance(Node.target.op, ast.Deref):
                BadOp: str | None = self.Trans.ExprHandler.GetUnaryOpSymbol(Node.target.op)
                raise Exception(
                    f"AugAssign 目标只支持单目 '*' 解引用（'{BadOp}' 不是合法左值）"
                    f"{Gen._get_node_info()}"
                )
            DerefPtr: ir.Value | None = self.HandleExprLlvm(Node.target.operand)
            if DerefPtr is None:
                raise Exception(f"解引用增量赋值 '*' 的操作数无法求值{Gen._get_node_info()}")
            if not isinstance(DerefPtr.type, ir.PointerType):
                raise Exception(
                    f"解引用增量赋值 '*' 的操作数必须是指针类型（实际 {DerefPtr.type}）"
                    f"{Gen._get_node_info()}"
                )
            DerefOp: str = self.Trans.ExprHandler.GetOpSymbol(Node.op)
            DerefOld: ir.Value = Gen._load(DerefPtr, name="old_deref_val")
            if isinstance(DerefOld.type, ir.IntType) and isinstance(Value.type, ir.IntType):
                if DerefOld.type.width < Value.type.width:
                    DerefOld = Gen.builder.zext(DerefOld, Value.type, name="zext_left_deref")
                elif DerefOld.type.width > Value.type.width:
                    Value = Gen.builder.zext(Value, DerefOld.type, name="zext_right_deref")
            elif isinstance(DerefOld.type, ir.PointerType) and isinstance(Value.type, ir.IntType):
                if DerefOp in ('+', '-'):
                    if Value.type.width < 64:
                        Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset_deref")
                    OffVal: ir.Value
                    if DerefOp == '-':
                        OffVal = Gen.builder.gep(DerefOld, [Gen.builder.neg(Value, name="neg_ptr_offset_deref")], name="ptr_sub_deref")
                    else:
                        OffVal = Gen.builder.gep(DerefOld, [Value], name="ptr_add_deref")
                    Gen._store(OffVal, DerefPtr)
                    return
            DerefNew: ir.Value = Gen.emit_binary_op(DerefOp, DerefOld, Value, is_unsigned=is_unsigned, is_aug=True)
            if isinstance(DerefNew.type, ir.IntType) and isinstance(DerefPtr.type.pointee, ir.IntType):
                if DerefNew.type.width > DerefPtr.type.pointee.width:
                    DerefNew = Gen.builder.trunc(DerefNew, DerefPtr.type.pointee, name="trunc_deref_result")
            Gen._store(DerefNew, DerefPtr)
        else:
            # fail-fast：目标形态未覆盖时绝不静默丢弃（对齐 ViperC
            # translate_augassign 的同名 fail-fast），避免"语句被整条丢掉、
            # 变量静默不变"。
            TargetOp: str | None = None
            if isinstance(Node.target, ast.UnaryOp):
                TargetOp = self.Trans.ExprHandler.GetUnaryOpSymbol(Node.target.op)
            raise Exception(
                f"AugAssign 目标类型不支持，仅支持 Name/Attribute/Subscript/*解引用"
                f"（当前 {type(Node.target).__name__}"
                + (f" op='{TargetOp}'" if TargetOp else "")
                + f"）{Gen._get_node_info()}"
            )