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 vast as ast

# 无信息 assert 使用的默认 fail_msg；须与 testcheck.check 的 fail_msg: str 形参一致
_ASSERT_DEFAULT_FAIL_MSG: str = "assertion failed"
# assert 自动依赖的检查库模块名与入口函数（见 includes/testcheck.vp）
_ASSERT_CHECK_MODULE: str = "testcheck"
_ASSERT_CHECK_FUNC: str = "check"

class AssertHandle(BaseHandle):
    def _HandleAssertLlvm(self, Node: ast.Assert) -> None:
        """把 `assert 条件[, 信息]` 译为一次普通函数调用 testcheck.check(条件, None, 信息)。

        语义：条件为真 → ok_msg 传 None（静默，不打印 PASS）；条件为假 → fail(信息)。
        不再中止程序（旧实现的 printf AssertionError + exit(1) 已移除），也不再走
        try/except 抛异常码的路径。
        """
        FailMsgNode: ast.expr = self._ResolveAssertFailMsgNode(Node)
        CallNode: ast.Call = ast.Call(
            func=ast.Attribute(
                value=ast.Name(id=_ASSERT_CHECK_MODULE, ctx=ast.Load()),
                attr=_ASSERT_CHECK_FUNC,
                ctx=ast.Load()),
            args=[Node.test, ast.Constant(value=None, kind=None), FailMsgNode],
            keywords=[],
        )
        ast.copy_location(CallNode, Node)
        # 复用既有跨模块调用生成路径（含符号解析、实参按形参类型收窄/扩充与 fail-fast 校验）
        self.HandleExprLlvm(CallNode)

    def _ResolveAssertFailMsgNode(self, Node: ast.Assert) -> ast.expr:
        """解析 assert 的信息表达式，并强制其为 str（fail-fast，禁止静默降级）。

        - 无信息 → 字符串字面量 "assertion failed"
        - 有信息 → 原节点（须为 str 形态：字面量 / f-string / str 变量或属性）
        """
        MsgNode: ast.expr | None = Node.msg
        if MsgNode is None:
            DefaultMsg: ast.Constant = ast.Constant(value=_ASSERT_DEFAULT_FAIL_MSG, kind=None)
            ast.copy_location(DefaultMsg, Node)
            return DefaultMsg
        # 旧实现允许把异常类名当信息（映射为异常码、不打印）。check() 的 fail_msg 形参
        # 是 str，这类写法已无对应语义：显式报错，不静默降级。
        _ExcName: str | None = None
        if isinstance(MsgNode, ast.Name):
            _ExcName = MsgNode.id
        elif isinstance(MsgNode, ast.Call) and isinstance(MsgNode.func, ast.Name):
            _ExcName = MsgNode.func.id
        if _ExcName is not None and _ExcName in EXCEPTION_CODE_MAP:
            self._RaiseAssertMsgTypeError(
                f"assert 的信息必须是字符串（str），不能是异常类 '{_ExcName}'")
        if not self._IsStrAssertMsg(MsgNode):
            self._RaiseAssertMsgTypeError(
                "assert 的信息必须是字符串（str）：仅支持字符串字面量、f-string 或 str 变量/属性"
                "（如需拼接请先写入 str 变量）")
        return MsgNode

    def _IsStrAssertMsg(self, MsgNode: ast.expr) -> bool:
        """静态判定 assert 信息是否为 str 形态。

        变量/属性（Name/Attribute）无法在此静态证伪为 str，放行并由调用点的
        `fail_msg: str` 形参做类型收口；其余非常量、非 f-string 的形态一律拒绝。
        """
        if isinstance(MsgNode, ast.Constant):
            return isinstance(MsgNode.value, str)
        if isinstance(MsgNode, ast.JoinedStr):
            return True
        if isinstance(MsgNode, (ast.Name, ast.Attribute)):
            return True
        return False

    def _RaiseAssertMsgTypeError(self, Message: str) -> None:
        from lib.core.LLVMCG.FuncGen import FailFastError
        Gen = self.Trans.LlvmGen
        raise FailFastError(Message + (Gen._get_node_info() if Gen else ''))