from __future__ import annotations
from typing import TYPE_CHECKING
import 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__',
    }
    _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__',
    }

    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 VarName in Gen.var_const_flags:
                src_info: str = Gen._get_node_info()
                _vlog().error(f"不能对 const 变量 '{VarName}' 增量赋值{src_info}")
                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 and VarName in Gen.module.globals:
                GVar: ir.GlobalVariable = 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)
                Gen._store(NewVal, GVar)
                return
            if VarName in Gen._reg_values:
                OldVal: ir.Value = Gen._reg_values[VarName]
                saved_block: ir.Block = Gen.builder.block
                entry_block: ir.Block = Gen.func.entry_basic_block
                Gen.builder.position_at_start(entry_block)
                var: ir.AllocaInstr = Gen.builder.alloca(OldVal.type, name=VarName)
                Gen._store(OldVal, var)
                Gen.builder.position_at_end(saved_block)
                Gen.variables[VarName] = var
                del Gen._reg_values[VarName]
            if VarName in Gen._direct_values:
                OldVal: ir.Value = Gen._direct_values.pop(VarName)
                saved_block: ir.Block = Gen.builder.block
                entry_block: ir.Block = Gen.func.entry_basic_block
                Gen.builder.position_at_start(entry_block)
                var: ir.AllocaInstr = Gen.builder.alloca(OldVal.type, name=VarName)
                Gen._store(OldVal, var)
                Gen.builder.position_at_end(saved_block)
                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)
                        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)
                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)
                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)
                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)
                    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)