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 AssertHandle(BaseHandle):
    def _HandleAssertLlvm(self, Node: ast.Assert) -> None:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        cond_val: ir.Value | None = self.HandleExprLlvm(Node.test)
        if not cond_val:
            return
        if not isinstance(cond_val.type, ir.IntType) or cond_val.type.width != 1:
            zero: ir.Constant = ir.Constant(cond_val.type, 0)
            cond_val = Gen.builder.icmp_signed('!=', cond_val, zero, name="assert_cond")
        AssertFailBB: ir.Block = Gen.func.append_basic_block(name="assert.fail")
        AssertOkBB: ir.Block = Gen.func.append_basic_block(name="assert.ok")
        Gen.builder.cbranch(cond_val, AssertOkBB, AssertFailBB)
        Gen.builder.position_at_start(AssertFailBB)
        exc_code: int = 9
        exc_msg: ir.Value | None = None
        if Node.msg:
            if isinstance(Node.msg, ast.Constant) and isinstance(Node.msg.value, str):
                exc_msg = self.HandleExprLlvm(Node.msg)
            elif isinstance(Node.msg, ast.Call) and isinstance(Node.msg.func, ast.Name):
                exc_code = EXCEPTION_CODE_MAP.get(Node.msg.func.id, 99)
                if Node.msg.args:
                    first_arg: ast.expr = Node.msg.args[0]
                    if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str):
                        exc_msg = self.HandleExprLlvm(first_arg)
            elif isinstance(Node.msg, ast.Name):
                exc_code = EXCEPTION_CODE_MAP.get(Node.msg.id, 99)
            else:
                exc_msg = self.HandleExprLlvm(Node.msg)
        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(ir.Constant(ir.IntType(32), exc_code), 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)
            Gen.builder.branch(ExceptBB)
        else:
            printf_func: ir.Function = Gen.get_or_declare_c_func('printf', ir.FunctionType(ir.IntType(32), [ir.PointerType(ir.IntType(8))], var_arg=True))
            if exc_msg:
                fmt: ir.Value = Gen.emit_constant("AssertionError: %s\n", 'string')
                Gen.builder.call(printf_func, [fmt, exc_msg], name="assert_fail_print")
            else:
                fmt: ir.Value = Gen.emit_constant("AssertionError\n", 'string')
                Gen.builder.call(printf_func, [fmt], name="assert_fail_print")
            exit_func: ir.Function = Gen.get_or_declare_c_func('exit', ir.FunctionType(ir.VoidType(), [ir.IntType(32)]))
            Gen.builder.call(exit_func, [ir.Constant(ir.IntType(32), 1)], name="call_exit")
            Gen.builder.unreachable()
        Gen.builder.position_at_start(AssertOkBB)