from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from lib.core.translator import Translator
import ast
import llvmlite.ir as ir
from lib.core.Handles.HandlesBase import BaseHandle, EXCEPTION_CODE_MAP
from lib.constants.config import mode as _config_mode


class RaiseHandle(BaseHandle):
    def _get_exception_code(self, ExcName: str) -> int:
        if ExcName in EXCEPTION_CODE_MAP:
            return EXCEPTION_CODE_MAP[ExcName]
        if ExcName in self.Trans.exception_registry:
            return self.Trans.exception_registry[ExcName]
        return 1

    def _HandleRaiseLlvm(self, Node: ast.Raise) -> None:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        exc_val: ir.Constant = ir.Constant(ir.IntType(32), 1)
        exc_msg: ir.Value | None = None
        IsStopIteration: bool = False
        if Node.exc:
            if isinstance(Node.exc, ast.Constant) and isinstance(Node.exc.value, int):
                exc_val = ir.Constant(ir.IntType(32), Node.exc.value)
            elif isinstance(Node.exc, ast.Name):
                ExcName: str = Node.exc.id
                if ExcName == 'StopIteration':
                    IsStopIteration = True
                exc_val = ir.Constant(ir.IntType(32), self._get_exception_code(ExcName))
            elif isinstance(Node.exc, ast.Call) and isinstance(Node.exc.func, ast.Name):
                ExcName: str = Node.exc.func.id
                if ExcName == 'StopIteration':
                    IsStopIteration = True
                exc_val = ir.Constant(ir.IntType(32), self._get_exception_code(ExcName))
                if Node.exc.args:
                    first_arg: ast.expr = Node.exc.args[0]
                    if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str):
                        exc_msg = self.HandleExprLlvm(first_arg)
                    elif isinstance(first_arg, (ast.Name, ast.Call, ast.Attribute)):
                        arg_val: ir.Value | None = self.HandleExprLlvm(first_arg)
                        if arg_val:
                            if isinstance(arg_val.type, ir.PointerType):
                                if self._is_char_pointer(arg_val):
                                    exc_msg = arg_val
                                else:
                                    try:
                                        exc_msg = Gen.builder.bitcast(arg_val, ir.PointerType(ir.IntType(8)), name="exc_msg_cast")
                                    except Exception as _e:
                                        if _config_mode == "strict":
                                            self.Trans.LogWarning(f"异常被忽略: {_e}")
                            else:
                                try:
                                    exc_msg = Gen.builder.inttoptr(arg_val, ir.PointerType(ir.IntType(8)), name="exc_msg_ptr")
                                except Exception as _e:
                                    if _config_mode == "strict":
                                        self.Trans.LogWarning(f"异常被忽略: {_e}")
            else:
                val: ir.Value | None = self.HandleExprLlvm(Node.exc)
                if val:
                    if isinstance(val.type, ir.IntType):
                        if val.type.width != 32:
                            try:
                                val = Gen.builder.trunc(val, ir.IntType(32), name="exc_trunc")
                            except Exception:  # 回退：trunc 失败时设默认值 1
                                val = ir.Constant(ir.IntType(32), 1)
                        exc_val = val
                    else:
                        try:
                            val = Gen.builder.ptrtoint(val, ir.IntType(32), name="exc_ptr2int")
                            exc_val = val
                        except Exception as _e:
                            if _config_mode == "strict":
                                self.Trans.LogWarning(f"异常被忽略: {_e}")
        if IsStopIteration:
            if Gen._stop_iter_flag_param is not None:
                Gen.builder.store(ir.Constant(ir.IntType(1), 1), Gen._stop_iter_flag_param)
            if Gen.func and Gen.func.type.pointee.return_type != ir.VoidType():
                ret_type: ir.Type = Gen.func.type.pointee.return_type
                if isinstance(ret_type, ir.PointerType):
                    Gen.builder.ret(ir.Constant(ret_type, None))
                elif isinstance(ret_type, ir.IntType):
                    Gen.builder.ret(ir.Constant(ret_type, 0))
                elif isinstance(ret_type, (ir.FloatType, ir.DoubleType)):
                    Gen.builder.ret(ir.Constant(ret_type, 0.0))
                elif isinstance(ret_type, ir.BaseStructType):
                    if ret_type.elements:
                        zero_val: ir.Constant = ir.Constant(ret_type, [
                            ir.Constant(et, None) if isinstance(et, (ir.PointerType, ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType))
                            else ir.Constant(et, 0) if isinstance(et, (ir.IntType, ir.FloatType, ir.DoubleType))
                            else ir.Constant(et, ir.Undefined) for et in ret_type.elements])
                    else:
                        zero_val = ir.Constant(ret_type, None)
                    Gen.builder.ret(zero_val)
                else:
                    Gen.builder.ret(ir.Constant(ret_type, 0))
            else:
                Gen.builder.ret_void()
            return
        if Gen.eh_except_block_stack:
            ExceptBB: ir.Block = Gen.eh_except_block_stack[-1][0]
            EndBB: ir.Block | None = Gen.eh_except_block_stack[-1][1]
            exception_code: ir.AllocaInstr = Gen.eh_except_block_stack[-1][2]
            eh_message: ir.AllocaInstr = Gen.eh_except_block_stack[-1][3]
            Gen._store(exc_val, exception_code)
            if exc_msg:
                Gen._store(exc_msg, eh_message)
            else:
                null_ptr: ir.Constant = ir.Constant(ir.PointerType(ir.IntType(8)), None)
                Gen._store(null_ptr, eh_message)
            if Gen.func and ExceptBB and ExceptBB.function is Gen.func:
                Gen.builder.branch(ExceptBB)
            else:
                eh_msg_arg: ir.Argument | None = None
                eh_code_arg: ir.Argument | None = None
                if Gen.func:
                    for arg in Gen.func.args:
                        if arg.name == '__eh_msg_out__':
                            eh_msg_arg = arg
                        elif arg.name == '__eh_code_out__':
                            eh_code_arg = arg
                if eh_msg_arg is not None:
                    if exc_msg:
                        Gen.builder.store(exc_msg, eh_msg_arg)
                    else:
                        null_ptr = ir.Constant(ir.PointerType(ir.IntType(8)), None)
                        Gen.builder.store(null_ptr, eh_msg_arg)
                if eh_code_arg is not None:
                    Gen.builder.store(exc_val, eh_code_arg)
                if Gen.func and Gen.func.type.pointee.return_type == ir.IntType(32):
                    Gen.builder.ret(ir.Constant(ir.IntType(32), 1))
                else:
                    Gen.builder.ret_void()
        else:
            eh_msg_arg: ir.Argument | None = None
            eh_code_arg: ir.Argument | None = None
            if Gen.func:
                for arg in Gen.func.args:
                    if arg.name == '__eh_msg_out__':
                        eh_msg_arg = arg
                    elif arg.name == '__eh_code_out__':
                        eh_code_arg = arg
            if eh_msg_arg is not None:
                if exc_msg:
                    Gen.builder.store(exc_msg, eh_msg_arg)
                else:
                    null_ptr: ir.Constant = ir.Constant(ir.PointerType(ir.IntType(8)), None)
                    Gen.builder.store(null_ptr, eh_msg_arg)
            if eh_code_arg is not None:
                Gen.builder.store(exc_val, eh_code_arg)
            if Gen.func and Gen.func.type.pointee.return_type == ir.IntType(32):
                Gen.builder.ret(ir.Constant(ir.IntType(32), 1))
            else:
                Gen.builder.ret_void()