import t, c
from stdint import *
import stdio
import string
import stdlib
import memhub
import w32.fileio as fileio
import w32.win32file
import w32.win32base
import subprocess
import viperlib
import ast
import lib.core.VLogger as VLogger
import lib.core.Handles.HandlesTranslator as HandlesTranslator
import lib.core.Handles.HandlesStruct as HandlesStruct
import lib.core.Handles.HandlesType as HandlesType
import lib.core.Handles.HandlesImports as HandlesImports
import lib.Projectrans.Config as Config

# ============================================================
# BuildPipeline - 编译管线
#
# 负责将 LLVM IR 编译为可执行文件：
#   1. 写 .ll 文件到 temp 目录
#   2. 调用 llc 编译 .ll → .obj
#   3. 调用 clang++ 链接 .obj → .exe
# ============================================================

# 全局 mbuddy 指针
_mbuddy: memhub.MemManager | t.CPtr

# 源代码缓冲区大小（1MB）
SRC_BUF_SIZE: t.CDefine = 1048576


# ============================================================
# TranslateFileGetTrans - 翻译单个文件，返回 Translator 对象
#
# 读取文件 → AST 解析 → LLVM IR 翻译，返回 Translator 对象供
# 调用者 dump_ir。命名空间隔离：includes 文件宽松模式（全部可见），
# 用户文件严格模式（仅本地+import）。
#
# Args:
#   mb: 内存池
#   file_path: 源文件路径
#   sha1_val: 文件内容的 SHA1 前16字符（用于函数名混淆）
#
# Returns:
#   Translator 对象（None 失败）
# ============================================================
def TranslateFileGetTrans(mb: memhub.MemBuddy | t.CPtr, file_path: str,
                          sha1_val: str,
                          current_package: str = None,
                          declare_only: int = 0) -> HandlesTranslator.Translator | t.CPtr:
    """翻译文件，返回 Translator 对象（调用者负责 dump_ir），None 失败

    current_package: 当前文件所属包名（用于解析相对导入），None 表示顶级模块
    declare_only: 0=全量翻译（默认），1=仅注册 struct/enum/union（不翻译方法体）
    """
    if file_path is None:
        return None
    f: fileio.File | t.CPtr = fileio.File(file_path, fileio.MODE.R)
    if f.closed:
        return None
    src_buf: bytes = stdlib.malloc(SRC_BUF_SIZE)
    if src_buf is None:
        f.close()
        return None
    bytes_read: LONG = f.read_all(src_buf, SRC_BUF_SIZE)
    f.close()
    if bytes_read <= 0:
        stdlib.free(src_buf)
        return None
    if bytes_read < SRC_BUF_SIZE:
        src_buf[bytes_read] = 0
    else:
        src_buf[SRC_BUF_SIZE - 1] = 0
    lx: ast.Lexer | t.CPtr = ast.new_lexer(mb)
    ast._lexer_init(lx, src_buf, mb)
    tokens: ast.Token | t.CPtr = ast.tokenize(lx)
    tree: ast.AST | t.CPtr = ast.parse_tokens(mb, tokens)
    # parse 完成后 AST 树已独立，释放源代码缓冲区（避免 1MB/文件泄漏）
    stdlib.free(src_buf)
    if tree is None:
        return None
    tr: HandlesTranslator.Translator | t.CPtr = mb.alloc(HandlesTranslator.Translator.__sizeof__())
    tr.__before_init__()
    tr.__init__()
    tr.ModuleSha1 = sha1_val
    tr.CurrentPackage = current_package
    tr._declare_only = declare_only
    # 设置当前文件名（供报错使用）
    HandlesType.set_current_file(file_path)
    # 模块切换：清空 CDefine 常量表，确保每个模块的 CDefine 常量正确隔离
    HandlesType.clear_cdefine_constants()
    # 命名空间隔离：includes 文件宽松模式（全部可见），用户文件严格模式（仅本地+import）
    strict_mode: int = 1
    if string.strstr(file_path, "includes") is not None:
        strict_mode = 0
    HandlesStruct.reset_visible_structs(mb, strict_mode)
    ret: int = tr.translate(tree)
    if ret != 0:
        return None
    return tr


class BuildResult:
    Success: t.CInt
    OutputPath: str
    ErrorMsg: str

    def __new__(self) -> t.CPtr:
        return t.CPtr(_mbuddy.alloc(BuildResult.__sizeof__()))

    def __init__(self):
        self.Success = 0
        self.OutputPath = None
        self.ErrorMsg = None


def ensure_dir(path: str) -> int:
    """递归创建目录（类似 mkdir -p），目录已存在视为成功

    Args:
        path: 目录路径（支持 / 或 \\ 分隔符）

    Returns:
        0 成功（包括目录已存在），非 0 失败
    """
    if path is None:
        return 1

    path_len: t.CSizeT = string.strlen(path)
    if path_len == 0:
        return 0

    # 复制路径到可写缓冲区（逐级截断用）
    buf: bytes = _mbuddy.alloc(path_len + 1)
    if buf is None:
        return 1
    string.strcpy(buf, path)

    # 遇到分隔符时临时截断，创建每一层目录
    # CreateDirectoryA 在目录已存在时返回 0（失败），忽略即可
    i: int = 0
    plen: int = path_len
    while i < plen:
        ch: int = buf[i]
        if ch == 47 or ch == 92:  # '/' = 47, '\\' = 92
            saved: int = ch
            buf[i] = '\0'
            w32.win32file.CreateDirectoryA(buf, None)
            buf[i] = saved
        i += 1

    # 创建最终目录
    w32.win32file.CreateDirectoryA(buf, None)
    return 0


# ============================================================
# build_sliced_path - 构建切片路径并确保目录存在
#
# 根据 Config.Sha1SliceLevel 将文件分散到 SHA1 前缀子目录:
#   level=0 → {base_dir}/{sha1}.{ext}
#   level=1 → {base_dir}/b7/{sha1}.{ext}
#   level=2 → {base_dir}/b7/90/{sha1}.{ext}
#
# 自动调用 ensure_dir 创建子目录。
#
# Args:
#   base_dir: 基础目录（如 temp_dir）
#   sha1: SHA1 字符串
#   ext: 文件扩展名（如 "stub.ll"）
#
# Returns:
#   完整文件路径（mbuddy 分配），None 失败
# ============================================================
def build_sliced_path(base_dir: str, sha1: str, ext: str) -> str:
    """构建切片路径并确保目录存在"""
    if base_dir is None or sha1 is None or ext is None:
        return None

    subdir: str = Config.slice_subdir(sha1, Config.Sha1SliceLevel)
    base_len: t.CSizeT = string.strlen(base_dir)
    sha1_len: t.CSizeT = string.strlen(sha1)
    ext_len: t.CSizeT = string.strlen(ext)

    if subdir is not None:
        sub_len: t.CSizeT = string.strlen(subdir)
        # ensure_dir({base_dir}/{subdir})
        dir_path: str = _mbuddy.alloc(base_len + sub_len + 2)
        if dir_path is not None:
            viperlib.snprintf(dir_path, base_len + sub_len + 2, "%s/%s", base_dir, subdir)
            ensure_dir(dir_path)
        # 文件路径: {base_dir}/{subdir}/{sha1}.{ext}
        path_len: t.CSizeT = base_len + sub_len + sha1_len + ext_len + 4
        path: str = _mbuddy.alloc(path_len)
        if path is None:
            return None
        viperlib.snprintf(path, path_len, "%s/%s/%s.%s", base_dir, subdir, sha1, ext)
        return path
    else:
        # 无切片: {base_dir}/{sha1}.{ext}
        path_len = base_len + sha1_len + ext_len + 3
        path = _mbuddy.alloc(path_len)
        if path is None:
            return None
        viperlib.snprintf(path, path_len, "%s/%s.%s", base_dir, sha1, ext)
        return path


def write_ir_to_file(ir_buf: bytes, ir_len: t.CSizeT, output_dir: str, module_name: str) -> int:
    """将 IR 缓冲区写入 .ll 文件

    Args:
        ir_buf: IR 文本缓冲区
        ir_len: IR 文本长度
        output_dir: 输出目录（temp 或 output）
        module_name: 模块名（SHA1）

    Returns:
        0 成功，非 0 失败
    """
    if ir_buf is None or output_dir is None or module_name is None:
        return 1

    # 构造切片路径: output_dir/{sha1前缀}/{module_name}.ll
    path: str = build_sliced_path(output_dir, module_name, "ll")
    if path is None:
        return 1

    # 打开文件写入
    f: fileio.File | t.CPtr = fileio.File(path, fileio.MODE.W)
    if f.closed:
        return 1

    written: t.CInt64T = f.write(ir_buf, ir_len)
    f.close()
    if written < 0:
        return 1
    return 0


def compile_ll_to_obj(ir_path: str, output_dir: str, module_name: str, cc_cmd: str, cc_flags: str) -> int:
    """调用 llc 将 .ll 编译为 .obj

    Args:
        ir_path: .ll 文件路径
        output_dir: 输出目录
        module_name: 模块名（SHA1，用于生成 .obj 文件名）
        cc_cmd: 编译器命令（如 "llc"）
        cc_flags: 编译器参数（如 "-filetype=obj -relocation-model=pic"）

    Returns:
        0 成功，非 0 失败
    """
    if ir_path is None or cc_cmd is None or output_dir is None or module_name is None:
        return 1

    # 构造切片 .obj 路径: output_dir/{sha1前缀}/{module_name}.obj
    obj_path: str = build_sliced_path(output_dir, module_name, "obj")
    if obj_path is None:
        return 1

    # 构造命令: llc -filetype=obj -o {obj_path} ir_path
    cmd_len: t.CSizeT = string.strlen(cc_cmd) + string.strlen(cc_flags) + string.strlen(obj_path) + string.strlen(ir_path) + 64
    cmd: bytes = _mbuddy.alloc(cmd_len)
    if cmd is None:
        return 1
    viperlib.snprintf(cmd, cmd_len, "%s %s -o %s %s", cc_cmd, cc_flags, obj_path, ir_path)

    result: subprocess.CompletedProcess | t.CPtr = subprocess.run(cmd, True, True)
    if result is None:
        fb: t.CChar | t.CPtr = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "subprocess.run 返回 None: %s", module_name)
            VLogger.error(fb, "LLC")
        return 1
    if result.returncode != 0:
        # 先直接输出 llc 的具体错误信息（VLogger.error 会 sys.exit，必须先输出）
        # 注意: subprocess 在 Windows 下将 stderr 合并到 stdout（si.hStdError = stdout_write）
        # 因此 result.stderr 总是 None，错误信息在 result.stdout 中
        # stdout 通常已含换行符，不再额外加 \n
        if result.stdout is not None:
            stdio.printf("%s", result.stdout)
            stdio.fflush(0)
        if result.stderr is not None:
            stdio.printf("%s", result.stderr)
            stdio.fflush(0)
        # 最后输出编译失败摘要（VLogger.error 会 sys.exit）
        fb = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "编译失败 (module=%s, cmd=%s)", module_name, cmd)
            VLogger.error(fb, "LLC")
        return 1
    return 0


def collect_obj_files(includes_binary_dir: str, out_buf: bytes, out_size: t.CSizeT) -> t.CSizeT:
    """扫描 includes_binary_dir 目录，收集所有 .obj 文件路径到 out_buf

    Args:
        includes_binary_dir: includes.binary 目录路径
        out_buf: 输出缓冲区（用于存放空格分隔的 .obj 文件路径）
        out_size: 输出缓冲区大小

    Returns:
        写入的字节数（不含 null 终止符），0 表示无文件或错误
    """
    if includes_binary_dir is None or out_buf is None or out_size == 0:
        return 0

    out_buf[0] = '\0'
    out_pos: t.CSizeT = 0

    # 构造搜索模式: dir/*.obj
    dir_len: t.CSizeT = string.strlen(includes_binary_dir)
    pattern: bytes = _mbuddy.alloc(dir_len + 8)
    if pattern is None:
        return 0
    viperlib.snprintf(pattern, dir_len + 8, "%s/*.obj", includes_binary_dir)

    # FindFirstFileA
    find_data: w32.win32file.WIN32_FIND_DATAA | t.CPtr = _mbuddy.alloc(w32.win32file.WIN32_FIND_DATAA.__sizeof__())
    if find_data is None:
        return 0
    string.memset(find_data, 0, w32.win32file.WIN32_FIND_DATAA.__sizeof__())

    handle: w32.win32base.HANDLE = w32.win32file.FindFirstFileA(pattern, find_data)
    if handle == w32.win32base.INVALID_HANDLE_VALUE:
        return 0

    # 需要 win32base 的 INVALID_HANDLE_VALUE
    while True:
        fname: str = find_data.cFileName
        if fname is not None:
            fname_len: t.CSizeT = string.strlen(fname)
            # 构造完整路径并追加到 out_buf
            # 路径格式: "dir/fname "
            need: t.CSizeT = dir_len + 1 + fname_len + 2
            if out_pos + need < out_size:
                if out_pos > 0:
                    out_buf[out_pos] = ' '
                    out_pos += 1
                viperlib.snprintf(out_buf + out_pos, need, "%s/%s", includes_binary_dir, fname)
                out_pos += dir_len + 1 + fname_len
            else:
                # 缓冲区不足，停止
                break

        if w32.win32file.FindNextFileA(handle, find_data) == 0:
            break

    w32.win32file.FindClose(handle)
    out_buf[out_pos] = '\0'
    return out_pos


def link_obj_to_exe(output_dir: str, module_name: str, linker_cmd: str, linker_flags: str, linker_output: str,
                    includes_binary_dir: str) -> int:
    """调用 clang++ 链接 .obj 为 .exe

    Args:
        output_dir: 输出目录（包含 .obj 文件）
        module_name: 模块名（SHA1）
        linker_cmd: 链接器命令（如 "clang++"）
        linker_flags: 链接器参数
        linker_output: 输出文件名（如 "test.exe"）
        includes_binary_dir: includes.binary 目录路径（链接时附加预编译 .obj）

    Returns:
        0 成功，非 0 失败
    """
    if output_dir is None or linker_cmd is None or module_name is None:
        return 1

    # 构造切片 .obj 路径: output_dir/{sha1前缀}/{module_name}.obj
    obj_path: str = build_sliced_path(output_dir, module_name, "obj")
    if obj_path is None:
        return 1

    # 收集 includes.binary 的 .obj 文件路径
    EXTRA_BUF_SIZE: t.CSizeT = 32768
    extra_objs: bytes = _mbuddy.alloc(EXTRA_BUF_SIZE)
    if extra_objs is None:
        return 1
    extra_len: t.CSizeT = 0
    if includes_binary_dir is not None:
        extra_len = collect_obj_files(includes_binary_dir, extra_objs, EXTRA_BUF_SIZE)
        if extra_len > 0:
            fb: t.CChar | t.CPtr = VLogger.fmt_buf()
            if fb is not None:
                viperlib.snprintf(fb, 1024, "附加 %d 字节的 includes.binary .obj 文件", extra_len)
                VLogger.info(fb, "link")
        else:
            fb = VLogger.fmt_buf()
            if fb is not None:
                viperlib.snprintf(fb, 1024, "includes.binary 无 .obj 文件: %s", includes_binary_dir)
                VLogger.warning(fb, "link")

    # 构造命令: clang++ {obj_path} extra_objs -o {output_dir}/{linker_output} linker_flags
    # 注意: .obj 文件必须在 -l 库标志之前，否则链接器无法解析符号依赖
    cmd_len: t.CSizeT = string.strlen(linker_cmd) + string.strlen(obj_path) + string.strlen(linker_flags) + string.strlen(output_dir) + string.strlen(linker_output) + extra_len + 128
    cmd: bytes = _mbuddy.alloc(cmd_len)
    if cmd is None:
        return 1
    if extra_len > 0:
        viperlib.snprintf(cmd, cmd_len, "%s %s %s -o %s/%s %s",
                          linker_cmd, obj_path, extra_objs,
                          output_dir, linker_output, linker_flags)
    else:
        viperlib.snprintf(cmd, cmd_len, "%s %s -o %s/%s %s",
                          linker_cmd, obj_path, output_dir, linker_output,
                          linker_flags)

    result: subprocess.CompletedProcess | t.CPtr = subprocess.run(cmd, True, True)
    if result is None:
        VLogger.error("subprocess.run 返回 None", "link")
        return 1
    if result.returncode != 0:
        fb: t.CChar | t.CPtr = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "链接失败，返回码: %d", result.returncode)
            VLogger.error(fb, "link")
        fb = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "命令: %s", cmd)
            VLogger.error(fb, "link")
        # 显示链接器错误输出
        # 注意: subprocess 在 Windows 下将 stderr 合并到 stdout，result.stderr 总是 None
        if result.stdout is not None:
            VLogger.error(result.stdout, "link")
        else:
            VLogger.error("无输出捕获", "link")
        if result.stderr is not None:
            VLogger.error(result.stderr, "link")
        return 1
    return 0


def compile_module_to_obj(ir_buf: bytes, ir_len: t.CSizeT,
                          temp_dir: str, output_dir: str, module_name: str,
                          cc_cmd: str, cc_flags: str) -> int:
    """编译 IR 到 .obj（不链接）

    Args:
        ir_buf: LLVM IR 文本缓冲区
        ir_len: IR 文本长度
        temp_dir: 临时目录
        output_dir: 输出目录
        module_name: 模块名（SHA1）
        cc_cmd: 编译器命令
        cc_flags: 编译器参数

    Returns:
        0 成功，非 0 失败
    """
    if ir_buf is None or temp_dir is None or output_dir is None or module_name is None:
        return 1

    # Step 1: 写 .ll 文件
    ret: int = write_ir_to_file(ir_buf, ir_len, temp_dir, module_name)
    if ret != 0:
        fb: t.CChar | t.CPtr = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "写 .ll 失败: %s", module_name)
            VLogger.error(fb, "compile")
        return 1

    # Step 2: 构造切片 .ll 路径并编译 → .obj
    ir_path: str = build_sliced_path(temp_dir, module_name, "ll")
    if ir_path is None:
        return 1

    ret = compile_ll_to_obj(ir_path, output_dir, module_name, cc_cmd, cc_flags)
    if ret != 0:
        fb = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "llc 编译失败: %s", module_name)
            VLogger.error(fb, "compile")
        return 1
    return 0


def link_objs_to_exe(obj_paths: str, obj_paths_len: t.CSizeT,
                     linker_cmd: str, linker_flags: str, linker_output: str,
                     includes_binary_dir: str) -> int:
    """链接多个 .obj 文件为 .exe

    Args:
        obj_paths: 空格分隔的 .obj 文件完整路径字符串
        obj_paths_len: obj_paths 长度
        linker_cmd: 链接器命令（如 "clang++"）
        linker_flags: 链接器参数
        linker_output: 输出文件完整路径
        includes_binary_dir: includes.binary 目录路径（链接时附加预编译 .obj）

    Returns:
        0 成功，非 0 失败
    """
    if obj_paths is None or obj_paths_len == 0 or linker_cmd is None:
        return 1

    # 收集 includes.binary 的 .obj 文件路径
    EXTRA_BUF_SIZE: t.CSizeT = 32768
    extra_objs: bytes = _mbuddy.alloc(EXTRA_BUF_SIZE)
    if extra_objs is None:
        return 1
    extra_len: t.CSizeT = 0
    if includes_binary_dir is not None:
        extra_len = collect_obj_files(includes_binary_dir, extra_objs, EXTRA_BUF_SIZE)

    # 构造命令: clang++ obj_paths extra_objs -o linker_output linker_flags
    cmd_len: t.CSizeT = string.strlen(linker_cmd) + obj_paths_len + string.strlen(linker_flags) + string.strlen(linker_output) + extra_len + 128
    cmd: bytes = _mbuddy.alloc(cmd_len)
    if cmd is None:
        return 1
    if extra_len > 0:
        viperlib.snprintf(cmd, cmd_len, "%s %s %s -o %s %s",
                          linker_cmd, obj_paths, extra_objs,
                          linker_output, linker_flags)
    else:
        viperlib.snprintf(cmd, cmd_len, "%s %s -o %s %s",
                          linker_cmd, obj_paths, linker_output, linker_flags)

    result: subprocess.CompletedProcess | t.CPtr = subprocess.run(cmd, True, True)
    if result is None:
        VLogger.error("subprocess.run 返回 None", "link")
        return 1
    if result.returncode != 0:
        # 先直接输出 linker 的具体错误信息（VLogger.error 会 sys.exit，必须先输出）
        # stdout/stderr 通常已含换行符，不再额外加 \n
        if result.stdout is not None:
            stdio.printf("%s", result.stdout)
            stdio.fflush(0)
        if result.stderr is not None:
            stdio.printf("%s", result.stderr)
            stdio.fflush(0)
        # 最后输出链接失败摘要（VLogger.error 会 sys.exit）
        fb: t.CChar | t.CPtr = VLogger.fmt_buf()
        if fb is not None:
            viperlib.snprintf(fb, 1024, "链接失败，返回码: %d, 命令: %s", result.returncode, cmd)
            VLogger.error(fb, "link")
        return 1
    return 0


def run_pipeline(ir_buf: bytes, ir_len: t.CSizeT,
                 temp_dir: str, output_dir: str, module_name: str,
                 cc_cmd: str, cc_flags: str,
                 linker_cmd: str, linker_flags: str, linker_output: str,
                 includes_binary_dir: str) -> BuildResult | t.CPtr:
    """执行完整编译管线

    Args:
        ir_buf: LLVM IR 文本缓冲区
        ir_len: IR 文本长度
        temp_dir: 临时目录
        output_dir: 输出目录
        module_name: 模块名
        cc_cmd: 编译器命令
        cc_flags: 编译器参数
        linker_cmd: 链接器命令
        linker_flags: 链接器参数
        linker_output: 输出文件名
        includes_binary_dir: includes.binary 目录路径（链接时附加预编译 .obj）

    Returns:
        BuildResult 对象
    """
    result: BuildResult | t.CPtr = BuildResult()
    if result is None:
        return None

    # Step 0: 确保 temp/output 目录存在（自动创建，避免写文件失败）
    ensure_dir(temp_dir)
    ensure_dir(output_dir)

    # Step 1: 写 .ll 文件
    ret: int = write_ir_to_file(ir_buf, ir_len, temp_dir, module_name)
    if ret != 0:
        result.Success = 0
        result.ErrorMsg = "写 .ll 文件失败"
        return result

    # Step 2: llc 编译 .ll → .obj
    # 构造切片 .ll 文件路径
    ir_path: str = build_sliced_path(temp_dir, module_name, "ll")
    if ir_path is None:
        result.Success = 0
        result.ErrorMsg = "内存分配失败"
        return result

    ret = compile_ll_to_obj(ir_path, output_dir, module_name, cc_cmd, cc_flags)
    if ret != 0:
        result.Success = 0
        result.ErrorMsg = "llc 编译失败"
        return result

    # Step 3: clang++ 链接 .obj → .exe（附加 includes.binary 预编译 .obj）
    ret = link_obj_to_exe(output_dir, module_name, linker_cmd, linker_flags, linker_output,
                          includes_binary_dir)
    if ret != 0:
        result.Success = 0
        result.ErrorMsg = "链接失败"
        return result

    result.Success = 1
    return result