import t, c
from stdint import *
import ast
import llvmlite
import memhub
import string
import stdio
import lib.core.Handles.HandlesTranslator as HT
import lib.core.Handles.HandlesVar as HandlesVar


# ============================================================
# HandlesNonlocal - nonlocal 变量访问（通过闭包 env）
#
# 闭包 env 结构: {i8* ptr0, i8* ptr1, ...} (每个 nonlocal 变量一个 i8* 指针)
# 指针指向原始变量（可能是 alloca 或全局变量）
#
# 在提升的嵌套函数中:
#   1. 函数签名: define i32 @__closure_{name}(i8* %env)
#   2. 入口处: %env_alloca = alloca i8*; store i8* %env, i8** %env_alloca
#   3. 访问 nonlocal var:
#      a. %env_ptr = load i8*, i8** %env_alloca
#      b. %addr = gep i8, i8* %env_ptr, i32 (index * 8)
#      c. %ptr_addr = bitcast i8* %addr to i8**
#      d. %var_ptr_raw = load i8*, i8** %ptr_addr
#      e. %var_ptr = bitcast i8* %var_ptr_raw to i32*
#      f. read: %val = load i32, i32* %var_ptr
#      write: store i32 %new, i32* %var_ptr
# ============================================================


# ============================================================
# 获取或创建 _env_ptr 变量（存储 env 参数的 alloca）
# ============================================================
def get_env_ptr_var(trans: HT.Translator | t.CPtr) -> llvmlite.Value | t.CPtr:
    """获取或创建 _env_ptr 变量（存储闭包 env 指针）"""
    env_var: llvmlite.Value | t.CPtr = HandlesVar.lookup_var(
        trans.SymTab, "_env_ptr")
    if env_var is not None:
        return env_var
    # 创建 alloca 存储 env
    pool: memhub.MemBuddy | t.CPtr = trans.Pool
    builder: llvmlite.IRBuilder | t.CPtr = trans._cur_builder
    i8_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int8(pool)
    i8_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, i8_ty)
    alloca: llvmlite.Value | t.CPtr = llvmlite.build_alloca(builder, i8_ptr_ty)
    if alloca is None:
        return None
    HandlesVar.define_var(trans.SymTab, "_env_ptr", alloca)
    return alloca


# ============================================================
# 加载 nonlocal 变量值 → 返回 i32 Value
#
# env 直接存储 i32 值（不是指针），每个 nonlocal 变量占 4 字节
# ============================================================
def load_nonlocal_var(trans: HT.Translator | t.CPtr,
                      name: str) -> llvmlite.Value | t.CPtr:
    """从闭包 env 加载 nonlocal 变量值（env 直接存储 i32 值）"""
    idx: int = HT.get_nonlocal_index(trans, name)
    if idx < 0:
        return None

    pool: memhub.MemBuddy | t.CPtr = trans.Pool
    builder: llvmlite.IRBuilder | t.CPtr = trans._cur_builder
    i8_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int8(pool)
    i8_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, i8_ty)
    i32_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
    i32_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, i32_ty)

    # 1. 加载 env_ptr
    env_alloca: llvmlite.Value | t.CPtr = get_env_ptr_var(trans)
    if env_alloca is None:
        return None
    env_ptr: llvmlite.Value | t.CPtr = llvmlite.build_load(builder, i8_ptr_ty, env_alloca)
    if env_ptr is None:
        return None

    # 2. GEP to offset (idx * 4) — env 直接存储 i32 值
    offset_val: llvmlite.Value | t.CPtr = llvmlite.const_int32(pool, idx * 4)
    addr: llvmlite.Value | t.CPtr = llvmlite.build_gep(builder, i8_ty, env_ptr, offset_val)
    if addr is None:
        return None

    # 3. Bitcast to i32* and load
    var_ptr: llvmlite.Value | t.CPtr = llvmlite.build_bitcast(builder, addr, i32_ptr_ty)
    if var_ptr is None:
        return None
    return llvmlite.build_load(builder, i32_ty, var_ptr)


# ============================================================
# 获取 nonlocal 变量指针（用于 store）
#
# 返回 env 中 i32 槽位的地址（i32*），用于直接 store
# ============================================================
def get_nonlocal_var_ptr(trans: HT.Translator | t.CPtr,
                         name: str) -> llvmlite.Value | t.CPtr:
    """获取 nonlocal 变量在 env 中的地址（i32*），用于 store 操作"""
    idx: int = HT.get_nonlocal_index(trans, name)
    if idx < 0:
        return None

    pool: memhub.MemBuddy | t.CPtr = trans.Pool
    builder: llvmlite.IRBuilder | t.CPtr = trans._cur_builder
    i8_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int8(pool)
    i8_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, i8_ty)
    i32_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool)
    i32_ptr_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Ptr(pool, i32_ty)

    # 1. 加载 env_ptr
    env_alloca: llvmlite.Value | t.CPtr = get_env_ptr_var(trans)
    if env_alloca is None:
        return None
    env_ptr: llvmlite.Value | t.CPtr = llvmlite.build_load(builder, i8_ptr_ty, env_alloca)
    if env_ptr is None:
        return None

    # 2. GEP to offset (idx * 4) — env 直接存储 i32 值
    offset_val: llvmlite.Value | t.CPtr = llvmlite.const_int32(pool, idx * 4)
    addr: llvmlite.Value | t.CPtr = llvmlite.build_gep(builder, i8_ty, env_ptr, offset_val)
    if addr is None:
        return None

    # 3. Bitcast to i32* and return (直接指向 env 中的 i32 槽位)
    return llvmlite.build_bitcast(builder, addr, i32_ptr_ty)