import t, c
from stdint import *
import stdio
import string
import stdlib
import memhub
import viperlib
import w32.fileio as fileio
import w32.win32file
import w32.win32base
import ast
import llvmlite
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.HandlesExprCall as HandlesExprCall
import lib.core.Handles.HandlesImports as HandlesImports
import lib.core.IncludesScanner as IncludesScanner
import lib.core.StubMerger as StubMerger
import lib.Projectrans.Config as Config
import lib.StubGen.Converter as StubConverter

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

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

# pyi 缓冲区大小（256KB）
PYI_BUF_SIZE: t.CSizeT = 262144


# ============================================================
# RunPhase1 - Phase1: 扫描 includes 目录，按需翻译并生成 stub
#
# 对每个 .py 文件检查 stub 是否已存在，若不存在则翻译并分离 stub。
# "按需翻译"：仅对 stub 不存在的文件执行翻译，避免重复工作。
#
# Args:
#   mb: 内存池
#   includes_dir: includes 目录路径
#   temp_dir: 临时目录（保存 stub）
#   log: 日志器
#
# Returns:
#   0 成功，非 0 失败
# ============================================================
def RunPhase1(mb: memhub.MemBuddy | t.CPtr, includes_dir: str, temp_dir: str,
              log: VLogger.Logger | t.CPtr) -> int:
    """Phase1: 扫描 includes 目录，按需翻译并生成 stub"""
    if includes_dir is None or temp_dir is None:
        stdio.printf("[Phase1] includes_dir 或 temp_dir 为空，跳过\n")
        return 1

    if log is not None:
        log.banner("Phase1: 扫描 includes（按需翻译）")

    # 扫描 includes 目录
    result: IncludesScanner.ScanResult | t.CPtr = IncludesScanner.scan_includes(mb, includes_dir)
    if result is None:
        stdio.printf("[Phase1] 扫描失败\n")
        return 1

    # 读取 _sha1_map.txt 获取需要的 includes SHA1 集合
    # 优先使用 Projectrans.py 生成的 _sha1_map.txt（含依赖分析，只包含需要的 includes）
    # 若不存在，则从扫描结果生成（包含所有 includes，可能导致结构体表溢出）
    sha1_set: str = stdlib.malloc(StubMerger.MAX_INCLUDES_SHA1 * 17)
    if sha1_set is None:
        stdio.printf("[Phase1] sha1_set 分配失败\n")
        return 1
    string.memset(sha1_set, 0, StubMerger.MAX_INCLUDES_SHA1 * 17)
    set_count: int = StubMerger._load_includes_sha1_set(mb, temp_dir, sha1_set)
    if set_count <= 0:
        stdio.printf("[Phase1] _sha1_map.txt 不存在或为空，从扫描结果生成\n")
        StubMerger.WriteIncludesSha1Map(mb, temp_dir, result, None, 0)
        set_count = StubMerger._load_includes_sha1_set(mb, temp_dir, sha1_set)
        if set_count < 0:
            stdio.printf("[Phase1] 无法加载 _sha1_map.txt，跳过 Phase1\n")
            return 1
    # 构建模块 SHA1 映射（供跨模块函数调用名混淆使用）
    td_len_p1map: t.CSizeT = string.strlen(temp_dir)
    p1_sha1_arr: bytes = stdlib.malloc(StubMerger.MAX_INCLUDES * 17)
    p1_mod_arr: bytes = stdlib.malloc(StubMerger.MAX_INCLUDES * 64)
    p1_inc_count: int = 0
    if p1_sha1_arr is not None and p1_mod_arr is not None:
        p1_inc_count = StubMerger._BuildIncludesSha1Map(temp_dir, td_len_p1map, p1_sha1_arr, p1_mod_arr)
    HandlesExprCall.set_module_sha1_map(p1_sha1_arr, p1_mod_arr, p1_inc_count)

    # 初始化 AST 表（只需一次）
    ast._init_tables(mb)

    # FileEntry 结构体大小（Phase 1a 和 1b 共用）
    entry_size: t.CSizeT = IncludesScanner.FileEntry.__sizeof__()

    # ============================================================
    # Phase 1a-pre: 扫描所有 includes 文件的 import 依赖
    #
    # 只处理 Import/ImportFrom 语句，跳过 ClassDef/FunctionDef 等，
    # 不调用 resolve_annotation_type，避免 list[...] 等不支持的语法触发 crash。
    # 生成 .deps.txt 供依赖图按需翻译使用。
    # ============================================================
    p1a_registered: int = 0
    p1a_skipped: int = 0
    p1a_failed: int = 0

    for i in range(result.Count):
        entry_addr_a: t.CUInt64T = t.CUInt64T(result.Entries) + i * entry_size
        entry_a: IncludesScanner.FileEntry | t.CPtr = (IncludesScanner.FileEntry | t.CPtr)(t.CVoid(entry_addr_a, t.CPtr))
        if entry_a is None:
            p1a_failed += 1
            continue

        sha1_a: str = entry_a.Sha1
        if sha1_a is None:
            p1a_failed += 1
            continue

        if StubMerger._is_in_sha1_set(sha1_a, sha1_set, set_count) == 0:
            p1a_skipped += 1
            continue

        # 读取文件内容
        file_path_a: str = entry_a.Path
        f_a: fileio.File | t.CPtr = fileio.File(file_path_a, fileio.MODE.R)
        if f_a.closed:
            p1a_failed += 1
            continue

        src_buf_a: bytes = stdlib.malloc(SRC_BUF_SIZE)
        if src_buf_a is None:
            f_a.close()
            p1a_failed += 1
            continue

        bytes_read_a: LONG = f_a.read_all(src_buf_a, SRC_BUF_SIZE)
        f_a.close()
        if bytes_read_a <= 0:
            stdlib.free(src_buf_a)
            p1a_failed += 1
            continue
        if bytes_read_a < SRC_BUF_SIZE:
            src_buf_a[bytes_read_a] = 0
        else:
            src_buf_a[SRC_BUF_SIZE - 1] = 0

        # 解析 AST
        lx_a: ast.Lexer | t.CPtr = ast.new_lexer(mb)
        if lx_a is None:
            stdlib.free(src_buf_a)
            p1a_failed += 1
            continue
        ast._lexer_init(lx_a, src_buf_a, mb)
        tokens_a: ast.Token | t.CPtr = ast.tokenize(lx_a)
        tree_a: ast.AST | t.CPtr = ast.parse_tokens(mb, tokens_a)
        if tree_a is None:
            stdlib.free(src_buf_a)
            p1a_failed += 1
            continue

        # 创建 Translator，设置 _declare_only=2（只扫描 import，跳过 ClassDef）
        tr_a: HandlesTranslator.Translator | t.CPtr = HandlesTranslator.Translator()
        if tr_a is None:
            stdlib.free(src_buf_a)
            p1a_failed += 1
            continue
        tr_a.ModuleSha1 = sha1_a
        tr_a._declare_only = 2
        tr_a.CurrentPackage = HandlesImports.compute_package_from_relpath(mb, entry_a.RelPath)
        HandlesType.set_current_file(file_path_a)
        HandlesType.clear_cdefine_constants()
        HandlesStruct.reset_visible_structs(mb, 0)
        ret_a: int = tr_a.translate(tree_a)
        if ret_a != 0:
            p1a_failed += 1
        else:
            p1a_registered += 1

        # 生成 .deps.txt（记录依赖模块名，供依赖图按需翻译使用）
        td_len_a: t.CSizeT = string.strlen(temp_dir)
        deps_path_a: bytes = stdlib.malloc(td_len_a + 32)
        if deps_path_a is not None:
            viperlib.snprintf(deps_path_a, td_len_a + 32, "%s/%s.deps.txt", temp_dir, sha1_a)
            df_a: fileio.File | t.CPtr = fileio.File(deps_path_a, fileio.MODE.W)
            if not df_a.closed:
                if tr_a._imported_modules is not None:
                    dl_a: t.CSizeT = string.strlen(tr_a._imported_modules)
                    df_a.write(tr_a._imported_modules, dl_a)
                df_a.close()
            stdlib.free(deps_path_a)

        # 释放 Translator 的 C malloc 资源
        if tr_a._global_names is not None:
            stdlib.free(tr_a._global_names)
        if tr_a._nonlocal_names is not None:
            stdlib.free(tr_a._nonlocal_names)

        stdlib.free(src_buf_a)

    # ============================================================
    # 依赖图按需翻译：构建可达 SHA1 集合
    #
    # 从 Config.SourceDir 的源文件开始，解析 import 语句，递归收集
    # 可达的 includes 文件 SHA1。Phase 1b 只翻译可达集合中的文件，
    # 避免翻译不需要的 includes（如 Test 不依赖 llvmlite，则不翻译）。
    # 如果 Config.SourceDir 不可用（仅运行 Phase1），回退到 sha1_set。
    # ============================================================
    reachable_set: str = None
    reachable_count: int = 0
    use_reachable: int = 0
    if Config.SourceDir is not None:
        reachable_set = stdlib.malloc(StubMerger.MAX_INCLUDES_SHA1 * 17)
        if reachable_set is not None:
            string.memset(reachable_set, 0, StubMerger.MAX_INCLUDES_SHA1 * 17)
            reachable_count = StubMerger._BuildReachableSha1Set(mb, Config.SourceDir, temp_dir, reachable_set)
            if reachable_count > 0:
                use_reachable = 1
                # 用可达集合重新生成 _sha1_map.txt（按图求索的最终产物）
                # Phase B+ 只遍历这些条目，避免编译不需要的 includes（如 asm.py）
                StubMerger.WriteIncludesSha1Map(mb, temp_dir, result, reachable_set, reachable_count)
                # 重新加载 sha1_set，使后续 Phase 1a-pre/1a/1b 的过滤也使用可达集合
                string.memset(sha1_set, 0, StubMerger.MAX_INCLUDES_SHA1 * 17)
                set_count = StubMerger._load_includes_sha1_set(mb, temp_dir, sha1_set)
            else:
                stdlib.free(reachable_set)
                reachable_set = None

    # ============================================================
    # Phase 1a: 注册可达 includes 文件的 struct/enum/union
    #
    # 解决字母序依赖问题：按字母顺序翻译时，后面的 struct 未注册
    # 前面就需要用。Phase 1a 先注册所有 struct/enum/union 到全局表，
    # Phase 1b 全量翻译时 struct 已注册，走 existing 路径只翻译方法体。
    # 只处理可达文件，避免翻译不需要的 includes（如 Test 不依赖 ast 模块）。
    # ============================================================
    p1a_reg: int = 0
    p1a_skp: int = 0
    p1a_fl: int = 0

    for i in range(result.Count):
        entry_addr_r: t.CUInt64T = t.CUInt64T(result.Entries) + i * entry_size
        entry_r: IncludesScanner.FileEntry | t.CPtr = (IncludesScanner.FileEntry | t.CPtr)(t.CVoid(entry_addr_r, t.CPtr))
        if entry_r is None:
            p1a_fl += 1
            continue

        sha1_r: str = entry_r.Sha1
        if sha1_r is None:
            p1a_fl += 1
            continue

        # 按需翻译过滤：只处理可达文件
        if use_reachable != 0:
            if StubMerger._is_in_sha1_set(sha1_r, reachable_set, reachable_count) == 0:
                p1a_skp += 1
                continue
        else:
            if StubMerger._is_in_sha1_set(sha1_r, sha1_set, set_count) == 0:
                p1a_skp += 1
                continue

        # 读取文件内容
        file_path_r: str = entry_r.Path
        f_r: fileio.File | t.CPtr = fileio.File(file_path_r, fileio.MODE.R)
        if f_r.closed:
            p1a_fl += 1
            continue

        src_buf_r: bytes = stdlib.malloc(SRC_BUF_SIZE)
        if src_buf_r is None:
            f_r.close()
            p1a_fl += 1
            continue

        bytes_read_r: LONG = f_r.read_all(src_buf_r, SRC_BUF_SIZE)
        f_r.close()
        if bytes_read_r <= 0:
            stdlib.free(src_buf_r)
            p1a_fl += 1
            continue
        if bytes_read_r < SRC_BUF_SIZE:
            src_buf_r[bytes_read_r] = 0
        else:
            src_buf_r[SRC_BUF_SIZE - 1] = 0

        # 解析 AST
        lx_r: ast.Lexer | t.CPtr = ast.new_lexer(mb)
        if lx_r is None:
            stdlib.free(src_buf_r)
            p1a_fl += 1
            continue
        ast._lexer_init(lx_r, src_buf_r, mb)
        tokens_r: ast.Token | t.CPtr = ast.tokenize(lx_r)
        tree_r: ast.AST | t.CPtr = ast.parse_tokens(mb, tokens_r)
        if tree_r is None:
            stdlib.free(src_buf_r)
            p1a_fl += 1
            continue

        # 创建 Translator，设置 _declare_only=1（注册 struct/enum/union）
        tr_r: HandlesTranslator.Translator | t.CPtr = HandlesTranslator.Translator()
        if tr_r is None:
            stdlib.free(src_buf_r)
            p1a_fl += 1
            continue
        tr_r.ModuleSha1 = sha1_r
        tr_r._declare_only = 1
        tr_r.CurrentPackage = HandlesImports.compute_package_from_relpath(mb, entry_r.RelPath)
        HandlesType.set_current_file(file_path_r)
        HandlesType.clear_cdefine_constants()
        HandlesStruct.reset_visible_structs(mb, 0)
        ret_r: int = tr_r.translate(tree_r)
        if ret_r != 0:
            p1a_fl += 1
        else:
            p1a_reg += 1

        # 释放 Translator 的 C malloc 资源
        if tr_r._global_names is not None:
            stdlib.free(tr_r._global_names)
        if tr_r._nonlocal_names is not None:
            stdlib.free(tr_r._nonlocal_names)
        stdlib.free(src_buf_r)

    # ============================================================
    # Phase 1b: 全量翻译（struct 已注册，走 existing 路径翻译方法体）
    # ============================================================
    # 遍历每个文件
    translated: int = 0
    skipped: int = 0
    failed: int = 0

    for i in range(result.Count):
        # 获取 entry
        entry_addr: t.CUInt64T = t.CUInt64T(result.Entries) + i * entry_size
        entry: IncludesScanner.FileEntry | t.CPtr = (IncludesScanner.FileEntry | t.CPtr)(t.CVoid(entry_addr, t.CPtr))
        if entry is None:
            failed += 1
            continue

        sha1: str = entry.Sha1
        if sha1 is None:
            failed += 1
            continue

        # 检查 SHA1 是否在翻译集合中（按需翻译：只翻译可达的文件）
        in_set: int = 0
        if use_reachable != 0:
            in_set = StubMerger._is_in_sha1_set(sha1, reachable_set, reachable_count)
        else:
            in_set = StubMerger._is_in_sha1_set(sha1, sha1_set, set_count)
        if in_set == 0:
            skipped += 1
            continue

        # 构造 stub 路径: {temp_dir}/{sha1}.stub.ll
        dir_len: t.CSizeT = string.strlen(temp_dir)
        sha1_len: t.CSizeT = string.strlen(sha1)
        stub_path: bytes = stdlib.malloc(dir_len + sha1_len + 16)
        if stub_path is None:
            failed += 1
            continue
        viperlib.snprintf(stub_path, dir_len + sha1_len + 16, "%s/%s.stub.ll", temp_dir, sha1)

        # 检查 stub 是否已存在（按需翻译：跳过已存在的）
        sf: fileio.File | t.CPtr = fileio.File(stub_path, fileio.MODE.R)
        if not sf.closed:
            sf.close()
            # 检查 text.ll 是否也存在（虚表扫描需要 text.ll）
            text_path: bytes = stdlib.malloc(dir_len + sha1_len + 16)
            if text_path is not None:
                viperlib.snprintf(text_path, dir_len + sha1_len + 16, "%s/%s.text.ll", temp_dir, sha1)
                tf: fileio.File | t.CPtr = fileio.File(text_path, fileio.MODE.R)
                if not tf.closed:
                    tf.close()
                    stdlib.free(text_path)
                    stdlib.free(stub_path)
                    continue
                stdlib.free(text_path)
            # text.ll 不存在，需要重新翻译
            rp: str = entry.RelPath
            stdio.printf("[Phase1] text.ll 不存在，重新翻译: %s (sha1=%s)\n", rp, sha1)
        else:
            # stub 不存在，需要翻译
            rp: str = entry.RelPath
            stdio.printf("[Phase1] 翻译: %s (sha1=%s)\n", rp, sha1)
            stdio.fflush(0)

        stdlib.free(stub_path)

        # 读取文件内容
        file_path: str = entry.Path
        f: fileio.File | t.CPtr = fileio.File(file_path, fileio.MODE.R)
        if f is None:
            stdio.printf("[Phase1] 无法打开(None): %s\n", file_path)
            stdio.fflush(0)
            failed += 1
            continue
        if f.closed:
            stdio.printf("[Phase1] 无法打开: %s\n", file_path)
            stdio.fflush(0)
            failed += 1
            continue

        src_buf: bytes = stdlib.malloc(SRC_BUF_SIZE)
        if src_buf is None:
            f.close()
            failed += 1
            continue

        bytes_read: LONG = f.read_all(src_buf, SRC_BUF_SIZE)
        f.close()
        if bytes_read <= 0:
            stdio.printf("[Phase1] 读取失败: %s\n", file_path)
            stdio.fflush(0)
            stdlib.free(src_buf)
            failed += 1
            continue
        if bytes_read < SRC_BUF_SIZE:
            src_buf[bytes_read] = 0
        else:
            src_buf[SRC_BUF_SIZE - 1] = 0

        # 解析 AST
        lx: ast.Lexer | t.CPtr = ast.new_lexer(mb)
        if lx is None:
            stdlib.free(src_buf)
            failed += 1
            continue
        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:
            stdio.printf("[Phase1] AST 解析失败: %s\n", file_path)
            stdio.fflush(0)
            stdlib.free(src_buf)
            failed += 1
            continue

        # 生成 .pyi 存根文件（直接遍历 AST，不依赖 PythonToStubConverter）
        pyi_buf: bytes = stdlib.malloc(PYI_BUF_SIZE)
        if pyi_buf is not None:
            pyi_pos: t.CSizeT = StubConverter._GeneratePyiFromAst(mb, tree, entry.RelPath, pyi_buf, PYI_BUF_SIZE)
            if pyi_pos > 0:
                pyi_path: bytes = stdlib.malloc(dir_len + 32)
                if pyi_path is not None:
                    viperlib.snprintf(pyi_path, dir_len + 32, "%s/%s.pyi", temp_dir, sha1)
                    pf: fileio.File | t.CPtr = fileio.File(pyi_path, fileio.MODE.W)
                    if not pf.closed:
                        pf.write(pyi_buf, pyi_pos)
                        pf.close()
                    stdlib.free(pyi_path)
            stdlib.free(pyi_buf)

        # 翻译 AST → LLVM IR
        tr: HandlesTranslator.Translator | t.CPtr = HandlesTranslator.Translator()
        if tr is None:
            stdlib.free(src_buf)
            failed += 1
            continue
        tr.ModuleSha1 = sha1
        tr.CurrentPackage = HandlesImports.compute_package_from_relpath(mb, entry.RelPath)
        HandlesType.set_current_file(file_path)
        HandlesType.clear_cdefine_constants()
        HandlesStruct.reset_visible_structs(mb, 0)
        ret: int = tr.translate(tree)
        if ret != 0:
            stdio.printf("[Phase1] 翻译失败: %s\n", file_path)
            stdlib.free(src_buf)
            if tr._global_names is not None:
                stdlib.free(tr._global_names)
            if tr._nonlocal_names is not None:
                stdlib.free(tr._nonlocal_names)
            failed += 1
            continue

        # dump stub IR (declarations only)
        PHASE1_IR_SIZE: t.CSizeT = 262144
        stub_buf: bytes = stdlib.malloc(PHASE1_IR_SIZE)
        if stub_buf is None:
            stdlib.free(src_buf)
            if tr._global_names is not None:
                stdlib.free(tr._global_names)
            if tr._nonlocal_names is not None:
                stdlib.free(tr._nonlocal_names)
            failed += 1
            continue
        tr.dump_ir(stub_buf, PHASE1_IR_SIZE, llvmlite.OUTPUT_STUB)
        stub_len: t.CSizeT = string.strlen(stub_buf)

        # save stub.ll
        dir_len_p1: t.CSizeT = string.strlen(temp_dir)
        stub_path_p1: bytes = stdlib.malloc(dir_len_p1 + 32)
        if stub_path_p1 is not None:
            viperlib.snprintf(stub_path_p1, dir_len_p1 + 32, "%s/%s.stub.ll", temp_dir, sha1)
            sf_p1: fileio.File | t.CPtr = fileio.File(stub_path_p1, fileio.MODE.W)
            if not sf_p1.closed:
                sf_p1.write(stub_buf, stub_len)
                sf_p1.close()
            stdlib.free(stub_path_p1)
        stdlib.free(stub_buf)

        # dump text IR (definitions only)
        text_buf: bytes = stdlib.malloc(PHASE1_IR_SIZE)
        if text_buf is None:
            stdlib.free(src_buf)
            if tr._global_names is not None:
                stdlib.free(tr._global_names)
            if tr._nonlocal_names is not None:
                stdlib.free(tr._nonlocal_names)
            failed += 1
            continue
        tr.dump_ir(text_buf, PHASE1_IR_SIZE, llvmlite.OUTPUT_TEXT)
        text_len: t.CSizeT = string.strlen(text_buf)

        # save text.ll
        text_path_p1: bytes = stdlib.malloc(dir_len_p1 + 32)
        if text_path_p1 is not None:
            viperlib.snprintf(text_path_p1, dir_len_p1 + 32, "%s/%s.text.ll", temp_dir, sha1)
            tf_p1: fileio.File | t.CPtr = fileio.File(text_path_p1, fileio.MODE.W)
            if not tf_p1.closed:
                tf_p1.write(text_buf, text_len)
                tf_p1.close()
            stdlib.free(text_path_p1)
        stdlib.free(text_buf)

        # save dependencies (_imported_modules) for Phase B
        deps_path_p1: bytes = stdlib.malloc(dir_len_p1 + 32)
        if deps_path_p1 is not None:
            viperlib.snprintf(deps_path_p1, dir_len_p1 + 32, "%s/%s.deps.txt", temp_dir, sha1)
            df_p1: fileio.File | t.CPtr = fileio.File(deps_path_p1, fileio.MODE.W)
            if not df_p1.closed:
                if tr._imported_modules is not None:
                    dl: t.CSizeT = string.strlen(tr._imported_modules)
                    df_p1.write(tr._imported_modules, dl)
                df_p1.close()
            stdlib.free(deps_path_p1)

        # 释放 Translator 的 C malloc 资源
        if tr._global_names is not None:
            stdlib.free(tr._global_names)
        if tr._nonlocal_names is not None:
            stdlib.free(tr._nonlocal_names)

        # 所有 dump_ir 完成，src_buf 不再需要
        stdlib.free(src_buf)

        translated += 1

    stdio.printf("[Phase1] 完成: 翻译=%d 跳过=%d 失败=%d\n", translated, skipped, failed)

    # 释放可达 SHA1 集合（如果分配了）
    if reachable_set is not None:
        stdlib.free(reachable_set)

    return 0