import t, c
from stdint import *
import ast
import memhub
import string
import stdlib
import viperlib
import stdio
import lib.core.Handles.HandlesBase as HandlesBase
import lib.core.Handles.HandlesTranslator as HT
import lib.core.Handles.HandlesStruct as HandlesStruct


# ============================================================
# HandlesImports - 导入语句处理（Mixin 继承模式）
#
# 管理 _imported_modules 和 _from_imports 字符串
# 注意：str = bytes = t.CChar | t.CPtr = i8*
# ============================================================


# ============================================================
# 模块级工具函数（保留供外部调用）
# ============================================================

# ============================================================
# 添加已导入模块名
# ============================================================
def add_imported_module(pool: memhub.MemBuddy | t.CPtr,
                        imported_modules: str,
                        name: str) -> str:
    """记录已导入的模块名，返回新的 imported_modules 字符串指针"""
    if name is None:
        return imported_modules
    if imported_modules is None:
        nlen: t.CSizeT = string.strlen(name)
        buf: t.CChar | t.CPtr = pool.alloc(nlen + 1)
        if buf is not None:
            string.strcpy(buf, name)
            return buf
        return None
    else:
        old_len: t.CSizeT = string.strlen(imported_modules)
        name_len: t.CSizeT = string.strlen(name)
        new_len: t.CSizeT = old_len + 1 + name_len
        buf2: t.CChar | t.CPtr = pool.alloc(new_len + 1)
        if buf2 is not None:
            string.strcpy(buf2, imported_modules)
            buf2[old_len] = ' '
            string.strcpy(buf2 + old_len + 1, name)
            return buf2
        return imported_modules


# ============================================================
# 检查模块是否已导入
# ============================================================
def is_module_imported(imported_modules: str, name: str) -> int:
    """检查模块是否已导入（词边界精确匹配，避免子串误匹配）"""
    if name is None or imported_modules is None:
        return 0
    name_len: t.CSizeT = string.strlen(name)
    cur: t.CChar | t.CPtr = imported_modules
    total_len: t.CSizeT = string.strlen(imported_modules)
    ci: t.CSizeT = 0
    while ci < total_len:
        while ci < total_len and cur[ci] == ' ':
            ci += 1
        if ci >= total_len:
            break
        word_start: t.CSizeT = ci
        while ci < total_len and cur[ci] != ' ':
            ci += 1
        word_len: t.CSizeT = ci - word_start
        if word_len == name_len:
            match: int = 1
            for ei in range(name_len):
                if cur[word_start + ei] != name[ei]:
                    match = 0
                    break
            if match == 1:
                return 1
    return 0


# ============================================================
# 添加 from-import 名称映射
# ============================================================
def add_from_import(pool: memhub.MemBuddy | t.CPtr,
                    from_imports: str,
                    local_name: str,
                    module_name: str,
                    original_name: str = None) -> str:
    """添加 from-import 映射，返回新的 from_imports 字符串

    格式:
        无别名: "local_name:module_name"
        有别名: "local_name:module_name:original_name"
        （original_name 为源模块中的真实函数名，用于跨模块 SHA1 修饰）
    """
    if local_name is None or module_name is None:
        return from_imports
    entry: t.CChar | t.CPtr = pool.alloc(256)
    if entry is None:
        return from_imports
    if original_name is not None and original_name != local_name:
        viperlib.snprintf(entry, 256, "%s:%s:%s", local_name, module_name, original_name)
    else:
        viperlib.snprintf(entry, 256, "%s:%s", local_name, module_name)
    if from_imports is None:
        return entry
    else:
        old_len: t.CSizeT = string.strlen(from_imports)
        entry_len: t.CSizeT = string.strlen(entry)
        new_len: t.CSizeT = old_len + 1 + entry_len
        buf: t.CChar | t.CPtr = pool.alloc(new_len + 1)
        if buf is not None:
            string.strcpy(buf, from_imports)
            buf[old_len] = ' '
            string.strcpy(buf + old_len + 1, entry)
            return buf
        return from_imports


# ============================================================
# 查找 from-import 名称 → 返回原始函数名（别名场景）
#
# 对于 from X import Y as Z，返回 Y（源模块中的真实函数名）。
# 无别名时返回 None（local_name 即为原始名）。
# ============================================================
def lookup_from_import_original(from_imports: str,
                                local_name: str) -> str:
    """查找 from-import 别名对应的原始函数名，返回 None=无别名或未找到"""
    if local_name is None or from_imports is None:
        return None
    name_len: t.CSizeT = string.strlen(local_name)
    cur: t.CChar | t.CPtr = from_imports
    ci: t.CSizeT = 0
    total_len: t.CSizeT = string.strlen(from_imports)
    while ci < total_len:
        # 跳过前导空格
        while ci < total_len and cur[ci] == ' ':
            ci += 1
        if ci >= total_len:
            break
        # 找到第一个 ':' 的位置
        colon1: t.CSizeT = ci
        while colon1 < total_len and cur[colon1] != ':' and cur[colon1] != ' ':
            colon1 += 1
        if colon1 >= total_len or cur[colon1] != ':':
            break
        entry_name_len: t.CSizeT = colon1 - ci
        # 跳过 star import
        if entry_name_len == 1 and cur[ci] == '*':
            ci = colon1
            while ci < total_len and cur[ci] != ' ':
                ci += 1
            continue
        # 比较名称
        if entry_name_len == name_len:
            match: int = 1
            ei: t.CSizeT = 0
            while ei < name_len:
                if cur[ci + ei] != local_name[ei]:
                    match = 0
                    break
                ei += 1
            if match == 1:
                # 找到匹配，检查是否有第三个字段（原始名）
                pos: t.CSizeT = colon1 + 1
                # 跳过模块名
                while pos < total_len and cur[pos] != ':' and cur[pos] != ' ' and cur[pos] != '\0':
                    pos += 1
                if pos < total_len and cur[pos] == ':':
                    # 有第三个字段: original_name
                    orig_start: t.CSizeT = pos + 1
                    orig_end: t.CSizeT = orig_start
                    while orig_end < total_len and cur[orig_end] != ' ' and cur[orig_end] != '\0':
                        orig_end += 1
                    orig_len: t.CSizeT = orig_end - orig_start
                    if orig_len > 0:
                        orig_buf: str = cur + orig_start
                        # 返回指向内部的指针（调用方需在使用期内保持 from_imports 有效）
                        return orig_buf
                    return None
                return None
        # 跳到下一个条目
        ci = colon1
        while ci < total_len and cur[ci] != ' ':
            ci += 1
    return None


# ============================================================
# 查找 from-import 名称 → 返回模块名或 None
#
# allow_star_fallback: 是否允许 star import 回退（默认 1=允许）。
#   当查询明确模块名（如 "BuildPipeline"）时，应传 0 禁用回退，
#   避免被 star import 模块名误导（如 from stdint import * 后
#   查询 "BuildPipeline" 错误回退到 "stdint"）。
# ============================================================
def lookup_from_import(from_imports: str, name: str,
                       allow_star_fallback: int = 1) -> str:
    """查找 from-import 名称，返回模块名或 None

    支持 star import: 如果 from_imports 中有 "*:module" 条目，
    且未找到精确名称匹配，且 allow_star_fallback != 0，则返回 star import 的模块名。
    """
    if name is None or from_imports is None:
        return None
    name_len: t.CSizeT = string.strlen(name)
    cur: t.CChar | t.CPtr = from_imports
    ci: t.CSizeT = 0
    total_len: t.CSizeT = string.strlen(from_imports)
    star_mod: str = None
    while ci < total_len:
        # 跳过前导空格
        while ci < total_len and cur[ci] == ' ':
            ci += 1
        if ci >= total_len:
            break
        # 找到 ':' 的位置
        colon_pos: t.CSizeT = ci
        while colon_pos < total_len and cur[colon_pos] != ':' and cur[colon_pos] != ' ':
            colon_pos += 1
        if colon_pos >= total_len or cur[colon_pos] != ':':
            break
        entry_name_len: t.CSizeT = colon_pos - ci
        # 检测 star import ("*:module")
        if entry_name_len == 1 and cur[ci] == '*':
            mod_start: t.CSizeT = colon_pos + 1
            mod_end: t.CSizeT = mod_start
            while mod_end < total_len and cur[mod_end] != ' ' and cur[mod_end] != '\0' and cur[mod_end] != ':':
                mod_end += 1
            star_mod = cur + mod_start
        # 比较名称
        elif entry_name_len == name_len:
            match: int = 1
            ei: t.CSizeT = 0
            while ei < name_len:
                if cur[ci + ei] != name[ei]:
                    match = 0
                    break
                ei += 1
            if match == 1:
                mod_start2: t.CSizeT = colon_pos + 1
                mod_end2: t.CSizeT = mod_start2
                # 模块名结束于: 空格、null、或第二个':'（别名格式的分隔符）
                while mod_end2 < total_len and cur[mod_end2] != ' ' and cur[mod_end2] != '\0' and cur[mod_end2] != ':':
                    mod_end2 += 1
                return cur + mod_start2
        # 跳到下一个条目
        ci = colon_pos
        while ci < total_len and cur[ci] != ' ':
            ci += 1
    # 未找到精确匹配，仅在允许时回退到 star import
    if allow_star_fallback == 0:
        return None
    return star_mod


# ============================================================
# 全局 re-export 映射表（module_sha1 + func_name → source_module）
#
# 当 ast/__init__.py 中有 from .lexer import _lexer_init 时，
# _lexer_init 被 re-export 为 ast._lexer_init。
# 跨模块调用 ast._lexer_init(...) 需使用 ast.lexer 的 SHA1 混淆，
# 而非 ast 的 SHA1。
#
# 数据布局：
#   _g_reexport_sha1s: 每条 17 字节（SHA1 16字符 + null）
#   _g_reexport_funcs: 每条 64 字节（函数名 + null）
#   _g_reexport_srcs:  每条 64 字节（源模块名 + null）
# ============================================================
MAX_REEXPORT: t.CDefine = 512

_g_reexport_sha1s: bytes = None
_g_reexport_funcs: bytes = None
_g_reexport_srcs: bytes = None
_g_reexport_count: int = 0


# ============================================================
# register_reexport - 注册 re-export 映射
#
# 在 HandleImportFromNames 中调用：当处理 from .X import Y 时，
# 记录当前模块（通过 ModuleSha1）re-export 了函数 Y，源模块为 X。
#
# 幂等：重复注册相同映射不会增加条目。
# 内存用 stdlib.malloc 分配（全局存储器，跨 Phase 持久化）。
# ============================================================
def register_reexport(mod_sha1: str, func_name: str, source_module: str) -> int:
    """注册 re-export 映射（幂等），返回 0 成功，-1 失败"""
    global _g_reexport_sha1s, _g_reexport_funcs, _g_reexport_srcs, _g_reexport_count
    if mod_sha1 is None or func_name is None or source_module is None:
        return -1
    # 懒初始化
    if _g_reexport_sha1s is None:
        _g_reexport_sha1s = stdlib.malloc(MAX_REEXPORT * 17)
        _g_reexport_funcs = stdlib.malloc(MAX_REEXPORT * 64)
        _g_reexport_srcs = stdlib.malloc(MAX_REEXPORT * 64)
        if _g_reexport_sha1s is None or _g_reexport_funcs is None or _g_reexport_srcs is None:
            return -1
        string.memset(_g_reexport_sha1s, 0, MAX_REEXPORT * 17)
        string.memset(_g_reexport_funcs, 0, MAX_REEXPORT * 64)
        string.memset(_g_reexport_srcs, 0, MAX_REEXPORT * 64)
    if _g_reexport_count >= MAX_REEXPORT:
        return -1
    # 幂等检查：查找是否已存在相同映射
    for i in range(_g_reexport_count):
        sidx: t.CSizeT = t.CSizeT(i) * 17
        fidx: t.CSizeT = t.CSizeT(i) * 64
        if string.strcmp(_g_reexport_sha1s + sidx, mod_sha1) == 0:
            if string.strcmp(_g_reexport_funcs + fidx, func_name) == 0:
                # 已存在，更新 source_module（以防变化）
                string.strcpy(_g_reexport_srcs + fidx, source_module)
                return 0
    # 添加新条目
    idx2: t.CSizeT = t.CSizeT(_g_reexport_count) * 17
    fidx2: t.CSizeT = t.CSizeT(_g_reexport_count) * 64
    string.strcpy(_g_reexport_sha1s + idx2, mod_sha1)
    string.strcpy(_g_reexport_funcs + fidx2, func_name)
    string.strcpy(_g_reexport_srcs + fidx2, source_module)
    _g_reexport_count += 1
    return 0


# ============================================================
# lookup_reexport - 查找 re-export 映射
#
# 给定模块 SHA1 和函数名，返回源模块名（如 "ast.lexer"）或 None。
# 调用方通过 _lookup_module_sha1(source_module) 获取源模块的 SHA1。
# ============================================================
def lookup_reexport(mod_sha1: str, func_name: str) -> str:
    """查找 re-export 映射，返回源模块名或 None"""
    if mod_sha1 is None or func_name is None:
        return None
    if _g_reexport_sha1s is None or _g_reexport_count <= 0:
        return None
    for i in range(_g_reexport_count):
        sidx: t.CSizeT = t.CSizeT(i) * 17
        fidx: t.CSizeT = t.CSizeT(i) * 64
        if string.strcmp(_g_reexport_sha1s + sidx, mod_sha1) == 0:
            if string.strcmp(_g_reexport_funcs + fidx, func_name) == 0:
                return _g_reexport_srcs + fidx
    return None


# ============================================================
# _resolve_relative_module - 解析相对导入为完整模块名
#
# 对于 from .X import Y（level=1, module="X"）在包 pkg 中：
#   解析为 "pkg.X"
# 对于 from . import Y（level=1, module=None）在包 pkg 中：
#   解析为 "pkg"
# 对于 from ..X import Y（level=2, module="X"）在包 pkg.sub 中：
#   解析为 "pkg.X"（先从 pkg.sub 上溯一级到 pkg，再追加 .X）
#
# Args:
#   pool: 内存池
#   current_package: 当前文件所属包名（如 "llvmlite"），None 表示无包
#   level: 相对导入级别（0=绝对，1=., 2=..）
#   module: ImportFrom 的 module 字段（可能为 None）
#
# Returns:
#   解析后的完整模块名；绝对导入（level<=0）直接返回 module；
#   无法解析时返回 module（回退到原始值）
# ============================================================
def _resolve_relative_module(pool: memhub.MemBuddy | t.CPtr,
                            current_package: str,
                            level: t.CInt,
                            module: str) -> str:
    """解析相对导入为完整模块名"""
    if level <= 0:
        return module
    if current_package is None:
        return module

    # 从 current_package 开始，上溯 (level-1) 级
    pkg: str = current_package
    up: t.CInt = level - 1
    while up > 0:
        pkg_len: t.CSizeT = string.strlen(pkg)
        last_dot: t.CSizeT = 0
        found: int = 0
        i: t.CSizeT = 0
        while i < pkg_len:
            if pkg[i] == '.':
                last_dot = i
                found = 1
            i += 1
        if found == 0:
            # 无更多父级，包变为空
            pkg = None
            break
        # 截断到最后一个 '.' 处
        new_pkg: str = pool.alloc(last_dot + 1)
        if new_pkg is None:
            return module
        string.strncpy(new_pkg, pkg, last_dot)
        new_pkg[last_dot] = '\0'
        pkg = new_pkg
        up -= 1

    if module is None:
        # from . import Y → 模块就是包本身
        return pkg
    if pkg is None:
        # 包已上溯到空，直接用 module
        return module
    # 拼接 pkg + "." + module
    pkg_len2: t.CSizeT = string.strlen(pkg)
    mod_len: t.CSizeT = string.strlen(module)
    buf: str = pool.alloc(pkg_len2 + 1 + mod_len + 1)
    if buf is None:
        return module
    string.strcpy(buf, pkg)
    buf[pkg_len2] = '.'
    string.strcpy(buf + pkg_len2 + 1, module)
    return buf


# ============================================================
# compute_package_from_relpath - 从相对路径计算包名
#
# 包名 = 文件所在目录路径，将 / 和 \ 替换为 .
# 对于顶级文件（无目录分隔符），返回 None
#
# 示例：
#   "llvmlite/__init__.py" → "llvmlite"
#   "llvmlite/__types.py"  → "llvmlite"
#   "ast/parser.py"        → "ast"
#   "ast.py"               → None（顶级文件）
# ============================================================
def compute_package_from_relpath(pool: memhub.MemBuddy | t.CPtr,
                                rel_path: str) -> str:
    """从相对路径计算包名（目录部分，分隔符替换为 .）"""
    if rel_path is None:
        return None
    rlen: t.CSizeT = string.strlen(rel_path)
    # 找最后一个 / 或 \
    last_sep: t.CSizeT = 0
    found: int = 0
    i: t.CSizeT = 0
    while i < rlen:
        ch: t.CChar = rel_path[i]
        if ch == '/' or ch == '\\':
            last_sep = i
            found = 1
        i += 1
    if found == 0:
        # 无目录分隔符 → 顶级文件，无包
        return None
    # 复制目录部分，将 / 和 \ 替换为 .
    dir_len: t.CSizeT = last_sep
    buf: str = pool.alloc(dir_len + 1)
    if buf is None:
        return None
    j: t.CSizeT = 0
    while j < dir_len:
        ch2: t.CChar = rel_path[j]
        if ch2 == '/' or ch2 == '\\':
            buf[j] = '.'
        else:
            buf[j] = ch2
        j += 1
    buf[dir_len] = '\0'
    return buf


# ============================================================
# ImportsHandle - 导入语句处理器（Mixin 继承模式）
#
# 方法版本：直接更新 trans._imported_modules / trans._from_imports
# ============================================================
@t.NoVTable
class ImportsHandle(HandlesBase.Mixin):
    """导入语句处理器：继承 Mixin 获得 Trans 回指针"""

    def __init__(self, trans: HT.Translator | t.CPtr):
        self.Trans = trans

    # ============================================================
    # HandleImport - 处理 import 语句，更新 trans._imported_modules
    # ============================================================
    def HandleImport(self, node: ast.AST | t.CPtr) -> int:
        """处理 import 语句，返回 0"""
        imp: ast.Import | t.CPtr = (ast.Import | t.CPtr)(node)
        if imp is None:
            return 0
        names: list[ast.AST | t.CPtr] | t.CPtr = imp.names
        if names is None:
            return 0
        nn: t.CSizeT = names.__len__()
        for ni in range(nn):
            alias: ast.Alias | t.CPtr = (ast.Alias | t.CPtr)(names.get(ni))
            if alias is not None and alias.name is not None:
                pool_val: memhub.MemBuddy | t.CPtr = self.Trans.Pool
                im_val: str = self.Trans._imported_modules
                self.Trans._imported_modules = add_imported_module(
                    pool_val, im_val, alias.name)
                # 别名也加入导入模块列表（用于模块限定构造器检查 Module.Class()）
                if alias.asname is not None:
                    self.Trans._imported_modules = add_imported_module(
                        pool_val, self.Trans._imported_modules, alias.asname)
                    # 别名也加入 from_imports（asname:fullname），
                    # 使 HT.func() 跨模块调用能通过 from_imports 查找别名→完整模块名→SHA1
                    # 否则 _lookup_module_sha1("HT") 找不到（mod_arr 只存文件名 HandlesTranslator）
                    self.Trans._from_imports = add_from_import(
                        pool_val, self.Trans._from_imports, alias.asname, alias.name)
                else:
                    # 无别名时（import a.b.c），将顶层包名 "a" 也加入 _imported_modules，
                    # 使 a.b.c 作为属性访问的 base 被翻译时能通过 is_module_imported("a") 检查
                    dot_ptr: str = string.strstr(alias.name, '.')
                    if dot_ptr is not None:
                        pkg_len: t.CSizeT = t.CSizeT(t.CUInt64T(dot_ptr) - t.CUInt64T(alias.name))
                        if pkg_len > 0 and pkg_len < 256:
                            pkg_buf: bytes = pool_val.alloc(pkg_len + 1)
                            if pkg_buf is not None:
                                string.strncpy(pkg_buf, alias.name, pkg_len)
                                pkg_buf[pkg_len] = '\0'
                                self.Trans._imported_modules = add_imported_module(
                                    pool_val, self.Trans._imported_modules, pkg_buf)
        return 0

    # ============================================================
    # HandleImportFromModule - 处理 from X import Y 的模块部分
    #
    # 对相对导入（level > 0），使用 trans.CurrentPackage 解析为
    # 完整模块名（如 __types → llvmlite.__types），确保 .deps.txt
    # 记录的模块名与 SHA1 映射表一致。
    # ============================================================
    def HandleImportFromModule(self, node: ast.AST | t.CPtr) -> int:
        """处理 from X import Y 的模块部分，更新 trans._imported_modules"""
        impf: ast.ImportFrom | t.CPtr = (ast.ImportFrom | t.CPtr)(node)
        if impf is None:
            return 0
        # 解析模块名：相对导入需补充包前缀
        resolved: str = _resolve_relative_module(
            self.Trans.Pool, self.Trans.CurrentPackage,
            impf.level, impf.module)
        if resolved is not None:
            self.Trans._imported_modules = add_imported_module(
                self.Trans.Pool, self.Trans._imported_modules, resolved)
        return 0

    # ============================================================
    # HandleImportFromNames - 处理 from X import Y 的名称部分
    #
    # 同样使用解析后的完整模块名，确保 from-import 映射
    # （如 LLVMType:llvmlite.__types）能正确查到 SHA1。
    # ============================================================
    def HandleImportFromNames(self, node: ast.AST | t.CPtr) -> int:
        """处理 from X import Y 的名称部分，更新 trans._from_imports"""
        impf: ast.ImportFrom | t.CPtr = (ast.ImportFrom | t.CPtr)(node)
        if impf is None:
            return 0
        # 解析模块名：相对导入需补充包前缀
        resolved: str = _resolve_relative_module(
            self.Trans.Pool, self.Trans.CurrentPackage,
            impf.level, impf.module)
        # [IFD] 诊断 from-import 名称处理
        mod_dbg: str = "(null)"
        if impf.module is not None:
            mod_dbg = impf.module
        pkg_dbg: str = "(null)"
        if self.Trans.CurrentPackage is not None:
            pkg_dbg = self.Trans.CurrentPackage
        res_dbg: str = "(null)"
        if resolved is not None:
            res_dbg = resolved
        names_dbg: t.CSizeT = 0
        if impf.names is not None:
            names_dbg = impf.names.__len__()
        stdio.printf("[IFD] level=%d module=%s pkg=%s resolved=%s names=%d\n",
                     impf.level, mod_dbg, pkg_dbg, res_dbg, names_dbg)
        if resolved is not None:
            names: list[ast.AST | t.CPtr] | t.CPtr = impf.names
            if names is not None:
                nn: t.CSizeT = names.__len__()
                for ni in range(nn):
                    alias: ast.Alias | t.CPtr = (ast.Alias | t.CPtr)(names.get(ni))
                    if alias is not None and alias.name is not None:
                        local_name: str = alias.name
                        orig_name: str = None
                        if alias.asname is not None:
                            local_name = alias.asname
                            orig_name = alias.name
                        self.Trans._from_imports = add_from_import(
                            self.Trans.Pool, self.Trans._from_imports,
                            local_name, resolved, orig_name)
                        # 命名空间隔离：from-import 的名称标记为可见结构体
                        HandlesStruct.add_visible_struct(self.Trans.Pool, local_name)
                        # 注册 re-export 映射：当前模块 re-export 了 local_name，源模块为 resolved
                        # 使跨模块调用 module.func() 能解析到正确的源模块 SHA1
                        if self.Trans.ModuleSha1 is not None:
                            register_reexport(self.Trans.ModuleSha1, local_name, resolved)
        return 0


# ============================================================
# NewImportsHandle - 工厂函数
# ============================================================
def NewImportsHandle(pool: memhub.MemBuddy | t.CPtr,
                     trans: HT.Translator | t.CPtr) -> ImportsHandle | t.CPtr:
    h: ImportsHandle | t.CPtr = pool.alloc(ImportsHandle.__sizeof__())
    if h is None:
        return None
    string.memset(h, 0, ImportsHandle.__sizeof__())
    h.Trans = trans
    return h