import t, c
from stdint import *
import llvmlite
import memhub
import string


# ============================================================
# HandlesVar - 嵌套作用域符号表 + 变量管理
#
# 使用 Scope 链（Parent 指针）管理嵌套作用域，
# 每个作用域包含一个 VarEntry 数组（扁平化，已验证可工作）。
#
# 作用域类型:
#   SCOPE_MODULE  — 模块级（根作用域）
#   SCOPE_FUNCTION — 函数级
#   SCOPE_BLOCK   — 块级（if/for/while，当前未使用块作用域）
# ============================================================

# 作用域类型常量
SCOPE_MODULE: t.CDefine = 0
SCOPE_FUNCTION: t.CDefine = 1
SCOPE_BLOCK: t.CDefine = 2

# 每个作用域最大变量数
MAX_VARS: t.CDefine = 256


# ============================================================
# VarEntry - 变量表条目
# ============================================================
@t.NoVTable
class VarEntry:
    """变量表条目"""
    Name: t.CChar | t.CPtr
    Alloca: llvmlite.Value | t.CPtr
    Used: t.CInt
    AnnotClassName: t.CChar | t.CPtr      # 原始类型注解的类名（str 别名在结构体字段中触发编译器 bug，改用显式联合类型）


# ============================================================
# Scope - 作用域节点
#
# Parent 指向父作用域（None 表示根），
# Vars 是 VarEntry 数组，VarCount 是当前变量数。
# ============================================================
@t.NoVTable
class Scope:
    """作用域节点"""
    Parent: Scope | t.CPtr
    Vars: VarEntry | t.CPtr
    VarCount: t.CInt
    Kind: t.CInt


# ============================================================
# SymbolTable - 嵌套作用域符号表
#
# Root 是模块级根作用域，Current 是当前作用域。
# enter_scope/exit_scope 管理作用域栈。
# ============================================================
@t.NoVTable
class SymbolTable:
    """嵌套作用域符号表"""
    Pool: memhub.MemBuddy | t.CPtr
    Root: Scope | t.CPtr
    Current: Scope | t.CPtr


# ============================================================
# init_vars — 分配并清零 VarEntry 数组
# ============================================================
def init_vars(pool: memhub.MemBuddy | t.CPtr) -> VarEntry | t.CPtr:
    """分配并清零变量表数组"""
    size: t.CSizeT = MAX_VARS * VarEntry.__sizeof__()
    vars_ptr: VarEntry | t.CPtr = pool.alloc(size)
    if vars_ptr is not None:
        string.memset(vars_ptr, 0, size)
    return vars_ptr


# ============================================================
# find_var — 在变量表中按名称查找
# ============================================================
def find_var(vars_ptr: VarEntry | t.CPtr,
             var_count: int,
             name: str) -> llvmlite.Value | t.CPtr:
    """在变量表中按名称查找"""
    if name is None or vars_ptr is None:
        return None
    entry_size: t.CSizeT = VarEntry.__sizeof__()
    for i in range(var_count):
        entry_addr: t.CUInt64T = t.CUInt64T(vars_ptr) + i * entry_size
        entry: VarEntry | t.CPtr = (VarEntry | t.CPtr)(t.CVoid(entry_addr, t.CPtr))
        if entry.Name is not None and entry.Used != 0:
            if string.strcmp(entry.Name, name) == 0:
                return entry.Alloca
    return None


# ============================================================
# find_var_entry — 在变量表中按名称查找，返回 VarEntry（含 AnnotTy）
# ============================================================
def find_var_entry(vars_ptr: VarEntry | t.CPtr,
                   var_count: int,
                   name: str) -> VarEntry | t.CPtr:
    """在变量表中按名称查找，返回 VarEntry 或 None"""
    if name is None or vars_ptr is None:
        return None
    entry_size: t.CSizeT = VarEntry.__sizeof__()
    for i in range(var_count):
        entry_addr: t.CUInt64T = t.CUInt64T(vars_ptr) + i * entry_size
        entry: VarEntry | t.CPtr = (VarEntry | t.CPtr)(t.CVoid(entry_addr, t.CPtr))
        if entry.Name is not None and entry.Used != 0:
            if string.strcmp(entry.Name, name) == 0:
                return entry
    return None


# ============================================================
# lookup_var_entry — 从当前作用域逐级向上查找变量，返回 VarEntry
# ============================================================
def lookup_var_entry(symtab: SymbolTable | t.CPtr,
                     name: str) -> VarEntry | t.CPtr:
    """从当前作用域逐级向上查找变量，返回 VarEntry 或 None"""
    if symtab is None or name is None:
        return None
    scope: Scope | t.CPtr = symtab.Current
    while scope is not None:
        result: VarEntry | t.CPtr = find_var_entry(scope.Vars, scope.VarCount, name)
        if result is not None:
            return result
        scope = scope.Parent
    return None


# ============================================================
# set_var_annot_class_name — 设置变量的原始类型注解类名
#
# 在函数参数定义后调用，存储原始类型注解的类名。
# 方法调用检测时，当 alloca 类型是 Ptr(i8)（联合类型简化），
# 通过 AnnotClassName 查找实际结构体类型。
# ============================================================
def set_var_annot_class_name(symtab: SymbolTable | t.CPtr,
                             name: str,
                             class_name: str) -> int:
    """设置变量的原始类型注解类名，返回 0=成功 / 1=失败"""
    if symtab is None or name is None:
        return 1
    entry: VarEntry | t.CPtr = lookup_var_entry(symtab, name)
    if entry is None:
        return 1
    entry.AnnotClassName = class_name
    return 0


# ============================================================
# add_var — 添加变量到变量表
# ============================================================
def add_var(vars_ptr: VarEntry | t.CPtr,
            var_count: int,
            name: str,
            alloca: llvmlite.Value | t.CPtr) -> int:
    """添加变量到变量表"""
    if name is None or alloca is None or vars_ptr is None:
        return 1
    if var_count >= MAX_VARS:
        return 1
    entry_size: t.CSizeT = VarEntry.__sizeof__()
    entry_addr: t.CUInt64T = t.CUInt64T(vars_ptr) + var_count * entry_size
    entry: VarEntry | t.CPtr = (VarEntry | t.CPtr)(t.CVoid(entry_addr, t.CPtr))
    entry.Name = name
    entry.Alloca = alloca
    entry.Used = 1
    return 0


# ============================================================
# get_or_create_var — 查找或创建变量 alloca（旧版兼容）
# ============================================================
def get_or_create_var(pool: memhub.MemBuddy | t.CPtr,
                      builder: llvmlite.IRBuilder | t.CPtr,
                      vars_ptr: VarEntry | t.CPtr,
                      var_count: int,
                      name: str,
                      ty: llvmlite.LLVMType | t.CPtr) -> llvmlite.Value | t.CPtr:
    """查找或创建变量 alloca"""
    existing: llvmlite.Value | t.CPtr = find_var(vars_ptr, var_count, name)
    if existing is not None:
        return existing
    alloca: llvmlite.Value | t.CPtr = llvmlite.build_alloca(builder, ty)
    if alloca is None:
        return None
    if add_var(vars_ptr, var_count, name, alloca) != 0:
        return None
    return alloca


# ============================================================
# _create_scope — 创建新作用域节点（内部辅助函数）
# ============================================================
def _create_scope(pool: memhub.MemBuddy | t.CPtr,
                  parent: Scope | t.CPtr,
                  kind: int) -> Scope | t.CPtr:
    """创建并初始化作用域节点"""
    scope: Scope | t.CPtr = pool.alloc(Scope.__sizeof__())
    if scope is None:
        return None
    string.memset(scope, 0, Scope.__sizeof__())
    scope.Parent = parent
    scope.Vars = init_vars(pool)
    scope.VarCount = 0
    scope.Kind = kind
    if scope.Vars is None:
        return None
    return scope


# ============================================================
# init_symbol_table — 创建符号表（含模块级根作用域）
# ============================================================
def init_symbol_table(pool: memhub.MemBuddy | t.CPtr) -> SymbolTable | t.CPtr:
    """创建并初始化符号表，包含模块级根作用域"""
    if pool is None:
        return None
    symtab: SymbolTable | t.CPtr = pool.alloc(SymbolTable.__sizeof__())
    if symtab is None:
        return None
    string.memset(symtab, 0, SymbolTable.__sizeof__())
    symtab.Pool = pool

    # 创建根作用域（模块级）
    root: Scope | t.CPtr = _create_scope(pool, None, SCOPE_MODULE)
    if root is None:
        return None

    symtab.Root = root
    symtab.Current = root
    return symtab


# ============================================================
# enter_scope — 进入新作用域
# ============================================================
def enter_scope(symtab: SymbolTable | t.CPtr,
                kind: int) -> Scope | t.CPtr:
    """进入新作用域，返回新创建的作用域"""
    if symtab is None:
        return None
    pool: memhub.MemBuddy | t.CPtr = symtab.Pool
    scope: Scope | t.CPtr = _create_scope(pool, symtab.Current, kind)
    if scope is None:
        return None
    symtab.Current = scope
    return scope


# ============================================================
# exit_scope — 退出当前作用域
# ============================================================
def exit_scope(symtab: SymbolTable | t.CPtr):
    """退出当前作用域，恢复到父作用域"""
    if symtab is not None and symtab.Current is not None:
        symtab.Current = symtab.Current.Parent


# ============================================================
# define_var — 在当前作用域定义变量
# ============================================================
def define_var(symtab: SymbolTable | t.CPtr,
               name: str,
               alloca: llvmlite.Value | t.CPtr) -> int:
    """在当前作用域定义变量，返回 0 成功"""
    if symtab is None or name is None or alloca is None:
        return 1
    scope: Scope | t.CPtr = symtab.Current
    if scope is None:
        return 1
    ret: int = add_var(scope.Vars, scope.VarCount, name, alloca)
    if ret == 0:
        scope.VarCount = scope.VarCount + 1
    return ret


# ============================================================
# define_module_var — 在模块作用域定义变量
# ============================================================
def define_module_var(symtab: SymbolTable | t.CPtr,
                      name: str,
                      alloca: llvmlite.Value | t.CPtr) -> int:
    """在模块作用域定义变量，返回 0 成功"""
    if symtab is None or name is None or alloca is None:
        return 1
    scope: Scope | t.CPtr = symtab.Root
    if scope is None:
        return 1
    ret: int = add_var(scope.Vars, scope.VarCount, name, alloca)
    if ret == 0:
        scope.VarCount = scope.VarCount + 1
    return ret


# ============================================================
# lookup_var — 从当前作用域逐级向上查找变量
# ============================================================
def lookup_var(symtab: SymbolTable | t.CPtr,
               name: str) -> llvmlite.Value | t.CPtr:
    """从当前作用域逐级向上查找变量，返回 alloca 或 None"""
    if symtab is None or name is None:
        return None
    scope: Scope | t.CPtr = symtab.Current
    while scope is not None:
        result: llvmlite.Value | t.CPtr = find_var(
            scope.Vars, scope.VarCount, name)
        if result is not None:
            return result
        scope = scope.Parent
    return None


# ============================================================
# lookup_current — 仅在当前作用域查找变量
# ============================================================
def lookup_current(symtab: SymbolTable | t.CPtr,
                   name: str) -> llvmlite.Value | t.CPtr:
    """仅在当前作用域查找变量"""
    if symtab is None or name is None:
        return None
    scope: Scope | t.CPtr = symtab.Current
    if scope is None:
        return None
    return find_var(scope.Vars, scope.VarCount, name)


# ============================================================
# lookup_module_var — 在模块作用域查找变量
# ============================================================
def lookup_module_var(symtab: SymbolTable | t.CPtr,
                      name: str) -> llvmlite.Value | t.CPtr:
    """在模块作用域查找变量"""
    if symtab is None or name is None:
        return None
    scope: Scope | t.CPtr = symtab.Root
    if scope is None:
        return None
    return find_var(scope.Vars, scope.VarCount, name)


# ============================================================
# get_or_create_sym — 查找或创建变量（在当前作用域）
# ============================================================
def get_or_create_sym(symtab: SymbolTable | t.CPtr,
                      pool: memhub.MemBuddy | t.CPtr,
                      builder: llvmlite.IRBuilder | t.CPtr,
                      name: str,
                      ty: llvmlite.LLVMType | t.CPtr) -> llvmlite.Value | t.CPtr:
    """查找或创建变量 alloca（在当前作用域）"""
    existing: llvmlite.Value | t.CPtr = lookup_current(symtab, name)
    if existing is not None:
        return existing
    alloca: llvmlite.Value | t.CPtr = llvmlite.build_alloca(builder, ty)
    if alloca is None:
        return None
    define_var(symtab, name, alloca)
    return alloca