from stdint import *
import mlirlite
import memhub
import string
import stdlib
import viperlib
import win32.win32file as win32file
import lib.core.VLogger as VLogger

# DebugInfoGen - emits function-level DWARF metadata (mlirlite text IR), mirroring TPC LLVMCG/DebugInfo.py node-for-node.
# Active only under --traceback (SetEnabled); hooks come from the Gen layer. When disabled every entry point returns before touching DI APIs, keeping IR byte-identical.

# ---- module-scope state (reset per module by BeginModule) ----
_g_enabled: int = 0
_g_finalized: int = 0
_g_mod: mlirlite.LLVMModule* = None
_g_file_path: str = None
_g_abs_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: mlirlite.IRBuilder* = None
_g_depth: int = 0

# Push/Pop save stack: enclosing (sp id + 1, entry DILocation ref, builder) per level; +1 lets 0 mean "no enclosing function".
# The saved builder makes SetStatement reject a builder not pushed by the current DISubprogram (e.g. an async state machine in a bracketed function); such instrs must carry no !dbg or LLVM's verifier fails.
define DBG_STACK_MAX = 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).
define DBG_LOC_SLOTS = 8192
define DBG_LOC_PROBE = 8
define DBG_LOC_KEY_STRIDE = 3
_g_loc_key: sizet* = None
_g_loc_ref: str = None
_g_loc_used: bytes = None


# SetEnabled / IsEnabled - mirror the --traceback switch (flag -> _g_enabled, forwarded to mlirlite's debug call-site recorder).
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
    # Forward the same switch to mlirlite's debug call-site recorder (under --traceback we emit a <sha1>.vpcall side file for WriteStackMap+objdump).
    mlirlite.DebugSetEnabled(flag)
    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: failed to allocate DWARF metadata buffer, aborting compilation", "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: mlirlite.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_abs_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_abs_file_path = None
    # Module-boundary fallback: clear any un-Pop'd debug function context (mid-translation failure) so the previous module's caller is not attributed here.
    mlirlite.DebugResetFunc()
    _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 - lazily create DIFile / !{ null } / DISubroutineType / DICompileUnit 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, _g_abs_file_path
    if _g_cu_id >= 0: return 0
    mod: mlirlite.LLVMModule* = _g_mod
    if mod is None: return 1
    # DIFile: absolute path (forward slashes), filename keeps the whole path.
    # Must be absolute: relative paths make LLVM's DWARF line records degrade to a bare "main.vp" that runtime cannot reload; absolutizing works whether LLVM keeps the full name or splits dir+basename.
    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: failed to allocate DIFile path buffer, aborting compilation", "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] == ':'):
        # Relative path -> prefix CWD via a win32file low-level call (as in Repl.vp), not os._mbuddy which the compiler process may not have initialized.
        CwdBuf: bytes = pool.alloc(4112)
        if CwdBuf is not None:
            CwdBuf[0] = '\0'
            nc: ULONG = win32file.GetCurrentDirectoryA(ULONG(4111), CwdBuf)
            if nc > 0:
                # Defensive: when the buffer is too small GetCurrentDirectoryA writes undefined content; force a NUL.
                CwdBuf[4110] = '\0'
                viperlib.snprintf(nbuf, cap_n, "%s/%s", CwdBuf, src)
    # Normalize all backslashes (the CWD segment may contain '\' too).
    flen: sizet = string.strlen(nbuf)
    ni: sizet = 0
    for ni in range(int(flen)):
        if nbuf[ni] == '\\': nbuf[ni] = '/'
    # Save the absolute path for mlirlite's debug call-site recorder (.vpcall file paths must match DIFile: absolute + forward slash).
    _g_abs_file_path = nbuf
    # filename = whole path; directory = its prefix (both NUL-terminated, non-overlapping).
    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: failed to allocate DIFile node buffer, aborting compilation", "DWARF")
        return 1
    viperlib.snprintf(dbuf, dsize, "!DIFile(filename: \"%s\", directory: \"%s\")", fname, directory)
    fid: int = mlirlite.module_add_di_node(mod, pool, dbuf)
    if fid < 0:
        VLogger.error("--traceback: failed to allocate DIFile metadata node, aborting compilation", "DWARF")
        return 1
    _g_file_id = fid
    # types: !{ null } placeholder list, then DISubroutineType referencing it
    nid: int = mlirlite.module_add_di_node(mod, pool, "!{ null }")
    if nid < 0:
        VLogger.error("--traceback: failed to allocate DISubroutineType type table, aborting compilation", "DWARF")
        return 1
    _g_null_id = nid
    tbuf: str = pool.alloc(64)
    if tbuf is None:
        VLogger.error("--traceback: failed to allocate DISubroutineType buffer, aborting compilation", "DWARF")
        return 1
    viperlib.snprintf(tbuf, 64, "!DISubroutineType(types: !%d)", _g_null_id)
    tid: int = mlirlite.module_add_di_node(mod, pool, tbuf)
    if tid < 0:
        VLogger.error("--traceback: failed to allocate DISubroutineType metadata node, aborting compilation", "DWARF")
        return 1
    _g_subtype_id = tid
    cbuf: str = pool.alloc(256)
    if cbuf is None:
        VLogger.error("--traceback: failed to allocate DICompileUnit buffer, aborting compilation", "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 = mlirlite.module_add_di_node(mod, pool, cbuf)
    if cid < 0:
        VLogger.error("--traceback: failed to allocate DICompileUnit metadata node, aborting compilation", "DWARF")
        return 1
    _g_cu_id = cid
    return 0


# _LocRefAt - read a cached DILocation ref string pointer.
def _LocRefAt(slot: int) -> str:
    slotp: str* = (str*)(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 = mlirlite.module_add_di_node(_g_mod, pool, lbuf)
    if lid < 0:
        VLogger.error("--traceback: failed to allocate DILocation metadata node, aborting compilation", "DWARF")
        return None
    lref: str = _RefOf(pool, lid)
    if lref is None:
        VLogger.error("--traceback: failed to allocate DILocation reference buffer, aborting compilation", "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* = (str*)(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: mlirlite.IRBuilder*, func: mlirlite.Function*, line: int, col: int) -> int:
    """Attach a DISubprogram to one define and set the builder's entry DILocation."""
    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 function nesting depth exceeds the limit, aborting compilation", "DWARF")
        return 0
    if _EnsureModuleNodes(pool) != 0: return 0
    name: str = mlirlite.function_get_name(func)
    if name is None:
        VLogger.error("--traceback: function name is empty, cannot emit DISubprogram, aborting compilation", "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: failed to allocate DISubprogram buffer, aborting compilation", "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 = mlirlite.module_add_di_node(_g_mod, pool, spbuf)
    if sp_id < 0:
        VLogger.error("--traceback: failed to allocate DISubprogram metadata node, aborting compilation", "DWARF")
        return 0
    sp_ref: str = _RefOf(pool, sp_id)
    if sp_ref is None:
        VLogger.error("--traceback: failed to allocate DISubprogram reference buffer, aborting compilation", "DWARF")
        return 0
    mlirlite.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* = (str*)(uint64(_g_stk_ref) + sizet(_g_depth) * 8)
    rslot[0] = _g_cur_entry_ref
    bslot: str* = (str*)(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: failed to generate entry DILocation, aborting compilation", "DWARF")
        return 0
    _g_cur_entry_ref = entry_ref
    mlirlite.builder_set_debug_location(builder, entry_ref)
    # Enter this define: tell mlirlite's call-site recorder the current symbol/file/def-line and write an F record (for the .vpcall file + WriteStackMap entry fallback line).
    mlirlite.DebugNoteFunc(name, _g_abs_file_path, dline)
    return 1


# PopFunction - restore the enclosing function's DI context.
def PopFunction(builder: mlirlite.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* = (str*)(uint64(_g_stk_ref) + sizet(_g_depth) * 8)
    _g_cur_entry_ref = rslot[0]
    bslot: str* = (str*)(uint64(_g_stk_bld) + sizet(_g_depth) * 8)
    _g_cur_builder = (mlirlite.IRBuilder*)(bslot[0])
    if builder is not None: mlirlite.builder_clear_debug_location(builder)
    # Leave this define: restore mlirlite's outer function context (for nested defines).
    mlirlite.DebugClearFunc()
    return 0


# SetStatement - point the current builder at a statement DILocation.
def SetStatement(pool: memhub.MemBuddy*, builder: mlirlite.IRBuilder*, line: int, col: int) -> int:
    """Update the builder's DILocation to the current statement (line precision)."""
    if _g_enabled == 0 or pool is None or builder is None: return 0
    # Record the current statement line: C records from build_call in this statement use it as the call-site line.
    mlirlite.DebugNoteLine(line)
    if _g_cur_sp < 0: return 0
    ref: str = _EnsureLocation(pool, _g_cur_sp, line, col)
    if ref is None: return 0
    mlirlite.builder_set_debug_location(builder, ref)
    return 1


# _ModuleHasDefine - does the module emit any define (IsDeclared == 0)?
def _ModuleHasDefine(mod: mlirlite.LLVMModule*) -> int:
    if mod is None: return 0
    cur: mlirlite.Function* = mod.FuncHead
    steps: int = 0
    while cur is not None:
        steps += 1
        if steps > 100000: break
        if mlirlite.function_is_declared(cur) == 0: return 1
        cur = mlirlite.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: mlirlite.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: this module emits a define but its IR carries no DWARF (!llvm.dbg.cu); its frames will show only addresses in the Ctrl+C stack", "DWARF")
    if _g_cu_id < 0: return 0
    cu_ref: str = _RefOf(pool, _g_cu_id)
    if cu_ref is None:
        VLogger.error("--traceback: failed to allocate DICompileUnit reference buffer, aborting compilation", "DWARF")
        return 1
    mlirlite.module_add_named_metadata(mod, pool, "llvm.dbg.cu", cu_ref)
    fbuf: str = pool.alloc(64)
    if fbuf is None:
        VLogger.error("--traceback: failed to allocate module-level debug flags buffer, aborting compilation", "DWARF")
        return 1
    viperlib.snprintf(fbuf, 64, "!{i32 2, !\"Debug Info Version\", i32 3}")
    flags_id: int = mlirlite.module_add_di_node(mod, pool, fbuf)
    if flags_id < 0:
        VLogger.error("--traceback: failed to allocate llvm.module.flags metadata node, aborting compilation", "DWARF")
        return 1
    fl_ref: str = _RefOf(pool, flags_id)
    if fl_ref is None:
        VLogger.error("--traceback: failed to allocate llvm.module.flags reference buffer, aborting compilation", "DWARF")
        return 1
    mlirlite.module_add_named_metadata(mod, pool, "llvm.module.flags", fl_ref)
    return 0