import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import stdlib
import lib.core.Handles.HandlesVar as HandlesVar
import lib.core.Handles.HandlesExprCall as HandlesExprCall
import lib.core.Handles.HandlesMain as HandlesMain
import lib.core.Handles.HandlesAssign as HandlesAssign
import lib.core.Handles.HandlesReturn as HandlesReturn
import lib.core.Handles.HandlesImports as HandlesImports
import lib.core.Handles.HandlesAnnAssign as HandlesAnnAssign
import lib.core.Handles.HandlesAugAssign as HandlesAugAssign
import lib.core.Handles.HandlesIf as HandlesIf
import lib.core.Handles.HandlesWhile as HandlesWhile
import lib.core.Handles.HandlesFor as HandlesFor
import lib.core.Handles.HandlesExpr as HandlesExpr
import lib.core.Handles.HandlesFunctions as HandlesFunctions
import lib.Projectrans.Config as Config

# ============================================================
# HandlesTranslator - 翻译器状态管理 + 主入口
#
# 从 translator.py 拆分出来，负责:
#   1. Translator 类 - 状态管理（Module, 变量表, 函数表, 导入）
#   2. translate() - 主翻译入口
#   3. dump_ir() - IR 输出
#
# 语句/表达式翻译全部委托给 Handles 模块:
#   HandlesMain    - wrapper main 创建 + children 遍历
#   HandlesBody    - 语句分派
#   HandlesExpr    - 表达式翻译
#   HandlesFunctions - 函数定义处理
# 注意：str = bytes = t.CChar | t.CPtr = i8*
# ============================================================

# 常量
MAX_VARS: t.CDefine = 256
MAX_FUNCS: t.CDefine = 256
MAX_GLOBAL_NAMES: t.CDefine = 32


@t.NoVTable
class Translator:
    """翻译器：管理编译状态，委托 Handles 模块执行翻译"""

    # LLVM 上下文（共享状态）
    Module: llvmlite.LLVMModule | t.CPtr
    Pool: memhub.MemBuddy | t.CPtr

    # 函数表
    _funcs: HandlesExprCall.FuncEntry | t.CPtr
    _func_count: t.CInt

    # 导入模块跟踪
    _imported_modules: str
    _from_imports: str

    # 当前翻译上下文
    _cur_func: llvmlite.Function | t.CPtr
    _cur_builder: llvmlite.IRBuilder | t.CPtr

    # 循环控制流目标（break/continue）
    _break_bb: llvmlite.BasicBlock | t.CPtr
    _continue_bb: llvmlite.BasicBlock | t.CPtr

    # 独立标签计数器（不与 builder.Counter 共享，避免 SSA 编号非单调）
    _label_counter: t.CInt

    # global 声明的变量名集合（当前函数内）
    _global_names: str
    _global_name_count: t.CInt

    # nonlocal 声明的变量名集合（当前函数内）
    _nonlocal_names: str
    _nonlocal_name_count: t.CInt

    # 闭包相关：当前函数名（用于嵌套函数提升命名）
    _cur_func_name: str

    # 模块 SHA1 前缀（16 字符），用于函数名混淆（SHA1 命名空间）
    ModuleSha1: str
    # 当前文件所属包名（如 "llvmlite"），用于解析相对导入（from . import ...）
    # None 表示当前文件是顶级模块（无包），相对导入无法解析
    CurrentPackage: str
    # 闭包 env 中的 nonlocal 变量偏移映射（每个 nonlocal 变量在 env 中的字节偏移）
    _closure_env_offsets: t.CInt
    _closure_env_count: t.CInt

    # 泛型特化上下文：当前正在特化的类型参数名/实参名列表
    # 由 _specialize_generic_class 设置，方法体翻译时用于将 T 替换为具体类型
    # None 表示不在泛型特化上下文中
    GenericTypeParamNames: list[str] | t.CPtr
    GenericTypeArgs: list[str] | t.CPtr

    # Phase 1a 声明模式标志：1=只注册 struct/enum/union 不翻译代码体，0=全量翻译
    _declare_only: t.CInt

    # 子 Handle 指针（每个 Handle 一个槽，通过 Mixin 回指针访问本结构体）
    AssignH: HandlesAssign.AssignHandle | t.CPtr
    ReturnH: HandlesReturn.ReturnHandle | t.CPtr
    ImportsH: HandlesImports.ImportsHandle | t.CPtr
    AnnAssignH: HandlesAnnAssign.AnnAssignHandle | t.CPtr
    AugAssignH: HandlesAugAssign.AugAssignHandle | t.CPtr
    IfH: HandlesIf.IfHandle | t.CPtr
    WhileH: HandlesWhile.WhileHandle | t.CPtr
    ForH: HandlesFor.ForHandle | t.CPtr
    ExprH: HandlesExpr.ExprHandle | t.CPtr
    ExprCallH: HandlesExprCall.ExprCallHandle | t.CPtr

    # 嵌套作用域符号表
    SymTab: HandlesVar.SymbolTable | t.CPtr

    def __init__(self):
        self.Module = None
        self.Pool = None
        self._funcs = None
        self._func_count = 0
        self._cur_func = None
        self._cur_builder = None
        self._break_bb = None
        self._continue_bb = None
        self._label_counter = 0
        self._imported_modules = None
        self._from_imports = None
        self._global_names = None
        self._global_name_count = 0
        self._nonlocal_names = None
        self._nonlocal_name_count = 0
        self._cur_func_name = None
        self.ModuleSha1 = None
        self.CurrentPackage = None
        self._closure_env_offsets = 0
        self._closure_env_count = 0
        self.GenericTypeParamNames = None
        self.GenericTypeArgs = None
        self._declare_only = 0
        self.AssignH = None
        self.ReturnH = None
        self.ImportsH = None
        self.AnnAssignH = None
        self.AugAssignH = None
        self.IfH = None
        self.WhileH = None
        self.ForH = None
        self.ExprH = None
        self.ExprCallH = None
        self.SymTab = None

    # ============================================================
    # 状态初始化
    # ============================================================
    def _init_state(self, pool: memhub.MemBuddy | t.CPtr):
        """初始化变量表、函数表和子 Handle"""
        self.Pool = pool
        if self._funcs is None:
            self._funcs = HandlesExprCall.init_func_table(pool, MAX_FUNCS)
        # 嵌套作用域符号表（新版）
        if self.SymTab is None:
            self.SymTab = HandlesVar.init_symbol_table(pool)
        # global/nonlocal 名称集合缓冲区（32 个 char* 指针 = 256 字节）
        if self._global_names is None:
            self._global_names = stdlib.malloc(MAX_GLOBAL_NAMES * 8)
            if self._global_names is not None:
                string.memset(self._global_names, 0, MAX_GLOBAL_NAMES * 8)
            self._global_name_count = 0
        if self._nonlocal_names is None:
            self._nonlocal_names = stdlib.malloc(MAX_GLOBAL_NAMES * 8)
            if self._nonlocal_names is not None:
                string.memset(self._nonlocal_names, 0, MAX_GLOBAL_NAMES * 8)
            self._nonlocal_name_count = 0
        # 创建子 Handle（传入 self 作为 Mixin 回指针）
        if self.AssignH is None:
            self.AssignH = HandlesAssign.NewAssignHandle(pool, self)
        if self.ReturnH is None:
            self.ReturnH = HandlesReturn.NewReturnHandle(pool, self)
        if self.ImportsH is None:
            self.ImportsH = HandlesImports.NewImportsHandle(pool, self)
        if self.AnnAssignH is None:
            self.AnnAssignH = HandlesAnnAssign.NewAnnAssignHandle(pool, self)
        if self.AugAssignH is None:
            self.AugAssignH = HandlesAugAssign.NewAugAssignHandle(pool, self)
        if self.IfH is None:
            self.IfH = HandlesIf.NewIfHandle(pool, self)
        if self.WhileH is None:
            self.WhileH = HandlesWhile.NewWhileHandle(pool, self)
        if self.ForH is None:
            self.ForH = HandlesFor.NewForHandle(pool, self)
        if self.ExprH is None:
            self.ExprH = HandlesExpr.NewExprHandle(pool, self)
        if self.ExprCallH is None:
            self.ExprCallH = HandlesExprCall.NewExprCallHandle(pool, self)

    # ============================================================
    # 主翻译入口
    # ============================================================
    def translate(self, tree: ast.AST | t.CPtr) -> int:
        """将 AST 翻译为 LLVM IR"""
        if tree is None:
            return 1
        if _mbuddy is None:
            return 1

        pool: memhub.MemBuddy | t.CPtr = _mbuddy

        # 初始化状态
        self._init_state(pool)

        # 创建 LLVM 模块
        mod: llvmlite.LLVMModule | t.CPtr = llvmlite.new_module(pool, "main")
        if mod is None:
            stdio.printf("[TR] NewModule returned NULL\n")
            return 1
        self.Module = mod

        # 设置目标（优先使用 project.vpj 中的配置）
        triple: str = Config.TargetTriple
        if triple is None:
            triple = "x86_64-pc-windows-msvc"
        llvmlite.module_set_target(mod, triple)

        # 类型
        i32_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
        i8_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int8(pool)
        i8_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, i8_ty)

        # 声明 printf
        printf_func: llvmlite.Function | t.CPtr = llvmlite.create_declare(
            pool, mod, "printf", i32_ty)
        if printf_func is None:
            stdio.printf("[TR] CreateDeclare printf returned NULL\n")
            return 1
        llvmlite.add_param(pool, printf_func, i8_ptr_ty, "fmt")
        printf_func.IsVarArg = 1

        # 声明 malloc（闭包分配用）
        malloc_func: llvmlite.Function | t.CPtr = llvmlite.create_declare(
            pool, mod, "malloc", i8_ptr_ty)
        if malloc_func is not None:
            i64_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int64(pool)
            llvmlite.add_param(pool, malloc_func, i64_ty, "size")

        # 检查用户是否定义了 main 函数
        has_user_main: int = 0
        ch: list[ast.AST | t.CPtr] | t.CPtr = tree.children
        if ch is not None:
            cn_count: t.CSizeT = ch.__len__()
            for ci in range(cn_count):
                child: ast.AST | t.CPtr = ch.get(ci)
                if child is not None and child.kind() == ast.ASTKind.FunctionDef:
                    fd: ast.FunctionDef | t.CPtr = (ast.FunctionDef | t.CPtr)(child)
                    if fd is not None and fd.name is not None:
                        if string.strcmp(fd.name, "main") == 0:
                            has_user_main = 1
                            break

        if self._declare_only != 0:
            # Phase 1a-pre(2=import扫描) / Phase 1a(1=struct注册): 只处理模块级，不创建 main 函数和 builder
            self._translate_module_level(pool, mod, tree)
        elif has_user_main == 0:
            # 无用户 main → 只翻译模块级语句，不创建 wrapper main
            # 修复：避免每个文件都生成 define i32 @main() 导致链接时 main 冲突
            # 只有包含 def main() 的入口文件才会有 define i32 @main()
            self._translate_module_level(pool, mod, tree)
        else:
            # 用户已定义 main → 委托 HandlesMain 翻译模块级
            self._translate_module_level(pool, mod, tree)

        return 0

    # ============================================================
    # 包装 main 函数（无用户 main 时）→ 委托 HandlesMain
    # ============================================================
    def _translate_wrapper_main(self, pool: memhub.MemBuddy | t.CPtr,
                                mod: llvmlite.LLVMModule | t.CPtr,
                                tree: ast.AST | t.CPtr,
                                i32_ty: llvmlite.LLVMType | t.CPtr):
        """委托 HandlesMain.create_wrapper_main() 创建包装 main（trans 单参）"""
        added: int = HandlesMain.create_wrapper_main(self, tree, i32_ty)

    # ============================================================
    # 模块级翻译（用户已定义 main）→ 委托 HandlesMain
    # ============================================================
    def _translate_module_level(self, pool: memhub.MemBuddy | t.CPtr,
                                mod: llvmlite.LLVMModule | t.CPtr,
                                tree: ast.AST | t.CPtr):
        """委托 HandlesMain.translate_children() 翻译模块级语句（trans 单参）"""
        # 全量翻译模式下，先处理导入语句，再创建前向声明，解决前向引用问题
        if self._declare_only == 0:
            # 预处理导入语句，确保 _imported_modules 和 _from_imports 已填充
            # （前向声明需要解析类型注解，如 t.CArray[str] 依赖 t 模块已导入）
            ch: list[ast.AST | t.CPtr] | t.CPtr = tree.children
            if ch is not None:
                cn_count: t.CSizeT = ch.__len__()
                for ci in range(cn_count):
                    child: ast.AST | t.CPtr = ch.get(ci)
                    if child is None:
                        continue
                    kd: int = child.kind()
                    if kd == ast.ASTKind.Import:
                        self.ImportsH.HandleImport(child)
                    elif kd == ast.ASTKind.ImportFrom:
                        self.ImportsH.HandleImportFromModule(child)
                        self.ImportsH.HandleImportFromNames(child)
            # 创建前向声明
            HandlesFunctions.forward_declare_functions(self, tree)
        added: int = HandlesMain.translate_children(self, tree)

    # ============================================================
    # IR 输出
    # ============================================================
    def dump_ir(self, buf: bytes, size: t.CSizeT, mode: int):
        """将 LLVM IR 输出到缓冲区

        mode: 0=完整, 1=stub(仅声明), 2=text(仅代码)
        """
        if self.Module is None or buf is None:
            return
        if _mbuddy is None:
            return
        buf[0] = '\0'
        llvmlite.LLVMModulePrint(buf, size, self.Module, _mbuddy, mode)


# ============================================================
# global/nonlocal 名称管理（模块级辅助函数）
# ============================================================
def add_global_name(trans: HT.Translator | t.CPtr, name: str):
    """添加一个 global 变量名"""
    if trans is None or name is None:
        return
    if trans._global_name_count >= MAX_GLOBAL_NAMES:
        return
    if trans._global_names is None:
        return
    # 检查是否已存在
    if is_global_name(trans, name) != 0:
        return
    entry_addr: t.CUInt64T = t.CUInt64T(trans._global_names) + trans._global_name_count * 8
    slot_ptr: str | t.CPtr = (t.CVoid | t.CPtr)(entry_addr)
    slot_ptr[0] = name
    trans._global_name_count += 1


def is_global_name(trans: HT.Translator | t.CPtr, name: str) -> int:
    """检查 name 是否在当前函数的 global 集合中"""
    if trans is None or name is None or trans._global_names is None:
        return 0
    for i in range(trans._global_name_count):
        entry_addr: t.CUInt64T = t.CUInt64T(trans._global_names) + i * 8
        slot_ptr: str | t.CPtr = (t.CVoid | t.CPtr)(entry_addr)
        entry: str = slot_ptr[0]
        if entry is not None:
            if string.strcmp(entry, name) == 0:
                return 1
    return 0


def add_nonlocal_name(trans: HT.Translator | t.CPtr, name: str):
    """添加一个 nonlocal 变量名"""
    if trans is None or name is None:
        return
    if trans._nonlocal_name_count >= MAX_GLOBAL_NAMES:
        return
    if trans._nonlocal_names is None:
        return
    if is_nonlocal_name(trans, name) != 0:
        return
    entry_addr: t.CUInt64T = t.CUInt64T(trans._nonlocal_names) + trans._nonlocal_name_count * 8
    slot_ptr: str | t.CPtr = (t.CVoid | t.CPtr)(entry_addr)
    slot_ptr[0] = name
    trans._nonlocal_name_count += 1
    trans._closure_env_count += 1


def is_nonlocal_name(trans: HT.Translator | t.CPtr, name: str) -> int:
    """检查 name 是否在当前函数的 nonlocal 集合中"""
    if trans is None or name is None or trans._nonlocal_names is None:
        return 0
    for i in range(trans._nonlocal_name_count):
        entry_addr: t.CUInt64T = t.CUInt64T(trans._nonlocal_names) + i * 8
        slot_ptr: str | t.CPtr = (t.CVoid | t.CPtr)(entry_addr)
        entry: str = slot_ptr[0]
        if entry is not None:
            if string.strcmp(entry, name) == 0:
                return 1
    return 0


def get_nonlocal_index(trans: HT.Translator | t.CPtr, name: str) -> int:
    """获取 nonlocal 变量在 env 中的索引（0-based），找不到返回 -1"""
    if trans is None or name is None or trans._nonlocal_names is None:
        return -1
    for i in range(trans._nonlocal_name_count):
        entry_addr: t.CUInt64T = t.CUInt64T(trans._nonlocal_names) + i * 8
        slot_ptr: str | t.CPtr = (t.CVoid | t.CPtr)(entry_addr)
        entry: str = slot_ptr[0]
        if entry is not None:
            if string.strcmp(entry, name) == 0:
                return i
    return -1


def clear_scope_names(trans: HT.Translator | t.CPtr):
    """清空当前函数的 global/nonlocal 名称集合（进入新函数时调用）"""
    if trans is None:
        return
    trans._global_name_count = 0
    trans._nonlocal_name_count = 0
    trans._closure_env_count = 0


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