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


class TryHandle(BaseHandle):
    def _get_exception_type_code(self, type_node: ast.expr) -> int:
        if isinstance(type_node, ast.Name):
            name: str = type_node.id
            if name in EXCEPTION_CODE_MAP:
                return EXCEPTION_CODE_MAP[name]
            if name in self.Trans.exception_registry:
                return self.Trans.exception_registry[name]
            return 1
        if isinstance(type_node, ast.Attribute):
            name: str = type_node.attr
            if name in EXCEPTION_CODE_MAP:
                return EXCEPTION_CODE_MAP[name]
            if name in self.Trans.exception_registry:
                return self.Trans.exception_registry[name]
            return 1
        if isinstance(type_node, ast.Tuple):
            for elt in type_node.elts:
                code: int = self._get_exception_type_code(elt)
                if code != 1:
                    return code
            return 1
        return 1

    def _get_exception_name(self, type_node: ast.expr) -> str | None:
        if isinstance(type_node, ast.Name):
            return type_node.id
        if isinstance(type_node, ast.Attribute):
            return type_node.attr
        return None

    def _get_all_match_codes(self, type_node: ast.expr) -> list[int]:
        name: str | None = self._get_exception_name(type_node)
        if name is None:
            return [self._get_exception_type_code(type_node)]
        if name not in self.Trans.exception_registry:
            return [self._get_exception_type_code(type_node)]
        codes: list[int] = [self.Trans.exception_registry[name]]
        self._collect_descendant_codes(name, codes)
        return codes

    def _collect_descendant_codes(self, parent_name: str, codes: list[int]) -> None:
        for child_name, parent in self.Trans.exception_parents.items():
            if parent == parent_name:
                child_code: int | None = self.Trans.exception_registry.get(child_name)
                if child_code is not None and child_code not in codes:
                    codes.append(child_code)
                self._collect_descendant_codes(child_name, codes)

    def _HandleTryLlvm(self, Node: ast.Try) -> None:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        if not Gen.func:
            return
        exception_code: ir.AllocaInstr = Gen._allocaEntry(ir.IntType(32), name="eh_code")
        eh_message: ir.AllocaInstr = Gen._allocaEntry(ir.PointerType(ir.IntType(8)), name="eh_message")
        Gen._store(ir.Constant(ir.IntType(32), 0), exception_code)
        Gen._store(ir.Constant(ir.PointerType(ir.IntType(8)), None), eh_message)
        HasFinally: bool = bool(Node.finalbody)
        if HasFinally:
            Gen.eh_finally_stack.append((None, None, None))
        else:
            Gen.eh_finally_stack.append(None)
        DispatchBB: ir.Block = Gen.func.append_basic_block(name="try.dispatch")
        AfterBB: ir.Block = Gen.func.append_basic_block(name="try.after")
        except_blocks: list[tuple[ir.Block, list[int], ast.ExceptHandler]] = []
        for i, handler in enumerate(Node.handlers):
            ExcTypeCodes: list[int] = [1]
            if handler.type:
                ExcTypeCodes = self._get_all_match_codes(handler.type)
            ExceptBB: ir.Block = Gen.func.append_basic_block(name=f"try.except_{i}")
            Gen.eh_except_block_stack.append((DispatchBB, None, exception_code, eh_message))
            except_blocks.append((ExceptBB, ExcTypeCodes, handler))
        TryBB: ir.Block = Gen.func.append_basic_block(name="try.body")
        ElseBB: ir.Block | None = None
        if Node.orelse:
            ElseBB = Gen.func.append_basic_block(name="try.else")
        Gen.builder.branch(TryBB)
        Gen.builder.position_at_start(TryBB)
        self.HandleBodyLlvm(Node.body)
        if not Gen.builder.block.is_terminated:
            Gen._store(ir.Constant(ir.IntType(32), 0), exception_code)
            if ElseBB:
                Gen.builder.branch(ElseBB)
            else:
                Gen.builder.branch(AfterBB)
        Gen.builder.position_at_start(DispatchBB)
        eh_code_val: ir.Value = Gen._load(exception_code, name="Load_eh_code")
        next_check_bb: ir.Block | None = None
        handler_var_map: dict[int, tuple[str, ir.AllocaInstr, ir.Value | None]] = {}
        UnhandledBB: ir.Block = Gen.func.append_basic_block(name="try.unhandled")
        for i, (ExceptBB, ExcTypeCodes, handler) in enumerate(except_blocks):
            MatchBB: ir.Block = Gen.func.append_basic_block(name=f"try.match_{i}")
            if i > 0 and next_check_bb is not None:
                Gen.builder.position_at_start(next_check_bb)
                eh_code_val = Gen._load(exception_code, name=f"Load_eh_code_{i}")
            if handler.type is None:
                Gen.builder.branch(MatchBB)
            else:
                match_cond: ir.Value | None = None
                for code in ExcTypeCodes:
                    is_code_match: ir.Value = Gen.builder.icmp_signed('==', eh_code_val, ir.Constant(ir.IntType(32), code), name=f"eh_match_{i}_c{code}")
                    if match_cond is None:
                        match_cond = is_code_match
                    else:
                        match_cond = Gen.builder.or_(match_cond, is_code_match, name=f"eh_match_any_{i}")
                if i < len(except_blocks) - 1:
                    next_check_bb = Gen.func.append_basic_block(name=f"try.check_{i}")
                else:
                    next_check_bb = UnhandledBB
                Gen.builder.cbranch(match_cond, MatchBB, next_check_bb)
            Gen.builder.position_at_start(MatchBB)
            if handler.name:
                var_alloca: ir.AllocaInstr = Gen._allocaEntry(ir.IntType(8).as_pointer(), name=handler.name)
                eh_msg_val: ir.Value = Gen._load(eh_message, name=f"Load_eh_msg_{handler.name}")
                Gen._store(eh_msg_val, var_alloca)
                handler_var_map[i] = (handler.name, var_alloca, Gen.variables.get(handler.name))
            Gen.builder.branch(ExceptBB)
        Gen.builder.position_at_start(UnhandledBB)
        eh_msg_out_arg: ir.Argument | None = None
        eh_code_out_arg: ir.Argument | None = None
        if Gen.func and len(Gen.func.args) > 0:
            for arg in Gen.func.args:
                if arg.name == '__eh_msg_out__':
                    eh_msg_out_arg = arg
                elif arg.name == '__eh_code_out__':
                    eh_code_out_arg = arg
        if eh_msg_out_arg is not None and eh_code_out_arg is not None:
            eh_msg_val: ir.Value = Gen._load(eh_message, name="Load_eh_msg_propagate")
            eh_code_val_prop: ir.Value = Gen._load(exception_code, name="Load_eh_code_propagate")
            Gen._store(eh_msg_val, eh_msg_out_arg)
            Gen._store(eh_code_val_prop, eh_code_out_arg)
            ret_type: ir.Type | None = getattr(Gen.func, 'ftype', None) and Gen.func.ftype.return_type if getattr(Gen, 'func', None) else None
            if isinstance(ret_type, ir.IdentifiedStructType):
                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)
            elif isinstance(ret_type, ir.PointerType):
                Gen.builder.ret(ir.Constant(ret_type, None))
            elif isinstance(ret_type, (ir.FloatType, ir.DoubleType)):
                Gen.builder.ret(ir.Constant(ret_type, 0.0))
            elif isinstance(ret_type, ir.IntType):
                Gen.builder.ret(ir.Constant(ret_type, 1))
            elif isinstance(ret_type, ir.VoidType):
                Gen.builder.ret_void()
            else:
                Gen.builder.ret(ir.Constant(ir.IntType(32), 1))
        else:
            Gen._store(ir.Constant(ir.IntType(32), 0), exception_code)
            Gen.builder.branch(AfterBB)
        for idx, (ExceptBB, ExcTypeCodes, handler) in enumerate(except_blocks):
            Gen.builder.position_at_start(ExceptBB)
            old_var: ir.Value | None = None
            if idx in handler_var_map:
                hname: str = handler_var_map[idx][0]
                var_alloca: ir.AllocaInstr = handler_var_map[idx][1]
                old_var = handler_var_map[idx][2]
                Gen.variables[hname] = var_alloca
            self.HandleBodyLlvm(handler.body)
            if idx in handler_var_map:
                hname = handler_var_map[idx][0]
                old_var = handler_var_map[idx][2]
                if old_var is not None:
                    Gen.variables[hname] = old_var
                elif hname in Gen.variables:
                    del Gen.variables[hname]
            if not Gen.builder.block.is_terminated:
                Gen.builder.branch(AfterBB)
            Gen.eh_except_block_stack.pop()
        if ElseBB:
            Gen.builder.position_at_start(ElseBB)
            self.HandleBodyLlvm(Node.orelse)
            if not Gen.builder.block.is_terminated:
                Gen.builder.branch(AfterBB)
        Gen.builder.position_at_start(AfterBB)
        if HasFinally:
            Gen.eh_finally_stack.pop()
        if Node.finalbody:
            FinallyBB: ir.Block = Gen.func.append_basic_block(name="try.finally")
            pending_ret_flag: ir.AllocaInstr = Gen._allocaEntry(ir.IntType(32), name="pending_ret_flag")
            pending_ret_val: ir.AllocaInstr | None = Gen._allocaEntry(ir.IntType(32), name="pending_ret_val") if Gen.func.type.pointee.return_type == ir.IntType(32) else None
            Gen.eh_finally_stack.append((FinallyBB, pending_ret_flag, pending_ret_val))
            Gen.builder.branch(FinallyBB)
            Gen.builder.position_at_start(FinallyBB)
            Gen._store(ir.Constant(ir.IntType(32), 0), pending_ret_flag)
            self.HandleBodyLlvm(Node.finalbody)
            if Gen.eh_finally_stack:
                _, pending_ret_flag, pending_ret_val = Gen.eh_finally_stack[-1]
                Gen.eh_finally_stack.pop()
            ret_flag_val: ir.Value = Gen._load(pending_ret_flag, name="Load_ret_flag")
            ret_flag_is_set: ir.Value = Gen.builder.icmp_signed('!=', ret_flag_val, ir.Constant(ir.IntType(32), 0), name="check_ret_flag")
            RetBB: ir.Block = Gen.func.append_basic_block(name="finally.ret")
            ContinueBB: ir.Block = Gen.func.append_basic_block(name="try.continue")
            Gen.builder.cbranch(ret_flag_is_set, RetBB, ContinueBB)
            Gen.builder.position_at_start(RetBB)
            if pending_ret_val:
                ret_val: ir.Value = Gen._load(pending_ret_val, name="Load_ret_val")
                Gen.builder.ret(ret_val)
            else:
                Gen.builder.ret_void()
            Gen.builder.position_at_start(ContinueBB)