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.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

# ============================================================
# 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) -> HandlesTranslator.Translator | t.CPtr:
    """翻译文件，返回 Translator 对象（调用者负责 dump_ir），None 失败

    current_package: 当前文件所属包名（用于解析相对导入），None 表示顶级模块
    """
    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:
        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)
    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
    # 设置当前文件名（供报错使用）
    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（失败），忽略即可
    for i in range(path_len):
        ch: int = c.Deref(buf + i)
        if ch == ord('/') or ch == ord('\\'):
            saved: int = ch
            buf[i] = '\0'
            w32.win32file.CreateDirectoryA(buf, None)
            buf[i] = saved

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


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: 模块名（如 "main"）

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

    # 构造文件路径: output_dir/module_name.ll
    dir_len: t.CSizeT = string.strlen(output_dir)
    name_len: t.CSizeT = string.strlen(module_name)
    path_len: t.CSizeT = dir_len + 1 + name_len + 4  # dir/module.ll\0
    path: bytes = _mbuddy.alloc(path_len)
    if path is None:
        return 1
    viperlib.snprintf(path, path_len, "%s/%s.ll", output_dir, module_name)

    # 打开文件写入
    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: 模块名（用于生成 .obj 文件名）
        cc_cmd: 编译器命令（如 "llc"）
        cc_flags: 编译器参数（如 "-filetype=obj -relocation-model=pic"）

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

    # 构造命令: llc -filetype=obj -o output_dir/module_name.obj ir_path
    cmd_len: t.CSizeT = string.strlen(cc_cmd) + string.strlen(cc_flags) + string.strlen(output_dir) + string.strlen(module_name) + 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.obj %s", cc_cmd, cc_flags, output_dir, module_name, ir_path)

    result: subprocess.CompletedProcess | t.CPtr = subprocess.run(cmd, True, True)
    if result is None:
        stdio.printf("[FATAL][LLC] subprocess.run 返回 None: %s\n", module_name)
        return 1
    if result.returncode != 0:
        stdio.printf("[FATAL][LLC] 编译失败 (module=%s, cmd=%s)\n", module_name, cmd)
        if result.stdout is not None:
            stdio.printf("[LLC] 输出:\n%s\n", result.stdout)
        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: 模块名
        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:
        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:
            stdio.printf("[link] 附加 %d 字节的 includes.binary .obj 文件\n", extra_len)
        else:
            stdio.printf("[link] 警告: includes.binary 无 .obj 文件: %s\n", includes_binary_dir)

    # 构造命令: clang++ main.obj extra_objs -o output linker_flags
    # 注意: .obj 文件必须在 -l 库标志之前，否则链接器无法解析符号依赖
    cmd_len: t.CSizeT = string.strlen(linker_cmd) + string.strlen(output_dir) + string.strlen(module_name) + 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.obj %s -o %s/%s %s",
                          linker_cmd, output_dir, module_name, extra_objs,
                          output_dir, linker_output, linker_flags)
    else:
        viperlib.snprintf(cmd, cmd_len, "%s %s/%s.obj -o %s/%s %s",
                          linker_cmd, output_dir, module_name, output_dir, linker_output,
                          linker_flags)

    result: subprocess.CompletedProcess | t.CPtr = subprocess.run(cmd, True, True)
    if result is None:
        stdio.printf("[link] subprocess.run 返回 None\n")
        return 1
    if result.returncode != 0:
        stdio.printf("[link] 链接失败，返回码: %d\n", result.returncode)
        stdio.printf("[link] 命令: %s\n", cmd)
        # 显示链接器错误输出（subprocess 将 stderr 合并到 stdout）
        if result.stdout is not None:
            stdio.printf("[link] 链接器输出:\n%s\n", result.stdout)
        else:
            stdio.printf("[link] 无输出捕获\n")
        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:
        return 1

    # Step 1: 写 .ll 文件
    ret: int = write_ir_to_file(ir_buf, ir_len, temp_dir, module_name)
    if ret != 0:
        stdio.printf("[compile] 写 .ll 失败: %s\n", module_name)
        return 1

    # Step 2: 构造 .ll 路径并编译 → .obj
    name_len: t.CSizeT = string.strlen(module_name)
    dir_len: t.CSizeT = string.strlen(temp_dir)
    ir_path: bytes = _mbuddy.alloc(dir_len + name_len + 5)
    if ir_path is None:
        return 1
    viperlib.snprintf(ir_path, dir_len + name_len + 5, "%s/%s.ll", temp_dir, module_name)

    ret = compile_ll_to_obj(ir_path, output_dir, module_name, cc_cmd, cc_flags)
    if ret != 0:
        stdio.printf("[compile] llc 编译失败: %s\n", module_name)
        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:
        stdio.printf("[link] subprocess.run 返回 None\n")
        return 1
    if result.returncode != 0:
        stdio.printf("[link] 链接失败，返回码: %d\n", result.returncode)
        stdio.printf("[link] 命令: %s\n", cmd)
        if result.stdout is not None:
            stdio.printf("[link] 链接器输出:\n%s\n", result.stdout)
        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 文件路径
    name_len: t.CSizeT = string.strlen(module_name)
    dir_len: t.CSizeT = string.strlen(temp_dir)
    ir_path: bytes = _mbuddy.alloc(dir_len + name_len + 5)
    if ir_path is None:
        result.Success = 0
        result.ErrorMsg = "内存分配失败"
        return result
    viperlib.snprintf(ir_path, dir_len + name_len + 5, "%s/%s.ll", temp_dir, module_name)

    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