from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from lib.core.translator import Translator
from lib.core.Handles.HandlesBase import BaseHandle, CTypeInfo
import ast
import llvmlite.ir as ir


class ReturnHandle(BaseHandle):
    def _find_active_finally(self, Gen: "Translator.LlvmGen") -> tuple[ir.Block, ir.AllocaInstr, ir.AllocaInstr | None] | None:
        for item in reversed(Gen.eh_finally_stack):
            if item is not None and item[0] is not None:
                return item
        return None

    def _HandleReturnLlvm(self, Node: ast.Return) -> None:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        # 闭包返回检测：当返回值是带 nonlocal 参数的嵌套函数时，创建闭包结构体
        # 闭包结构体 {i8* env, i8* fn} 在堆上分配，env 指向 nonlocal 变量的堆存储
        if Node.value and isinstance(Node.value, ast.Name):
            ret_name: str = Node.value.id
            if ret_name in Gen.nonlocal_params and ret_name in Gen.functions:
                nested_fn: ir.Function = Gen.functions[ret_name]
                nonlocal_params: list[tuple[str, ir.Type]] = Gen.nonlocal_params[ret_name]
                fn_return_type: ir.Type = nested_fn.function_type.return_type
                nonlocal_types: list[ir.Type] = [t for _, t in nonlocal_params]
                # 构建 env：收集所有 nonlocal 变量指针
                env_ptrs: list[ir.Value] = []
                for var_name, _ in nonlocal_params:
                    vp: ir.Value | None = Gen.variables.get(var_name)
                    if vp is not None:
                        env_ptrs.append(vp)
                if env_ptrs:
                    # env 值：单个指针直接用，多个指针打包成结构体
                    i8_ptr_type: ir.PointerType = ir.IntType(8).as_pointer()
                    if len(env_ptrs) == 1:
                        env_val: ir.Value = Gen.builder.bitcast(env_ptrs[0], i8_ptr_type, name="closure_env")
                    else:
                        env_struct_type: ir.LiteralStructType = ir.LiteralStructType([p.type for p in env_ptrs])
                        env_size: int = len(env_ptrs) * 8
                        malloc_type: ir.FunctionType = ir.FunctionType(i8_ptr_type, [ir.IntType(64)])
                        malloc_func: ir.Function = Gen._get_or_declare_func('malloc', malloc_type)
                        env_raw: ir.Value = Gen.builder.call(malloc_func, [ir.Constant(ir.IntType(64), env_size)], name="env_struct_raw")
                        env_struct_ptr: ir.Value = Gen.builder.bitcast(env_raw, ir.PointerType(env_struct_type), name="env_struct_ptr")
                        for i, ptr in enumerate(env_ptrs):
                            gep: ir.Value = Gen.builder.gep(env_struct_ptr, [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), i)], name=f"env_field_{i}")
                            Gen._store(ptr, gep)
                        env_val = Gen.builder.bitcast(env_struct_ptr, i8_ptr_type, name="closure_env")
                    # fn 指针
                    fn_ptr: ir.Value = Gen.builder.bitcast(nested_fn, i8_ptr_type, name="closure_fn")
                    # 创建闭包结构体 {i8* env, i8* fn}
                    closure_type: ir.LiteralStructType = ir.LiteralStructType([i8_ptr_type, i8_ptr_type])
                    malloc_type2: ir.FunctionType = ir.FunctionType(i8_ptr_type, [ir.IntType(64)])
                    malloc_func2: ir.Function = Gen._get_or_declare_func('malloc', malloc_type2)
                    closure_raw: ir.Value = Gen.builder.call(malloc_func2, [ir.Constant(ir.IntType(64), 16)], name="closure_raw")
                    closure_ptr: ir.Value = Gen.builder.bitcast(closure_raw, ir.PointerType(closure_type), name="closure_ptr")
                    # 填充闭包结构体
                    env_gep: ir.Value = Gen.builder.gep(closure_ptr, [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 0)], name="closure_env_slot")
                    Gen._store(env_val, env_gep)
                    fn_gep: ir.Value = Gen.builder.gep(closure_ptr, [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 1)], name="closure_fn_slot")
                    Gen._store(fn_ptr, fn_gep)
                    # 记录闭包返回类型
                    current_name: str = Gen.func.name if Gen.func else ''
                    short_name: str = current_name.split('.', 1)[1] if '.' in current_name else current_name
                    Gen._closure_return_types[short_name] = (fn_return_type, nonlocal_types)
                    # 返回闭包结构体指针（bitcast 到 i8*）
                    closure_ret: ir.Value = Gen.builder.bitcast(closure_ptr, i8_ptr_type, name="closure_ret")
                    Gen.emit_return(closure_ret)
                    return
        active_finally: tuple[ir.Block, ir.AllocaInstr, ir.AllocaInstr | None] | None = self._find_active_finally(Gen)
        if getattr(self.Trans, 'CurrentCReturnTypes', None) and self.Trans.CurrentCReturnTypes and Node.value:
            return_values: list[ast.expr] = []
            if isinstance(Node.value, ast.Tuple):
                return_values = Node.value.elts
            else:
                return_values = [Node.value]
            if len(return_values) == len(self.Trans.CurrentCReturnTypes):
                llvm_vals: list[ir.Value] = []
                expected_types: list[ir.Type] = []
                for i, val in enumerate(return_values):
                    ReturnTypeInfo: CTypeInfo | None = CTypeInfo.FromNode(self.Trans.CurrentCReturnTypes[i], self.Trans.SymbolTable)
                    IsRetPtr: bool = ReturnTypeInfo.IsPtr if ReturnTypeInfo else False
                    if (isinstance(ReturnTypeInfo, CTypeInfo) and ReturnTypeInfo.IsStr) or (isinstance(self.Trans.CurrentCReturnTypes[i], ast.Name) and self.Trans.CurrentCReturnTypes[i].id == 'str'):
                        IsRetPtr = True
                    ExpectedType: ir.Type = ReturnTypeInfo.ToLLVM(Gen) if ReturnTypeInfo and ReturnTypeInfo.BaseType else Gen._CType2LLVM('int', IsRetPtr)
                    expected_types.append(ExpectedType)
                    Val: ir.Value | None = self.HandleExprLlvm(val)
                    if Val is None:
                        if isinstance(ExpectedType, ir.PointerType):
                            Val = ir.Constant(ExpectedType, None)
                        elif isinstance(ExpectedType, ir.IntType):
                            Val = ir.Constant(ExpectedType, 0)
                        else:
                            Val = ir.Constant(ExpectedType, ir.Undefined)
                    else:
                        if Val.type != ExpectedType:
                            if isinstance(ExpectedType, ir.IntType) and isinstance(Val.type, ir.IntType):
                                if Val.type.width < ExpectedType.width:
                                    Val = Gen.builder.zext(Val, ExpectedType, name=f"zext_ret_{i}")
                                elif Val.type.width > ExpectedType.width:
                                    Val = Gen.builder.trunc(Val, ExpectedType, name=f"trunc_ret_{i}")
                            elif isinstance(ExpectedType, ir.PointerType) and isinstance(Val.type, ir.PointerType):
                                Val = Gen.builder.bitcast(Val, ExpectedType, name=f"cast_ret_{i}")
                    llvm_vals.append(Val)
                if llvm_vals:
                    struct_type: ir.LiteralStructType = ir.LiteralStructType(expected_types)
                    result: ir.Value = ir.Constant(struct_type, ir.Undefined)
                    for i, val in enumerate(llvm_vals):
                        result = Gen.builder.insert_value(result, val, i, name=f"ret_field_{i}")
                    if active_finally:
                        FinallyBB: ir.Block = active_finally[0]
                        pending_ret_flag: ir.AllocaInstr = active_finally[1]
                        pending_ret_val: ir.AllocaInstr | None = active_finally[2]
                        if pending_ret_val:
                            Gen._store(result, pending_ret_val)
                        Gen._store(ir.Constant(ir.IntType(32), 1), pending_ret_flag)
                        Gen.builder.branch(FinallyBB)
                    else:
                        Gen.emit_return(result)
                return
        if Node.value:
            Val: ir.Value | None = self.HandleExprLlvm(Node.value)
            if Val:
                ret_type: ir.Type | None = getattr(Gen.func, 'ftype', None) and Gen.func.ftype.return_type if getattr(Gen, 'func', None) else None
                if isinstance(ret_type, ir.IdentifiedStructType) and isinstance(Val, ir.Constant) and isinstance(Val.type, ir.PointerType) and Val.constant is None:
                    if ret_type.elements:
                        zero_val: ir.Constant = ir.Constant(ret_type, [
                            ir.Constant(et, None) if isinstance(et, (ir.PointerType, ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType))
                            else ir.Constant(et, 0) if isinstance(et, (ir.IntType, ir.FloatType, ir.DoubleType))
                            else ir.Constant(et, ir.Undefined) for et in ret_type.elements])
                    else:
                        zero_val = ir.Constant(ret_type, None)
                    if active_finally:
                        FinallyBB: ir.Block = active_finally[0]
                        pending_ret_flag: ir.AllocaInstr = active_finally[1]
                        pending_ret_val: ir.AllocaInstr | None = active_finally[2]
                        if pending_ret_val:
                            Gen._store(zero_val, pending_ret_val)
                        Gen._store(ir.Constant(ir.IntType(32), 1), pending_ret_flag)
                        Gen.builder.branch(FinallyBB)
                    else:
                        Gen.emit_return(zero_val)
                    return
                if getattr(self.Trans, 'CurrentCReturnTypes', None) and self.Trans.CurrentCReturnTypes:
                    struct_type: ir.Type = Gen.func.ftype.return_type
                    if isinstance(struct_type, ir.LiteralStructType):
                        result: ir.Value = ir.Constant(struct_type, ir.Undefined)
                        if isinstance(Val.type, ir.PointerType) and len(struct_type.elements) > 0 and isinstance(struct_type.elements[0], ir.PointerType):
                            result = Gen.builder.insert_value(result, Val, 0, name="ret_field_0")
                            for j in range(1, len(struct_type.elements)):
                                elem_type: ir.Type = struct_type.elements[j]
                                zero_val: ir.Constant = ir.Constant(elem_type, 0 if isinstance(elem_type, ir.IntType) else None)
                                result = Gen.builder.insert_value(result, zero_val, j, name=f"ret_field_{j}")
                        else:
                            for j, elem_type in enumerate(struct_type.elements):
                                if j == 0:
                                    if Val.type == elem_type:
                                        result = Gen.builder.insert_value(result, Val, j, name=f"ret_field_{j}")
                                    elif isinstance(elem_type, ir.IntType) and isinstance(Val.type, ir.IntType):
                                        if Val.type.width < elem_type.width:
                                            Val = Gen.builder.zext(Val, elem_type, name=f"zext_ret_{j}")
                                        elif Val.type.width > elem_type.width:
                                            Val = Gen.builder.trunc(Val, elem_type, name=f"trunc_ret_{j}")
                                        result = Gen.builder.insert_value(result, Val, j, name=f"ret_field_{j}")
                                    else:
                                        zero_val = ir.Constant(elem_type, 0 if isinstance(elem_type, ir.IntType) else None)
                                        result = Gen.builder.insert_value(result, zero_val, j, name=f"ret_field_{j}")
                                else:
                                    zero_val = ir.Constant(elem_type, 0 if isinstance(elem_type, ir.IntType) else None)
                                    result = Gen.builder.insert_value(result, zero_val, j, name=f"ret_field_{j}")
                        if active_finally:
                            FinallyBB: ir.Block = active_finally[0]
                            pending_ret_flag: ir.AllocaInstr = active_finally[1]
                            pending_ret_val: ir.AllocaInstr | None = active_finally[2]
                            if pending_ret_val:
                                Gen._store(result, pending_ret_val)
                            Gen._store(ir.Constant(ir.IntType(32), 1), pending_ret_flag)
                            Gen.builder.branch(FinallyBB)
                        else:
                            Gen.emit_return(result)
                        return
                if active_finally:
                    FinallyBB: ir.Block = active_finally[0]
                    pending_ret_flag: ir.AllocaInstr = active_finally[1]
                    pending_ret_val: ir.AllocaInstr | None = active_finally[2]
                    if pending_ret_val:
                        Gen._store(Val, pending_ret_val)
                    Gen._store(ir.Constant(ir.IntType(32), 1), pending_ret_flag)
                    Gen.builder.branch(FinallyBB)
                else:
                    Gen.emit_return(Val)
        else:
            if active_finally:
                FinallyBB: ir.Block = active_finally[0]
                pending_ret_flag: ir.AllocaInstr = active_finally[1]
                pending_ret_val: ir.AllocaInstr | None = active_finally[2]
                Gen._store(ir.Constant(ir.IntType(32), 1), pending_ret_flag)
                Gen.builder.branch(FinallyBB)
            else:
                ret_type: ir.Type | None = getattr(Gen.func, 'ftype', None) and Gen.func.ftype.return_type if getattr(Gen, 'func', None) else None
                if isinstance(ret_type, ir.IntType):
                    Gen.builder.ret(ir.Constant(ret_type, 0))
                elif isinstance(ret_type, ir.PointerType):
                    Gen.builder.ret(ir.Constant(ret_type, None))
                elif isinstance(ret_type, (ir.FloatType, ir.DoubleType)):
                    Gen.builder.ret(ir.Constant(ret_type, 0.0))
                elif isinstance(ret_type, ir.BaseStructType):
                    if ret_type.elements:
                        zero_val: ir.Constant = ir.Constant(ret_type, [ir.Constant(et, None) if isinstance(et, (ir.PointerType, ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType)) else ir.Constant(et, 0) if isinstance(et, (ir.IntType, ir.FloatType, ir.DoubleType)) else ir.Constant(et, ir.Undefined) for et in ret_type.elements])
                    else:
                        zero_val = ir.Constant(ret_type, None)
                    Gen.builder.ret(zero_val)
                elif isinstance(ret_type, ir.ArrayType):
                    Gen.builder.ret(ir.Constant(ret_type, None))
                else:
                    Gen.emit_return()