import t, c
from stdint import *
import llvmlite
import memhub
import string
import stdlib
import viperlib
import w32.win32file as w32file
import lib.core.VLogger as VLogger

# ============================================================
# DebugInfoGen - function-level DWARF metadata emitter (llvmlite text IR)
#
# Node layout mirrors TPC lib/core/LLVMCG/DebugInfo.py node-for-node:
#   !DIFile -> !{ null } -> !DISubroutineType -> distinct !DICompileUnit
#     -> distinct !DISubprogram (per define) -> !DILocation (per statement)
#   module tail: !llvm.dbg.cu + !llvm.module.flags (Debug Info Version 3)
#
# Enabled only while --traceback is on: the Phase layer mirrors
# BuildPipeline._traceback_mode here via SetEnabled(). While disabled every
# entry point returns before touching any llvmlite DI API, so the emitted IR
# stays byte-identical to the pre-DWARF output.
#
# Scope/statement hooks are driven from the Gen layer:
#   Translator.translate          -> BeginModule / Finalize
#   FuncGenDef.translate_function_def / ClassGenVTable._translate_method /
#   FuncGenMixin (nested + lambda) -> PushFunction / PopFunction
#   StmtGenMixin.translate_stmt   -> SetStatement
# ============================================================

# ---- module-scope state (reset per module by BeginModule) ----
_g_enabled: int = 0
_g_finalized: int = 0
_g_mod: llvmlite.LLVMModule* = None
_g_file_path: str = None
_g_file_id: int = -1
_g_null_id: int = -1
_g_subtype_id: int = -1
_g_cu_id: int = -1
_g_cur_sp: int = -1
_g_cur_entry_ref: str = None
_g_cur_builder: llvmlite.IRBuilder* = None
_g_depth: int = 0

# Push/Pop save stack: enclosing (sp id + 1, entry DILocation ref, builder)
# per level. +1 so 0 means "no enclosing function" (sp ids are >= 0). The saved
# builder is what makes SetStatement reject a builder that is not the one the
# current DISubprogram was pushed on (e.g. an async state machine emitted inside
# a bracketed function): such instructions must carry no !dbg, otherwise LLVM's
# verifier rejects "!dbg attachment points at wrong subprogram for function".
DBG_STACK_MAX: t.CDefine = 256
_g_stk_sp: sizet* = None
_g_stk_ref: str = None
_g_stk_bld: str = None

# DILocation cache: direct-mapped, up to 8-way linear probe, key (sp, line, col).
DBG_LOC_SLOTS: t.CDefine = 8192
DBG_LOC_PROBE: t.CDefine = 8
DBG_LOC_KEY_STRIDE: t.CDefine = 3
_g_loc_key: sizet* = None
_g_loc_ref: str = None
_g_loc_used: bytes = None


# ============================================================
# SetEnabled / IsEnabled - mirror the --traceback switch
# ============================================================
def SetEnabled(flag: bool) -> int:
    """Enable/disable DI emission (mirrors BuildPipeline._traceback_mode)."""
    global _g_enabled
    if flag:
        _g_enabled = 1
    else:
        _g_enabled = 0
    return 0


def IsEnabled() -> int:
    return _g_enabled


# ============================================================
# _EnsureStorage - lazily allocate the long-lived stack / cache buffers
# ============================================================
def _EnsureStorage() -> int:
    global _g_stk_sp, _g_stk_ref, _g_stk_bld, _g_loc_key, _g_loc_ref, _g_loc_used
    if _g_stk_sp is not None:
        return 0
    sbuf1: bytes = stdlib.malloc(DBG_STACK_MAX * 8)
    sbuf2: bytes = stdlib.malloc(DBG_STACK_MAX * 8)
    sbuf3: bytes = stdlib.malloc(DBG_STACK_MAX * 8)
    kbuf: bytes = stdlib.malloc(DBG_LOC_SLOTS * DBG_LOC_KEY_STRIDE * 8)
    rbuf: bytes = stdlib.malloc(DBG_LOC_SLOTS * 8)
    ubuf: bytes = stdlib.malloc(DBG_LOC_SLOTS)
    if sbuf1 is None or sbuf2 is None or sbuf3 is None or kbuf is None or rbuf is None or ubuf is None:
        VLogger.error("--traceback: DWARF 元数据缓冲分配失败，终止编译", "DWARF")
        return 1
    string.memset(sbuf1, 0, DBG_STACK_MAX * 8)
    string.memset(sbuf2, 0, DBG_STACK_MAX * 8)
    string.memset(sbuf3, 0, DBG_STACK_MAX * 8)
    string.memset(kbuf, 0, DBG_LOC_SLOTS * DBG_LOC_KEY_STRIDE * 8)
    string.memset(rbuf, 0, DBG_LOC_SLOTS * 8)
    string.memset(ubuf, 0, DBG_LOC_SLOTS)
    _g_stk_sp = (sizet*)(sbuf1)
    _g_stk_ref = sbuf2
    _g_stk_bld = sbuf3
    _g_loc_key = (sizet*)(kbuf)
    _g_loc_ref = rbuf
    _g_loc_used = ubuf
    return 0


def _RefOf(pool: memhub.MemBuddy*, node_id: int) -> str:
    """Allocate a persistent "!N" reference string for a DI node id."""
    buf: str = pool.alloc(24)
    if buf is None:
        return None
    viperlib.snprintf(buf, 24, "!%d", node_id)
    return buf


# ============================================================
# BeginModule - start of one module translation
# ============================================================
def BeginModule(mod: llvmlite.LLVMModule*, file_path: str) -> int:
    """Reset per-module DI state. Call once at the start of each module."""
    global _g_mod, _g_file_path, _g_finalized, _g_file_id, _g_null_id
    global _g_subtype_id, _g_cu_id, _g_cur_sp, _g_cur_entry_ref, _g_cur_builder, _g_depth
    if _g_enabled == 0:
        return 0
    if _EnsureStorage() != 0:
        return 1
    _g_mod = mod
    _g_file_path = file_path
    _g_finalized = 0
    _g_file_id = -1
    _g_null_id = -1
    _g_subtype_id = -1
    _g_cu_id = -1
    _g_cur_sp = -1
    _g_cur_entry_ref = None
    _g_cur_builder = None
    _g_depth = 0
    if _g_loc_used is not None:
        string.memset(_g_loc_used, 0, DBG_LOC_SLOTS)
    return 0


# ============================================================
# _EnsureModuleNodes - DIFile / !{ null } / DISubroutineType / DICompileUnit
# (created lazily at the first define, once per module)
# ============================================================
def _EnsureModuleNodes(pool: memhub.MemBuddy*) -> int:
    global _g_file_id, _g_null_id, _g_subtype_id, _g_cu_id
    if _g_cu_id >= 0:
        return 0
    mod: llvmlite.LLVMModule* = _g_mod
    if mod is None:
        return 1
    # DIFile: **绝对**路径（正斜杠），filename 保留整条路径。
    # 为什么必须绝对：本工程给 CurrentFile 的是**项目相对**路径（如
    # "./Test/LLvmLiteTest/App/main.vp"），LLVM 会把该目录写成 CU 的 DW_AT_comp_dir，
    # 而主文件条目用 `dir_index: 0` 指向 comp_dir；stackmap 写入端只解析
    # `llvm-dwarfdump --debug-line`（其中**没有** comp_dir）⇒ 行记录只能退化成裸文件名
    # "main.vp"：运行期既显示不全，也无法按路径重读源码行（实证：曾整段栈都打印
    # `File "main.vp", line N` 而无源码行）。
    # 绝对化后，无论 LLVM 保留整条 name（写入端 `_tb_is_abs_path` 命中 ⇒ 原样采用）
    # 还是把它拆成 dir+basename（dir_index 指向同一个绝对目录 ⇒ 拼回绝对路径），
    # 结果都是绝对路径 —— 故本改动不需要动写入端。
    src: str = _g_file_path
    if src is None:
        src = ""
    slen: sizet = string.strlen(src)
    cap_n: sizet = slen + 4112
    nbuf: str = pool.alloc(cap_n)
    if nbuf is None:
        VLogger.error("--traceback: DIFile 路径缓冲分配失败，终止编译", "DWARF")
        return 1
    if slen > 0:
        string.strncpy(nbuf, src, slen)
    nbuf[slen] = '\0'
    if slen > 0 and nbuf[0] != '/' and not (slen >= 2 and nbuf[1] == ':'):
        # 相对路径 → 前置当前工作目录。用 w32file 低层调用（同 Repl.vp），
        # 不经 os._mbuddy：编译器进程未必初始化过那个内存管理器。
        CwdBuf: bytes = pool.alloc(4112)
        if CwdBuf is not None:
            CwdBuf[0] = '\0'
            nc: ULONG = w32file.GetCurrentDirectoryA(ULONG(4111), CwdBuf)
            if nc > 0:
                # 防御：缓冲区不足时 GetCurrentDirectoryA 的写入内容未定义，强制补 NUL
                CwdBuf[4110] = '\0'
                viperlib.snprintf(nbuf, cap_n, "%s/%s", CwdBuf, src)
    # 归一所有反斜杠（CWD 段同样可能含 '\'）
    flen: sizet = string.strlen(nbuf)
    ni: sizet = 0
    for ni in range(flen):
        if nbuf[ni] == '\\':
            nbuf[ni] = '/'
    # filename 取整条路径；directory 另存其目录前缀（两段各自 NUL 结尾，互不污染）
    fname: str = nbuf
    directory: str = ""
    slash: str = string.strrchr(nbuf, '/')
    if slash is not None:
        dlen: sizet = string.strlen(nbuf) - string.strlen(slash + 1) - 1
        dirbuf: str = pool.alloc(dlen + 1)
        if dirbuf is not None:
            string.strncpy(dirbuf, nbuf, dlen)
            dirbuf[dlen] = '\0'
            directory = dirbuf
    dsize: sizet = string.strlen(fname) + string.strlen(directory) + 96
    dbuf: str = pool.alloc(dsize)
    if dbuf is None:
        VLogger.error("--traceback: DIFile 节点缓冲分配失败，终止编译", "DWARF")
        return 1
    viperlib.snprintf(dbuf, dsize, "!DIFile(filename: \"%s\", directory: \"%s\")", fname, directory)
    fid: int = llvmlite.module_add_di_node(mod, pool, dbuf)
    if fid < 0:
        VLogger.error("--traceback: DIFile 元数据节点分配失败，终止编译", "DWARF")
        return 1
    _g_file_id = fid
    # types: !{ null } placeholder list, then DISubroutineType referencing it
    nid: int = llvmlite.module_add_di_node(mod, pool, "!{ null }")
    if nid < 0:
        VLogger.error("--traceback: DISubroutineType 类型表分配失败，终止编译", "DWARF")
        return 1
    _g_null_id = nid
    tbuf: str = pool.alloc(64)
    if tbuf is None:
        VLogger.error("--traceback: DISubroutineType 缓冲分配失败，终止编译", "DWARF")
        return 1
    viperlib.snprintf(tbuf, 64, "!DISubroutineType(types: !%d)", _g_null_id)
    tid: int = llvmlite.module_add_di_node(mod, pool, tbuf)
    if tid < 0:
        VLogger.error("--traceback: DISubroutineType 元数据节点分配失败，终止编译", "DWARF")
        return 1
    _g_subtype_id = tid
    cbuf: str = pool.alloc(256)
    if cbuf is None:
        VLogger.error("--traceback: DICompileUnit 缓冲分配失败，终止编译", "DWARF")
        return 1
    viperlib.snprintf(cbuf, 256,
        "distinct !DICompileUnit(language: DW_LANG_C99, file: !%d, producer: \"TPC\", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug)",
        _g_file_id)
    cid: int = llvmlite.module_add_di_node(mod, pool, cbuf)
    if cid < 0:
        VLogger.error("--traceback: DICompileUnit 元数据节点分配失败，终止编译", "DWARF")
        return 1
    _g_cu_id = cid
    return 0


# ============================================================
# _LocRefAt - read a cached DILocation ref string pointer
# ============================================================
def _LocRefAt(slot: int) -> str:
    slotp: str | t.CPtr = (t.CPtr)(uint64(_g_loc_ref) + sizet(slot) * 8)
    r: str = slotp[0]
    return r


# ============================================================
# _EnsureLocation - DILocation ref for (sp, line, col), created + cached once
# ============================================================
def _EnsureLocation(pool: memhub.MemBuddy*, sp: int, line: int, col: int) -> str:
    if line <= 0:
        return _g_cur_entry_ref
    c: int = col
    if c <= 0:
        c = 1
    if _g_loc_used is None or _g_mod is None:
        return None
    h: uint64 = (uint64(sp) * 1000003 + uint64(line) * 31 + uint64(c)) & uint64(DBG_LOC_SLOTS - 1)
    pi: int = 0
    while pi < DBG_LOC_PROBE:
        slot: int = int((h + uint64(pi)) & uint64(DBG_LOC_SLOTS - 1))
        if _g_loc_used[slot] == 0:
            break
        kbase: int = slot * DBG_LOC_KEY_STRIDE
        if int(_g_loc_key[kbase]) == sp and int(_g_loc_key[kbase + 1]) == line \
                and int(_g_loc_key[kbase + 2]) == c:
            return _LocRefAt(slot)
        pi += 1
    lbuf: str = pool.alloc(96)
    if lbuf is None:
        return None
    viperlib.snprintf(lbuf, 96, "!DILocation(line: %d, column: %d, scope: !%d)", line, c, sp)
    lid: int = llvmlite.module_add_di_node(_g_mod, pool, lbuf)
    if lid < 0:
        VLogger.error("--traceback: DILocation 元数据节点分配失败，终止编译", "DWARF")
        return None
    lref: str = _RefOf(pool, lid)
    if lref is None:
        VLogger.error("--traceback: DILocation 引用缓冲分配失败，终止编译", "DWARF")
        return None
    if pi < DBG_LOC_PROBE:
        slot2: int = int((h + uint64(pi)) & uint64(DBG_LOC_SLOTS - 1))
        kbase2: int = slot2 * DBG_LOC_KEY_STRIDE
        _g_loc_key[kbase2] = sizet(sp)
        _g_loc_key[kbase2 + 1] = sizet(line)
        _g_loc_key[kbase2 + 2] = sizet(c)
        rslot: str | t.CPtr = (t.CPtr)(uint64(_g_loc_ref) + sizet(slot2) * 8)
        rslot[0] = lref
        _g_loc_used[slot2] = 1
    return lref


# ============================================================
# PushFunction - create the DISubprogram for one define
# ============================================================
def PushFunction(pool: memhub.MemBuddy*, builder: llvmlite.IRBuilder*,
                 func: llvmlite.Function*, line: int, col: int) -> int:
    """Attach a DISubprogram to one define and set the builder's entry DILocation.

    Returns 1 when pushed (caller must call PopFunction), 0 when DI is disabled
    or not applicable. Every instruction emitted afterwards on `builder` carries
    ", !dbg !N" until the next SetStatement.
    """
    global _g_cur_sp, _g_cur_entry_ref, _g_cur_builder, _g_depth
    if _g_enabled == 0:
        return 0
    if pool is None or builder is None or func is None or _g_mod is None:
        return 0
    if _g_depth >= DBG_STACK_MAX:
        VLogger.error("--traceback: DWARF 函数嵌套深度超上限，终止编译", "DWARF")
        return 0
    if _EnsureModuleNodes(pool) != 0:
        return 0
    name: str = llvmlite.function_get_name(func)
    if name is None:
        VLogger.error("--traceback: 函数名为空，无法发射 DISubprogram，终止编译", "DWARF")
        return 0
    dline: int = line
    if dline <= 0:
        dline = 1
    dcol: int = col
    if dcol <= 0:
        dcol = 1
    ssize: sizet = string.strlen(name) + 320
    spbuf: str = pool.alloc(ssize)
    if spbuf is None:
        VLogger.error("--traceback: DISubprogram 缓冲分配失败，终止编译", "DWARF")
        return 0
    viperlib.snprintf(spbuf, ssize,
        "distinct !DISubprogram(name: \"%s\", scope: !%d, file: !%d, type: !%d, line: %d, scopeLine: %d, unit: !%d, spFlags: DISPFlagDefinition)",
        name, _g_file_id, _g_file_id, _g_subtype_id, dline, dline, _g_cu_id)
    sp_id: int = llvmlite.module_add_di_node(_g_mod, pool, spbuf)
    if sp_id < 0:
        VLogger.error("--traceback: DISubprogram 元数据节点分配失败，终止编译", "DWARF")
        return 0
    sp_ref: str = _RefOf(pool, sp_id)
    if sp_ref is None:
        VLogger.error("--traceback: DISubprogram 引用缓冲分配失败，终止编译", "DWARF")
        return 0
    llvmlite.function_set_dbg(func, sp_ref)
    # save the enclosing function's DI context (LIFO), then switch to this one
    _g_stk_sp[_g_depth] = sizet(_g_cur_sp + 1)
    rslot: str | t.CPtr = (t.CPtr)(uint64(_g_stk_ref) + sizet(_g_depth) * 8)
    rslot[0] = _g_cur_entry_ref
    bslot: str | t.CPtr = (t.CPtr)(uint64(_g_stk_bld) + sizet(_g_depth) * 8)
    bslot[0] = _g_cur_builder
    _g_depth += 1
    _g_cur_sp = sp_id
    _g_cur_builder = builder
    entry_ref: str = _EnsureLocation(pool, sp_id, dline, dcol)
    if entry_ref is None:
        VLogger.error("--traceback: 入口 DILocation 生成失败，终止编译", "DWARF")
        return 0
    _g_cur_entry_ref = entry_ref
    llvmlite.builder_set_debug_location(builder, entry_ref)
    return 1


# ============================================================
# PopFunction - restore the enclosing function's DI context
# ============================================================
def PopFunction(builder: llvmlite.IRBuilder*) -> int:
    global _g_cur_sp, _g_cur_entry_ref, _g_cur_builder, _g_depth
    if _g_enabled == 0:
        return 0
    if _g_depth <= 0:
        return 0
    _g_depth -= 1
    _g_cur_sp = int(_g_stk_sp[_g_depth]) - 1
    rslot: str | t.CPtr = (t.CPtr)(uint64(_g_stk_ref) + sizet(_g_depth) * 8)
    _g_cur_entry_ref = rslot[0]
    bslot: str | t.CPtr = (t.CPtr)(uint64(_g_stk_bld) + sizet(_g_depth) * 8)
    _g_cur_builder = (llvmlite.IRBuilder*)(bslot[0])
    if builder is not None:
        llvmlite.builder_clear_debug_location(builder)
    return 0


# ============================================================
# SetStatement - point the current builder at a statement DILocation
# ============================================================
def SetStatement(pool: memhub.MemBuddy*, builder: llvmlite.IRBuilder*,
                 line: int, col: int) -> int:
    """Update the builder's DILocation to the current statement (line precision).

    No-op outside a DI-attached define (module-level code / synthetic defines),
    so instructions outside a function body never reference a function's scope.
    """
    if _g_enabled == 0 or pool is None or builder is None:
        return 0
    if _g_cur_sp < 0:
        return 0
    ref: str = _EnsureLocation(pool, _g_cur_sp, line, col)
    if ref is None:
        return 0
    llvmlite.builder_set_debug_location(builder, ref)
    return 1


# ============================================================
# _ModuleHasDefine - does the module emit any define (IsDeclared == 0)?
# ============================================================
def _ModuleHasDefine(mod: llvmlite.LLVMModule*) -> int:
    if mod is None:
        return 0
    cur: llvmlite.Function* = mod.FuncHead
    steps: int = 0
    while cur is not None:
        steps += 1
        if steps > 100000:
            break
        if llvmlite.function_is_declared(cur) == 0:
            return 1
        cur = llvmlite.function_get_next(cur)
    return 0


# ============================================================
# Finalize - register !llvm.dbg.cu / !llvm.module.flags before dump_ir
# ============================================================
def Finalize(pool: memhub.MemBuddy*) -> int:
    """Module-level DI tail; idempotent per module, no-op when no SP was emitted."""
    global _g_finalized
    if _g_enabled == 0:
        return 0
    if _g_finalized != 0:
        return 0
    mod: llvmlite.LLVMModule* = _g_mod
    if mod is None or pool is None:
        return 0
    _g_finalized = 1
    # Self-check (mirrors TPC Phase2Translator.py:3444): a module that emits a
    # define but carries no DI at all loses line info for its frames.
    if mod.DIMetaCount == 0 and _ModuleHasDefine(mod) != 0:
        VLogger.warning("--traceback: 本模块发射了 define 但 IR 未携带 DWARF (!llvm.dbg.cu)——该模块帧在 Ctrl+C 栈里将只显示地址", "DWARF")
    if _g_cu_id < 0:
        return 0
    cu_ref: str = _RefOf(pool, _g_cu_id)
    if cu_ref is None:
        VLogger.error("--traceback: DICompileUnit 引用缓冲分配失败，终止编译", "DWARF")
        return 1
    llvmlite.module_add_named_metadata(mod, pool, "llvm.dbg.cu", cu_ref)
    fbuf: str = pool.alloc(64)
    if fbuf is None:
        VLogger.error("--traceback: 模块级调试标志缓冲分配失败，终止编译", "DWARF")
        return 1
    viperlib.snprintf(fbuf, 64, "!{i32 2, !\"Debug Info Version\", i32 3}")
    flags_id: int = llvmlite.module_add_di_node(mod, pool, fbuf)
    if flags_id < 0:
        VLogger.error("--traceback: llvm.module.flags 元数据节点分配失败，终止编译", "DWARF")
        return 1
    fl_ref: str = _RefOf(pool, flags_id)
    if fl_ref is None:
        VLogger.error("--traceback: llvm.module.flags 引用缓冲分配失败，终止编译", "DWARF")
        return 1
    llvmlite.module_add_named_metadata(mod, pool, "llvm.module.flags", fl_ref)
    return 0