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


class ExprOpsHandle(BaseHandle):
    def _HandleBinOpLlvm(self, Node, VarType=None):
        Gen = self.Trans.LlvmGen
        Gen._set_node_info(Node, "BinOp")
        LeftVal = self.HandleExprLlvm(Node.left, VarType)
        if not LeftVal:
            return None
        RightVal = self.HandleExprLlvm(Node.right, VarType)
        if not RightVal:
            return None
        Op = self.Trans.ExprHandler.GetOpSymbol(Node.op)
        if not Op:
            return None
        OP_OVERLOAD_MAP = {
            '+': '__add__', '-': '__sub__', '*': '__mul__',
            '/': '__truediv__', '//': '__floordiv__', '%': '__mod__',
            '**': '__pow__',
        }
        if Op in OP_OVERLOAD_MAP:
            LeftClassName = None
            if isinstance(LeftVal.type, ir.PointerType):
                pointee = LeftVal.type.pointee
                if isinstance(pointee, (ir.IdentifiedStructType, ir.LiteralStructType)):
                    for CN, ST in Gen.structs.items():
                        if pointee == ST:
                            LeftClassName = CN
                            break
                elif isinstance(pointee, ir.IntType) and pointee.width == 8:
                    LeftClassName = self.Trans.ExprUtils._get_var_class(Node.left, Gen)
                    if LeftClassName and LeftClassName in Gen.structs:
                        TargetStructPtr = ir.PointerType(Gen.structs[LeftClassName])
                        LeftVal = Gen.builder.bitcast(LeftVal, TargetStructPtr, name=f"bitcast_to_{LeftClassName}")
            if LeftClassName is None and isinstance(Node.left, ast.Name):
                LeftClassName = Gen.var_struct_class.get(Node.left.id)
                if LeftClassName and LeftClassName in Gen.structs:
                    if isinstance(LeftVal.type, ir.PointerType) and isinstance(LeftVal.type.pointee, ir.IntType) and LeftVal.type.pointee.width == 8:
                        TargetStructPtr = ir.PointerType(Gen.structs[LeftClassName])
                        LeftVal = Gen.builder.bitcast(LeftVal, TargetStructPtr, name=f"bitcast_to_{LeftClassName}")
            if LeftClassName:
                op_name = OP_OVERLOAD_MAP[Op]
                result = self.Trans.ExprUtils._try_operator_overLoad(LeftClassName, op_name, LeftVal, Gen, RightVal)
                if result is not None:
                    return result
        if Op in ('+', 'Add') and isinstance(LeftVal.type, ir.PointerType) and isinstance(RightVal.type, ir.IntType):
            if isinstance(LeftVal.type.pointee, ir.IntType) and LeftVal.type.pointee.width == 8:
                if isinstance(Node.left, ast.Constant) and isinstance(Node.left.value, str) and len(Node.left.value) == 1:
                    CharVal = Gen._load(LeftVal, name="Load_char_const")
                    CharAsInt = Gen.builder.zext(CharVal, RightVal.type, name="char_to_int")
                    result = Gen.builder.add(CharAsInt, RightVal, name="char_add")
                    TruncResult = Gen.builder.trunc(result, ir.IntType(8), name="add_to_char")
                    return TruncResult
        is_unsigned = Gen._check_node_unsigned(Node.left) or Gen._check_node_unsigned(Node.right)
        if isinstance(LeftVal.type, ir.IntType) and isinstance(RightVal.type, ir.IntType):
            if LeftVal.type.width < RightVal.type.width:
                need_zext = is_unsigned or Gen._check_node_unsigned(Node.left) or LeftVal.type.width == 8
                if not need_zext and isinstance(RightVal, ir.Constant):
                    signed_max = (1 << (LeftVal.type.width - 1)) - 1
                    try:
                        if RightVal.constant > signed_max:
                            need_zext = True
                    except (TypeError, ValueError):
                        pass
                if need_zext:
                    LeftVal = Gen.builder.zext(LeftVal, RightVal.type, name="zext_left")
                else:
                    LeftVal = Gen.builder.sext(LeftVal, RightVal.type, name="sext_left")
            elif LeftVal.type.width > RightVal.type.width:
                need_zext = is_unsigned or Gen._check_node_unsigned(Node.right) or RightVal.type.width == 8
                if not need_zext and isinstance(LeftVal, ir.Constant):
                    signed_max = (1 << (RightVal.type.width - 1)) - 1
                    try:
                        if LeftVal.constant > signed_max:
                            need_zext = True
                    except (TypeError, ValueError):
                        pass
                if need_zext:
                    RightVal = Gen.builder.zext(RightVal, LeftVal.type, name="zext_right")
                else:
                    RightVal = Gen.builder.sext(RightVal, LeftVal.type, name="zext_right")
        elif isinstance(LeftVal.type, ir.PointerType) and isinstance(RightVal.type, ir.IntType):
            if Op not in ('+', '-', 'Add', 'Sub') and isinstance(LeftVal.type.pointee, ir.IntType):
                LeftVal = Gen._load(LeftVal, name="deref_ptr_binop")
                if LeftVal.type.width < RightVal.type.width:
                    LeftVal = Gen.builder.zext(LeftVal, RightVal.type, name="zext_deref_left")
                elif LeftVal.type.width > RightVal.type.width:
                    RightVal = Gen.builder.zext(RightVal, LeftVal.type, name="zext_right_deref")
        elif isinstance(RightVal.type, ir.PointerType) and isinstance(LeftVal.type, ir.IntType):
            if Op not in ('+', '-', 'Add', 'Sub') and isinstance(RightVal.type.pointee, ir.IntType):
                RightVal = Gen._load(RightVal, name="deref_ptr_binop_r")
                if RightVal.type.width < LeftVal.type.width:
                    RightVal = Gen.builder.zext(RightVal, LeftVal.type, name="zext_deref_right")
                elif RightVal.type.width > LeftVal.type.width:
                    LeftVal = Gen.builder.zext(LeftVal, RightVal.type, name="zext_left_deref")
        if Op in ('>>', 'RShift') and isinstance(Node.left, ast.Name):
            vname = Node.left.id
            if vname in Gen.var_signedness and Gen.var_signedness[vname]:
                is_unsigned = True
        result = Gen.emit_binary_op(Op, LeftVal, RightVal, is_unsigned=is_unsigned)
        return result

    def _HandleBoolOpLlvm(self, Node):
        Gen = self.Trans.LlvmGen

        def _to_bool(val, name):
            if isinstance(val.type, (ir.FloatType, ir.DoubleType)):
                return Gen.builder.fcmp_ordered('!=', val, ir.Constant(val.type, 0.0), name=name)
            if not isinstance(val.type, ir.IntType) or val.type.width != 1:
                return Gen.builder.icmp_signed('!=', val, Gen._ZeroConst(val.type), name=name)
            return val

        values = Node.values
        if len(values) < 2:
            return self.HandleExprLlvm(values[0]) if values else None
        if isinstance(Node.op, ast.And):
            phi_incomings = []
            end_bb = Gen.func.append_basic_block(name="and.end")
            for idx in range(len(values) - 1):
                val = self.HandleExprLlvm(values[idx])
                if not val:
                    if not end_bb.is_terminated:
                        saved_block = Gen.builder.block
                        Gen.builder.position_at_start(end_bb)
                        Gen.builder.unreachable()
                        Gen.builder.position_at_end(saved_block)
                    return None
                if not isinstance(val.type, ir.IntType) or val.type.width != 1:
                    val = _to_bool(val, f"andcond{idx}")
                cur_bb = Gen.builder.block
                next_bb = Gen.func.append_basic_block(name=f"and.rhs{idx}")
                Gen.builder.cbranch(val, next_bb, end_bb)
                phi_incomings.append((ir.Constant(ir.IntType(1), 0), cur_bb))
                Gen.builder.position_at_start(next_bb)
            last_val = self.HandleExprLlvm(values[-1])
            if not last_val:
                last_val = ir.Constant(ir.IntType(1), 0)
            if not isinstance(last_val.type, ir.IntType) or last_val.type.width != 1:
                last_val = _to_bool(last_val, "andcond_last")
            last_bb = Gen.builder.block
            Gen.builder.branch(end_bb)
            Gen.builder.position_at_start(end_bb)
            result_phi = Gen.builder.phi(ir.IntType(1), name="and.result")
            for const_val, bb in phi_incomings:
                result_phi.add_incoming(const_val, bb)
            result_phi.add_incoming(last_val, last_bb)
            return result_phi
        elif isinstance(Node.op, ast.Or):
            phi_incomings = []
            end_bb = Gen.func.append_basic_block(name="or.end")
            for idx in range(len(values) - 1):
                val = self.HandleExprLlvm(values[idx])
                if not val:
                    if not end_bb.is_terminated:
                        saved_block = Gen.builder.block
                        Gen.builder.position_at_start(end_bb)
                        Gen.builder.unreachable()
                        Gen.builder.position_at_end(saved_block)
                    return None
                if not isinstance(val.type, ir.IntType) or val.type.width != 1:
                    val = _to_bool(val, f"orcond{idx}")
                cur_bb = Gen.builder.block
                next_bb = Gen.func.append_basic_block(name=f"or.rhs{idx}")
                Gen.builder.cbranch(val, end_bb, next_bb)
                phi_incomings.append((ir.Constant(ir.IntType(1), 1), cur_bb))
                Gen.builder.position_at_start(next_bb)
            last_val = self.HandleExprLlvm(values[-1])
            if not last_val:
                last_val = ir.Constant(ir.IntType(1), 1)
            if not isinstance(last_val.type, ir.IntType) or last_val.type.width != 1:
                last_val = _to_bool(last_val, "orcond_last")
            last_bb = Gen.builder.block
            Gen.builder.branch(end_bb)
            Gen.builder.position_at_start(end_bb)
            result_phi = Gen.builder.phi(ir.IntType(1), name="or.result")
            for const_val, bb in phi_incomings:
                result_phi.add_incoming(const_val, bb)
            result_phi.add_incoming(last_val, last_bb)
            return result_phi
        return None

    def _HandleUnaryOpLlvm(self, Node):
        Gen = self.Trans.LlvmGen
        OperandVal = self.HandleExprLlvm(Node.operand)
        if not OperandVal:
            return None
        OpSymbol = self.Trans.ExprHandler.GetUnaryOpSymbol(Node.op)
        if not OpSymbol:
            return None
        UNARY_OVERLOAD_MAP = {
            '-': '__neg__',
            '+': '__pos__',
            '~': '__invert__',
        }
        if OpSymbol in UNARY_OVERLOAD_MAP:
            ClassName = None
            if isinstance(OperandVal.type, ir.PointerType):
                pointee = OperandVal.type.pointee
                if isinstance(pointee, (ir.IdentifiedStructType, ir.LiteralStructType)):
                    for CN, ST in Gen.structs.items():
                        if pointee == ST:
                            ClassName = CN
                            break
                elif isinstance(pointee, ir.IntType) and pointee.width == 8:
                    ClassName = self.Trans.ExprUtils._get_var_class(Node.operand, Gen)
                    if ClassName and ClassName in Gen.structs:
                        TargetStructPtr = ir.PointerType(Gen.structs[ClassName])
                        OperandVal = Gen.builder.bitcast(OperandVal, TargetStructPtr, name=f"bitcast_to_{ClassName}")
            if ClassName is None and isinstance(Node.operand, ast.Name):
                ClassName = Gen.var_struct_class.get(Node.operand.id)
                if ClassName and ClassName in Gen.structs:
                    if isinstance(OperandVal.type, ir.PointerType) and isinstance(OperandVal.type.pointee, ir.IntType) and OperandVal.type.pointee.width == 8:
                        TargetStructPtr = ir.PointerType(Gen.structs[ClassName])
                        OperandVal = Gen.builder.bitcast(OperandVal, TargetStructPtr, name=f"bitcast_to_{ClassName}")
            if ClassName:
                op_name = UNARY_OVERLOAD_MAP[OpSymbol]
                result = self.Trans.ExprUtils._try_operator_overLoad(ClassName, op_name, OperandVal, Gen)
                if result is not None:
                    return result
        if OpSymbol == 'not':
            if isinstance(OperandVal.type, ir.IntType) and OperandVal.type.width == 1:
                return Gen.builder.not_(OperandVal, name="not_result")
            if isinstance(OperandVal.type, ir.PointerType):
                null_val = ir.Constant(OperandVal.type, None)
                return Gen.builder.icmp_signed('==', OperandVal, null_val, name="not_result")
            zero = ir.Constant(OperandVal.type, 0)
            return Gen.builder.icmp_signed('==', OperandVal, zero, name="not_result")
        elif OpSymbol == '~':
            if isinstance(OperandVal.type, ir.IntType):
                mask = (1 << OperandVal.type.width) - 1
                return Gen.builder.xor(OperandVal, ir.Constant(OperandVal.type, mask), name="bnot_result")
            return Gen.builder.xor(OperandVal, ir.Constant(OperandVal.type, -1), name="bnot_result")
        elif OpSymbol == '+':
            return OperandVal
        elif OpSymbol == '-':
            if isinstance(OperandVal.type, ir.IntType):
                return Gen.builder.neg(OperandVal, name="neg_result")
            elif isinstance(OperandVal.type, (ir.FloatType, ir.DoubleType)):
                return Gen.builder.fneg(OperandVal, name="fneg_result")
        return None

    def _HandleCompareLlvm(self, Node):
        Gen = self.Trans.LlvmGen
        LeftVal = self.HandleExprLlvm(Node.left)
        if not LeftVal:
            return None
        is_unsigned = Gen._check_node_unsigned(Node.left) or (Gen._check_node_unsigned(Node.comparators[0]) if Node.comparators else False)
        if len(Node.ops) == 1:
            Op = Node.ops[0]
            ComparatorSymbol = self.Trans.ExprHandler.GetComparatorSymbol(Op)
            if not ComparatorSymbol:
                return None
            RightVal = self.HandleExprLlvm(Node.comparators[0])
            if not RightVal:
                return None
            if isinstance(LeftVal, ir.Constant) and isinstance(RightVal, ir.Constant):
                if isinstance(LeftVal.type, ir.IntType) and isinstance(RightVal.type, ir.IntType):
                    lv, rv = LeftVal.constant, RightVal.constant
                    lw, rw = LeftVal.type.width, RightVal.type.width
                    if lw < rw:
                        lv = lv if lv >= 0 else lv + (1 << lw)
                    elif rw < lw:
                        rv = rv if rv >= 0 else rv + (1 << rw)
                    if is_unsigned:
                        mask_l = (1 << max(lw, rw)) - 1
                        lv = lv & mask_l
                        rv = rv & mask_l
                    cmp_map = {
                        '==': lv == rv, '!=': lv != rv,
                        '<': lv < rv, '>': lv > rv,
                        '<=': lv <= rv, '>=': lv >= rv,
                    }
                    result = cmp_map.get(ComparatorSymbol, False)
                    return ir.Constant(ir.IntType(1), 1 if result else 0)
            if isinstance(LeftVal.type, ir.IntType) and isinstance(RightVal.type, ir.IntType):
                if LeftVal.type.width < RightVal.type.width:
                    # 如果宽类型常量值超出窄类型有符号范围，则必须用zext
                    # i8 类型默认使用 zext（系统编程中 i8 几乎总是无符号字节）
                    need_zext = is_unsigned or Gen._check_node_unsigned(Node.left) or LeftVal.type.width == 8
                    if not need_zext and isinstance(RightVal, ir.Constant):
                        signed_max = (1 << (LeftVal.type.width - 1)) - 1
                        try:
                            if RightVal.constant > signed_max or RightVal.constant < 0:
                                need_zext = True
                        except (TypeError, ValueError):
                            pass
                    if need_zext:
                        LeftVal = Gen.builder.zext(LeftVal, RightVal.type, name="zext_left_cmp")
                    else:
                        LeftVal = Gen.builder.sext(LeftVal, RightVal.type, name="sext_left_cmp")
                elif LeftVal.type.width > RightVal.type.width:
                    need_zext = is_unsigned or Gen._check_node_unsigned(Node.comparators[0]) or RightVal.type.width == 8
                    if not need_zext and isinstance(LeftVal, ir.Constant):
                        signed_max = (1 << (RightVal.type.width - 1)) - 1
                        try:
                            if LeftVal.constant > signed_max or LeftVal.constant < 0:
                                need_zext = True
                        except (TypeError, ValueError):
                            pass
                    if need_zext:
                        RightVal = Gen.builder.zext(RightVal, LeftVal.type, name="zext_right_cmp")
                    else:
                        RightVal = Gen.builder.sext(RightVal, LeftVal.type, name="sext_right_cmp")
            elif isinstance(LeftVal.type, ir.PointerType) and isinstance(RightVal.type, ir.IntType):
                if isinstance(LeftVal.type.pointee, ir.IntType) and LeftVal.type.pointee.width == 8 and isinstance(RightVal.type, ir.IntType) and RightVal.type.width == 8:
                    LeftVal = Gen._load(LeftVal, name="Load_char_ptr_cmp")
                else:
                    if isinstance(RightVal, ir.Constant) and RightVal.constant == 0:
                        RightVal = ir.Constant(LeftVal.type, None)
                    else:
                        RightVal = Gen.builder.inttoptr(RightVal, LeftVal.type, name="int2ptr_cmp")
            elif isinstance(RightVal.type, ir.PointerType) and isinstance(LeftVal.type, ir.IntType):
                if isinstance(RightVal.type.pointee, ir.IntType) and RightVal.type.pointee.width == 8 and isinstance(LeftVal.type, ir.IntType) and LeftVal.type.width == 8:
                    RightVal = Gen._load(RightVal, name="Load_char_ptr_cmp")
                else:
                    if isinstance(LeftVal, ir.Constant) and LeftVal.constant == 0:
                        LeftVal = ir.Constant(RightVal.type, None)
                    else:
                        LeftVal = Gen.builder.inttoptr(LeftVal, RightVal.type, name="int2ptr_cmp")
            if isinstance(LeftVal.type, ir.PointerType) and isinstance(RightVal.type, ir.PointerType):
                result = Gen.builder.icmp_unsigned(ComparatorSymbol, LeftVal, RightVal, name="cmpresult")
            elif isinstance(LeftVal.type, (ir.FloatType, ir.DoubleType)) or isinstance(RightVal.type, (ir.FloatType, ir.DoubleType)):
                if isinstance(LeftVal.type, (ir.FloatType, ir.DoubleType)) and isinstance(RightVal.type, ir.IntType):
                    RightVal = Gen.builder.sitofp(RightVal, LeftVal.type, name="int_to_float_cmp")
                elif isinstance(RightVal.type, (ir.FloatType, ir.DoubleType)) and isinstance(LeftVal.type, ir.IntType):
                    LeftVal = Gen.builder.sitofp(LeftVal, RightVal.type, name="int_to_float_cmp")
                result = Gen.builder.fcmp_unordered(ComparatorSymbol, LeftVal, RightVal, name="cmpresult")
            elif is_unsigned:
                result = Gen.builder.icmp_unsigned(ComparatorSymbol, LeftVal, RightVal, name="cmpresult")
            else:
                result = Gen.builder.icmp_signed(ComparatorSymbol, LeftVal, RightVal, name="cmpresult")
            return result
        EndBB = Gen.func.append_basic_block(name="cmp.end")
        result_val = None
        PrevBB = Gen.builder.block
        for i, (Op, Comparator) in enumerate(zip(Node.ops, Node.comparators)):
            ComparatorSymbol = self.Trans.ExprHandler.GetComparatorSymbol(Op)
            if not ComparatorSymbol:
                continue
            RightVal = self.HandleExprLlvm(Comparator)
            if not RightVal:
                continue
            NextBB = Gen.func.append_basic_block(name="cmp.next")
            if isinstance(LeftVal.type, ir.IntType) and isinstance(RightVal.type, ir.IntType):
                if LeftVal.type.width < RightVal.type.width:
                    need_zext = is_unsigned or Gen._check_node_unsigned(Node.left) or LeftVal.type.width == 8
                    if not need_zext and isinstance(RightVal, ir.Constant):
                        signed_max = (1 << (LeftVal.type.width - 1)) - 1
                        try:
                            if RightVal.constant > signed_max or RightVal.constant < 0:
                                need_zext = True
                        except (TypeError, ValueError):
                            pass
                    if need_zext:
                        LeftVal = Gen.builder.zext(LeftVal, RightVal.type, name=f"zext_left_{i}")
                    else:
                        LeftVal = Gen.builder.sext(LeftVal, RightVal.type, name=f"sext_left_{i}")
                elif LeftVal.type.width > RightVal.type.width:
                    need_zext = is_unsigned or Gen._check_node_unsigned(Comparator) or RightVal.type.width == 8
                    if not need_zext and isinstance(LeftVal, ir.Constant):
                        signed_max = (1 << (RightVal.type.width - 1)) - 1
                        try:
                            if LeftVal.constant > signed_max or LeftVal.constant < 0:
                                need_zext = True
                        except (TypeError, ValueError):
                            pass
                    if need_zext:
                        RightVal = Gen.builder.zext(RightVal, LeftVal.type, name=f"zext_right_{i}")
                    else:
                        RightVal = Gen.builder.sext(RightVal, LeftVal.type, name=f"sext_right_{i}")
            if is_unsigned:
                cmp_result = Gen.builder.icmp_unsigned(ComparatorSymbol, LeftVal, RightVal, name=f"cmp_{i}")
            elif isinstance(LeftVal.type, (ir.FloatType, ir.DoubleType)) or isinstance(RightVal.type, (ir.FloatType, ir.DoubleType)):
                if isinstance(LeftVal.type, (ir.FloatType, ir.DoubleType)) and isinstance(RightVal.type, ir.IntType):
                    RightVal = Gen.builder.sitofp(RightVal, LeftVal.type, name=f"int_to_float_cmp_{i}")
                elif isinstance(RightVal.type, (ir.FloatType, ir.DoubleType)) and isinstance(LeftVal.type, ir.IntType):
                    LeftVal = Gen.builder.sitofp(LeftVal, RightVal.type, name=f"int_to_float_cmp_{i}")
                cmp_result = Gen.builder.fcmp_unordered(ComparatorSymbol, LeftVal, RightVal, name=f"cmp_{i}")
            else:
                cmp_result = Gen.builder.icmp_signed(ComparatorSymbol, LeftVal, RightVal, name=f"cmp_{i}")
            if result_val is None:
                result_val = cmp_result
            else:
                result_val = Gen.builder.and_(result_val, cmp_result, name=f"and_cmp_{i}")
            Gen.builder.cbranch(cmp_result, NextBB, EndBB)
            Gen.builder.position_at_start(NextBB)
            LeftVal = RightVal
        Gen.builder.branch(EndBB)
        Gen.builder.position_at_start(EndBB)
        final_result = Gen.builder.phi(ir.IntType(1), name="final_cmp")
        final_result.add_incoming(ir.Constant(ir.IntType(1), 0), PrevBB)
        for i in range(len(Node.ops)):
            bb = Gen.func.append_basic_block(name=f"cmp.end.{i}")
            final_result.add_incoming(ir.Constant(ir.IntType(1), 1), bb)
        return final_result