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 llvmlite.ir as ir
import ast


class IfHandle(BaseHandle):
    def _get_attr_full_name(self, node):
        if isinstance(node, ast.Attribute):
            parts = []
            cur = node
            while isinstance(cur, ast.Attribute):
                parts.append(cur.attr)
                cur = cur.value
            if isinstance(cur, ast.Name):
                parts.append(cur.id)
            parts.reverse()
            return '.'.join(parts)
        return None

    def _is_cif_call(self, node):
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Attribute):
                if isinstance(node.func.value, ast.Name) and node.func.value.id == 'c':
                    return node.func.attr in ('CIf', 'CIfdef', 'CIfndef')
            if isinstance(node.func, ast.Name):
                return node.func.id in ('CIf', 'CIfdef', 'CIfndef')
        return False

    def _resolve_macro_name(self, arg):
        if isinstance(arg, ast.Name):
            return arg.id
        if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
            return arg.value
        if isinstance(arg, ast.Attribute):
            full_name = self._get_attr_full_name(arg)
            if full_name:
                return full_name
        return None

    def _evaluate_cif_condition(self, node):
        if not isinstance(node, ast.Call):
            return None
        args = node.args
        if not args:
            return False
        if isinstance(node.func, ast.Attribute):
            kind = node.func.attr
        elif isinstance(node.func, ast.Name):
            kind = node.func.id
        else:
            return None
        if kind == 'CIfdef':
            name = self._resolve_macro_name(args[0])
            if name is None:
                return False
            return self._is_macro_defined(name)
        elif kind == 'CIfndef':
            name = self._resolve_macro_name(args[0])
            if name is None:
                return True
            return not self._is_macro_defined(name)
        elif kind == 'CIf':
            return self._eval_const_expr(args[0])
        return None

    def _is_macro_defined(self, name):
        Gen = self.Trans.LlvmGen
        if hasattr(Gen, '_define_constants') and name in Gen._define_constants:
            return True
        if name in self.Trans.SymbolTable:
            info = self.Trans.SymbolTable[name]
            if getattr(info, 'IsDefine', False):
                return True
        platform_macros = self._get_platform_macros()
        return name in platform_macros

    def _get_macro_value(self, name):
        Gen = self.Trans.LlvmGen
        if hasattr(Gen, '_define_constants') and name in Gen._define_constants:
            val = Gen._define_constants[name]
            if isinstance(val, (int, float)):
                return val
        if name in self.Trans.SymbolTable:
            info = self.Trans.SymbolTable[name]
            if getattr(info, 'IsDefine', False):
                val = getattr(info, 'DefineValue', 0)
                if isinstance(val, (int, float)):
                    return val
        platform_macros = self._get_platform_macros()
        if name in platform_macros:
            return platform_macros[name]
        return None

    def _get_platform_macros(self):
        Gen = self.Trans.LlvmGen
        macros = {}
        pi = getattr(Gen, '_platform_info', {})
        ptr_size = getattr(Gen, 'ptr_size', 8)
        # 根据目标平台设置宏（从三元组推导）
        if pi.get('is_windows'):
            macros['_WIN32'] = 1
            if not pi.get('is_32bit'):
                macros['_WIN64'] = 1
                macros['WIN64'] = 1
            macros['WIN32'] = 1
        if pi.get('is_linux'):
            macros['__linux__'] = 1
        if pi.get('is_macos'):
            macros['__APPLE__'] = 1
            macros['__MACH__'] = 1
        if pi.get('is_x86_64'):
            macros['__x86_64__'] = 1
            macros['__x86_64'] = 1
            macros['_M_X64'] = 1
        elif pi.get('is_arm64'):
            macros['__aarch64__'] = 1
            macros['_M_ARM64'] = 1
        if not pi.get('is_32bit') and pi.get('is_linux'):
            macros['__LP64__'] = 1
        elif pi.get('is_32bit'):
            macros['_ILP32'] = 1
        macros['SIZEOF_VOID_P'] = ptr_size
        macros['__SIZEOF_POINTER__'] = ptr_size
        macros['NDEBUG'] = 0
        return macros

    def _eval_const_expr(self, node):
        if isinstance(node, ast.Constant):
            val = node.value
            if isinstance(val, bool):
                return 1 if val else 0
            if isinstance(val, int):
                return val
            if isinstance(val, float):
                return 1 if val != 0.0 else 0
            return 0
        if isinstance(node, ast.Name):
            name = node.id
            val = self._get_macro_value(name)
            if val is not None:
                return val
            return None
        if isinstance(node, ast.Attribute):
            full_key = self._get_attr_full_name(node)
            if full_key:
                val = self._get_macro_value(full_key)
                if val is not None:
                    return val
                short_name = full_key.split('.')[-1] if '.' in full_key else full_key
                val = self._get_macro_value(short_name)
                if val is not None:
                    return val
            return None
        if isinstance(node, ast.UnaryOp):
            operand = self._eval_const_expr(node.operand)
            if operand is None:
                return None
            if isinstance(node.op, ast.Not):
                return 1 if not operand else 0
            if isinstance(node.op, ast.USub):
                return -operand
            if isinstance(node.op, ast.Invert):
                return ~operand
            return None
        if isinstance(node, ast.BinOp):
            left = self._eval_const_expr(node.left)
            right = self._eval_const_expr(node.right)
            if left is None or right is None:
                return None
            try:
                if isinstance(node.op, ast.Add):
                    return left + right
                elif isinstance(node.op, ast.Sub):
                    return left - right
                elif isinstance(node.op, ast.Mult):
                    return left * right
                elif isinstance(node.op, ast.FloorDiv):
                    return left // right if right != 0 else 0
                elif isinstance(node.op, ast.Mod):
                    return left % right if right != 0 else 0
                elif isinstance(node.op, ast.LShift):
                    return left << right
                elif isinstance(node.op, ast.RShift):
                    return left >> right
                elif isinstance(node.op, ast.BitOr):
                    return left | right
                elif isinstance(node.op, ast.BitAnd):
                    return left & right
                elif isinstance(node.op, ast.BitXor):
                    return left ^ right
            except Exception:
                return None
            return None
        if isinstance(node, ast.BoolOp):
            if isinstance(node.op, ast.And):
                for val in node.values:
                    result = self._eval_const_expr(val)
                    if result is None:
                        return None
                    if not result:
                        return 0
                return 1
            elif isinstance(node.op, ast.Or):
                for val in node.values:
                    result = self._eval_const_expr(val)
                    if result is None:
                        return None
                    if result:
                        return 1
                return 0
        if isinstance(node, ast.Compare):
            left = self._eval_const_expr(node.left)
            if left is None:
                return None
            for op, comparator in zip(node.ops, node.comparators):
                right = self._eval_const_expr(comparator)
                if right is None:
                    return None
                if isinstance(op, ast.Eq):
                    result = left == right
                elif isinstance(op, ast.NotEq):
                    result = left != right
                elif isinstance(op, ast.Lt):
                    result = left < right
                elif isinstance(op, ast.LtE):
                    result = left <= right
                elif isinstance(op, ast.Gt):
                    result = left > right
                elif isinstance(op, ast.GtE):
                    result = left >= right
                else:
                    return None
                if not result:
                    return 0
            return 1
        if isinstance(node, ast.Call):
            if self._is_cif_call(node):
                return self._evaluate_cif_condition(node)
        return None

    def _HandleIfLlvm(self, Node):
        Gen = self.Trans.LlvmGen
        if self._is_cif_call(Node.test):
            compile_time_val = self._evaluate_cif_condition(Node.test)
            if compile_time_val is not None:
                if compile_time_val:
                    self.Trans.BodyHandler.HandleBodyLlvm(Node.body)
                elif Node.orelse:
                    if isinstance(Node.orelse, list) and len(Node.orelse) == 1 and isinstance(Node.orelse[0], ast.If):
                        self._HandleIfLlvm(Node.orelse[0])
                    else:
                        self.Trans.BodyHandler.HandleBodyLlvm(Node.orelse)
                return
        Cond = self.Trans.ExprHandler.HandleExprLlvm(Node.test)
        if not Cond:
            return
        is_const = isinstance(Cond, ir.Constant) and isinstance(Cond.type, ir.IntType) and Cond.type.width == 1
        if is_const:
            if Cond.constant:
                self.Trans.BodyHandler.HandleBodyLlvm(Node.body)
            elif Node.orelse:
                if isinstance(Node.orelse, list) and len(Node.orelse) == 1 and isinstance(Node.orelse[0], ast.If):
                    self._HandleIfLlvm(Node.orelse[0])
                else:
                    self.Trans.BodyHandler.HandleBodyLlvm(Node.orelse)
            return
        if not isinstance(Cond.type, ir.IntType) or Cond.type.width != 1:
            ClassName = self.Trans.ExprHandler._get_var_class(Node.test, Gen)
            if ClassName and Gen._has_function(f'{ClassName}.__bool__'):
                obj_val = Cond
                if isinstance(obj_val.type, ir.PointerType) and isinstance(obj_val.type.pointee, ir.IntType) and obj_val.type.pointee.width == 8:
                    obj_val = Gen.builder.bitcast(obj_val, ir.PointerType(Gen.structs[ClassName]), name=f"cast_{ClassName}")
                Cond = Gen.builder.call(Gen._get_function(f'{ClassName}.__bool__'), [obj_val], name=f"call_{ClassName}.__bool__")
                if isinstance(Cond.type, ir.IntType) and Cond.type.width != 1:
                    Cond = Gen.builder.icmp_signed('!=', Cond, Gen._ZeroConst(Cond.type), name="bool_result")
            else:
                if isinstance(Cond.type, (ir.FloatType, ir.DoubleType)):
                    Cond = Gen.builder.fcmp_ordered('!=', Cond, ir.Constant(Cond.type, 0.0), name="ifcond")
                else:
                    Cond = Gen.builder.icmp_signed('!=', Cond, Gen._ZeroConst(Cond.type), name="ifcond")
        ThenBB = Gen.func.append_basic_block(name="then")
        ElseBB = Gen.func.append_basic_block(name="else") if Node.orelse else None
        MergeBB = Gen.func.append_basic_block(name="endif")
        if not Gen.builder.block.is_terminated:
            Gen.builder.cbranch(Cond, ThenBB, ElseBB if ElseBB else MergeBB)
        Gen.builder.position_at_start(ThenBB)
        self.Trans.BodyHandler.HandleBodyLlvm(Node.body)
        if not Gen.builder.block.is_terminated:
            Gen.builder.branch(MergeBB)
        if ElseBB:
            Gen.builder.position_at_start(ElseBB)
            self.Trans.BodyHandler.HandleBodyLlvm(Node.orelse)
            if not Gen.builder.block.is_terminated:
                Gen.builder.branch(MergeBB)
        if not MergeBB.is_terminated:
            Gen.builder.position_at_start(MergeBB)
        else:
            Gen.builder.position_at_end(MergeBB)