use crate::common::wtf8::{Wtf8Buf, wtf8_concat};

use crate::contextvars::PyContext;

use crate::{

common::lock::PyRwLock,

vm::{

AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,

builtins::{

PyBaseException, PyBaseExceptionRef, PyDict, PyGenericAlias, PyList, PyListRef,

PyModule, PySet, PyTraceback, PyTuple, PyType, PyTypeRef, alloc_default,

},

extend_module,

function::{FuncArgs, KwArgs, OptionalArg, PySetterValue},

protocol::PyIterReturn,

recursion::ReprGuard,

types::{

Callable, Constructor, Destructor, Initializer, IterNext, Iterable, Representable,

SelfIter,

},

warn,

},

};

use core::sync::atomic::{AtomicBool, AtomicI32, AtomicU64, Ordering};

use crossbeam_utils::atomic::AtomicCell;

pub(crate) fn module_exec(vm: &VirtualMachine, module: &Py<PyModule>) -> PyResult<()> {

__module_exec(vm, module);

// Initialize module-level state

let weakref_module = vm.import("weakref", 0)?;

let weak_set_class = vm

.get_attribute_opt(&weakref_module, vm.ctx.intern_str("WeakSet"))?

.ok_or_else(|| vm.new_attribute_error("WeakSet not found"))?;

let scheduled_tasks = weak_set_class.call((), vm)?;

let eager_tasks = PySet::default().into_ref(&vm.ctx);

let current_tasks = PyDict::default().into_ref(&vm.ctx);

extend_module!(vm, module, {

"_scheduled_tasks" => scheduled_tasks,

"_eager_tasks" => eager_tasks,

"_current_tasks" => current_tasks,

});

// Register fork handler to clear task state in child process

#[cfg(unix)]

{

let on_fork = vm

.get_attribute_opt(module.as_object(), vm.ctx.intern_str("_on_fork"))?

.expect("_on_fork not found in _asyncio module");

vm.state.after_forkers_child.lock().push(on_fork);

}

Ok(())

}

#[derive(FromArgs)]

struct AddDoneCallbackArgs {

#[pyarg(positional, name = "fn")]

func: PyObjectRef,

#[pyarg(named, optional)]

context: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct CancelArgs {

#[pyarg(any, optional)]

msg: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct LoopArg {

#[pyarg(any, name = "loop", optional)]

loop_: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct GetStackArgs {

#[pyarg(named, optional)]

limit: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct PrintStackArgs {

#[pyarg(named, optional)]

limit: Option<PyObjectRef>,

#[pyarg(named, optional)]

file: Option<PyObjectRef>,

}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]

enum FutureState {

Pending,

Cancelled,

Finished,

}

impl FutureState {

const fn as_str(self) -> &'static str {

match self {

Self::Pending => "PENDING",

Self::Cancelled => "CANCELLED",

Self::Finished => "FINISHED",

}

}

}

#[pyattr]

#[pyclass(name = "Future", module = "_asyncio", traverse)]

#[derive(Debug, PyPayload)]

struct PyFuture {

fut_loop: PyRwLock<Option<PyObjectRef>>,

fut_callback0: PyRwLock<Option<PyObjectRef>>,

fut_context0: PyRwLock<Option<PyObjectRef>>,

fut_callbacks: PyRwLock<Option<PyObjectRef>>,

fut_exception: PyRwLock<Option<PyObjectRef>>,

fut_exception_tb: PyRwLock<Option<PyObjectRef>>,

fut_result: PyRwLock<Option<PyObjectRef>>,

fut_source_tb: PyRwLock<Option<PyObjectRef>>,

fut_cancel_msg: PyRwLock<Option<PyObjectRef>>,

fut_cancelled_exc: PyRwLock<Option<PyObjectRef>>,

fut_awaited_by: PyRwLock<Option<PyObjectRef>>,

#[pytraverse(skip)]

fut_state: AtomicCell<FutureState>,

#[pytraverse(skip)]

fut_awaited_by_is_set: AtomicBool,

#[pytraverse(skip)]

fut_log_tb: AtomicBool,

#[pytraverse(skip)]

fut_blocking: AtomicBool,

}

impl Constructor for PyFuture {

type Args = FuncArgs;

fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {

Ok(Self::new_empty())

}

}

#[derive(FromArgs)]

struct FutureInitArgs {

#[pyarg(named, name = "loop", optional)]

loop_: Option<PyObjectRef>,

}

impl Initializer for PyFuture {

type Args = FutureInitArgs;

fn init(zelf: &Py<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {

Self::py_init(zelf, args.loop_, vm)

}

}

#[pyclass(

impl PyFuture {

fn new_empty() -> Self {

Self {

fut_loop: PyRwLock::new(None),

fut_callback0: PyRwLock::new(None),

fut_context0: PyRwLock::new(None),

fut_callbacks: PyRwLock::new(None),

fut_exception: PyRwLock::new(None),

fut_exception_tb: PyRwLock::new(None),

fut_result: PyRwLock::new(None),

fut_source_tb: PyRwLock::new(None),

fut_cancel_msg: PyRwLock::new(None),

fut_cancelled_exc: PyRwLock::new(None),

fut_awaited_by: PyRwLock::new(None),

fut_state: AtomicCell::new(FutureState::Pending),

fut_awaited_by_is_set: AtomicBool::new(false),

fut_log_tb: AtomicBool::new(false),

fut_blocking: AtomicBool::new(false),

}

}

fn py_init(

zelf: &Py<Self>,

loop_: Option<PyObjectRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

if zelf.fut_loop.read().is_some() {

return Err(vm.new_runtime_error(wtf8_concat!(

zelf.class().fully_qualified_name(vm)?,

" object is already initialized"

)));

}

// Get the event loop

let loop_obj = match loop_ {

Some(l) if !vm.is_none(&l) => l,

_ => get_event_loop(vm)?,

};

*zelf.fut_loop.write() = Some(loop_obj.clone());

// Check if loop has get_debug method and call it

if let Ok(Some(get_debug)) =

vm.get_attribute_opt(&loop_obj, vm.ctx.intern_str("get_debug"))

&& let Ok(debug) = get_debug.call((), vm)

&& debug.try_to_bool(vm).unwrap_or(false)

{

// Get source traceback

if let Ok(tb_module) = vm.import("traceback", 0)

&& let Ok(Some(extract_stack)) =

vm.get_attribute_opt(&tb_module, vm.ctx.intern_str("extract_stack"))

&& let Ok(tb) = extract_stack.call((), vm)

{

*zelf.fut_source_tb.write() = Some(tb);

}

}

Ok(())

}

#[pymethod]

fn result(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {

match zelf.fut_state.load() {

FutureState::Pending => Err(new_invalid_state_error(vm, "Result is not ready.")),

FutureState::Cancelled => {

let exc = zelf.make_cancelled_error_impl(vm);

Err(exc)

}

FutureState::Finished => {

zelf.fut_log_tb.store(false, Ordering::Relaxed);

let fut_exception = zelf.fut_exception.read().clone();

if let Some(exc) = fut_exception {

let exc: PyBaseExceptionRef = exc.downcast().unwrap();

// Restore the original traceback to prevent traceback accumulation

let fut_exception_tb = zelf.fut_exception_tb.read().clone();

if let Some(tb) = fut_exception_tb

&& let Ok(tb) = tb.downcast::<PyTraceback>()

{

exc.set_traceback(Some(tb));

}

Err(exc)

} else {

Ok(zelf

.fut_result

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none()))

}

}

}

}

#[pymethod]

fn exception(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {

match zelf.fut_state.load() {

FutureState::Pending => Err(new_invalid_state_error(vm, "Exception is not set.")),

FutureState::Cancelled => {

let exc = zelf.make_cancelled_error_impl(vm);

Err(exc)

}

FutureState::Finished => {

zelf.fut_log_tb.store(false, Ordering::Relaxed);

Ok(zelf

.fut_exception

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none()))

}

}

}

#[pymethod]

fn set_result(zelf: PyRef<Self>, result: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

if zelf.fut_loop.read().is_none() {

return Err(vm.new_runtime_error("Future object is not initialized."));

}

if zelf.fut_state.load() != FutureState::Pending {

return Err(new_invalid_state_error(vm, "invalid state"));

}

*zelf.fut_result.write() = Some(result);

zelf.fut_state.store(FutureState::Finished);

Self::schedule_callbacks(&zelf, vm)?;

Ok(())

}

#[pymethod]

fn set_exception(

zelf: PyRef<Self>,

exception: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult<()> {

if zelf.fut_loop.read().is_none() {

return Err(vm.new_runtime_error("Future object is not initialized."));

}

if zelf.fut_state.load() != FutureState::Pending {

return Err(new_invalid_state_error(vm, "invalid state"));

}

// Normalize the exception

let exc = if let Some(cls) = exception.downcast_ref::<PyType>()

&& cls.fast_issubclass(vm.ctx.exceptions.base_exception_type)

{

let exc = exception.call((), vm)?;

if zelf.fut_state.load() != FutureState::Pending {

return Err(new_invalid_state_error(vm, "invalid state"));

}

exc

} else {

exception

};

if !exc.fast_isinstance(vm.ctx.exceptions.base_exception_type) {

return Err(vm.new_type_error(format!(

"exception must be a BaseException, not {}",

exc.class().name()

)));

}

// Wrap StopIteration in RuntimeError

let exc = if exc.fast_isinstance(vm.ctx.exceptions.stop_iteration) {

let msg = "StopIteration interacts badly with generators and cannot be raised into a Future";

let runtime_err = vm.new_runtime_error(msg.to_string());

// Set cause and context to the original StopIteration

let stop_iter: PyRef<PyBaseException> = exc.downcast().unwrap();

runtime_err.set_cause(Some(stop_iter.clone()));

runtime_err.set_context(Some(stop_iter));

runtime_err.into()

} else {

exc

};

// Save the original traceback for later restoration

if let Ok(exc_ref) = exc.clone().downcast::<PyBaseException>() {

let tb = exc_ref.traceback().map(|tb| tb.into());

*zelf.fut_exception_tb.write() = tb;

}

*zelf.fut_exception.write() = Some(exc);

zelf.fut_state.store(FutureState::Finished);

Self::schedule_callbacks(&zelf, vm)?;

zelf.fut_log_tb.store(true, Ordering::Relaxed);

Ok(())

}

#[pymethod]

fn add_done_callback(

zelf: PyRef<Self>,

args: AddDoneCallbackArgs,

vm: &VirtualMachine,

) -> PyResult<()> {

zelf.add_done_callback_impl(zelf.as_object(), args, vm)

}

fn add_done_callback_impl(

&self,

owner: &PyObject,

args: AddDoneCallbackArgs,

vm: &VirtualMachine,

) -> PyResult<()> {

if self.fut_loop.read().is_none() {

return Err(vm.new_runtime_error("Future object is not initialized."));

}

let ctx = match args.context {

Some(c) => c,

None => get_copy_context(vm)?,

};

if self.fut_state.load() != FutureState::Pending {

self.call_soon_with_context(owner, args.func, Some(ctx), vm)?;

} else {

let callbacks = self.fut_callbacks.read().clone();

// Once a list exists, append to it even if callback0 was removed.

// Refilling callback0 would move the new callback ahead of the list.

if callbacks.is_none() && self.fut_callback0.read().is_none() {

*self.fut_callback0.write() = Some(args.func);

*self.fut_context0.write() = Some(ctx);

} else {

let tuple = vm.ctx.new_tuple(vec![args.func, ctx]);

if let Some(callbacks) = callbacks {

callbacks

.downcast_ref::<PyList>()

.unwrap()

.borrow_vec_mut()

.push(tuple.into());

} else {

*self.fut_callbacks.write() =

Some(vm.ctx.new_list(vec![tuple.into()]).into());

}

}

}

Ok(())

}

#[pymethod]

fn remove_done_callback(

zelf: &Py<Self>,

RemoveDoneCallbackArgs { func }: RemoveDoneCallbackArgs,

vm: &VirtualMachine,

) -> PyResult<isize> {

zelf.remove_done_callback_impl(&func, vm)

}

fn remove_done_callback_impl(

&self,

func: &PyObject,

vm: &VirtualMachine,

) -> PyResult<isize> {

if self.fut_loop.read().is_none() {

return Err(vm.new_runtime_error("Future object is not initialized."));

}

let mut cleared_callback0 = 0;

let cb0 = self.fut_callback0.read().clone();

if let Some(cb0) = cb0 {

let cmp = vm.identical_or_equal(&cb0, func);

drop(cb0);

if cmp? {

// Dropping callbacks or contexts may run finalizers. Release

// each field lock before dropping the previous reference.

let callback = self.fut_callback0.write().take();

drop(callback);

let context = self.fut_context0.write().take();

drop(context);

cleared_callback0 = 1;

}

}

let len = {

let callbacks = self.fut_callbacks.read().clone();

let Some(callbacks) = callbacks else {

return Ok(cleared_callback0);

};

callbacks

.downcast_ref::<PyList>()

.unwrap()

.borrow_vec()

.len()

};

if len == 0 {

let callbacks = self.fut_callbacks.write().take();

drop(callbacks);

return Ok(cleared_callback0);

}

if len == 1 {

let item = {

let callbacks = self.fut_callbacks.read().clone().unwrap();

callbacks.downcast_ref::<PyList>().unwrap().borrow_vec()[0].clone()

};

let tuple = item.downcast_ref::<PyTuple>().unwrap();

let cmp = vm.identical_or_equal(&tuple.as_slice()[0], func);

drop(item);

if cmp? {

let callbacks = self.fut_callbacks.write().take();

drop(callbacks);

return Ok(1 + cleared_callback0);

}

return Ok(cleared_callback0);

}

let mut new_callbacks = Vec::with_capacity(len);

let mut i = 0;

loop {

// Comparisons can remove, append, or schedule the callbacks.

let item = {

let callbacks = self.fut_callbacks.read().clone();

let Some(callbacks) = callbacks else { break };

callbacks

.downcast_ref::<PyList>()

.unwrap()

.borrow_vec()

.get(i)

.cloned()

};

let Some(item) = item else { break };

let tuple = item.downcast_ref::<PyTuple>().unwrap();

if !vm.identical_or_equal(&tuple.as_slice()[0], func)? {

new_callbacks.push(item);

}

i += 1;

}

let callbacks = self.fut_callbacks.read().clone();

if new_callbacks.is_empty() || callbacks.is_none() {

let callbacks = self.fut_callbacks.write().take();

drop(callbacks);

return Ok(len as isize + cleared_callback0);

}

let callbacks = callbacks.unwrap();

let list = callbacks.downcast_ref::<PyList>().unwrap();

let current_len = list.borrow_vec().len();

let remaining = new_callbacks.len();

if remaining != current_len {

// Keep the filtered snapshot alive while removed callbacks run

// finalizers: they can recursively remove the surviving callbacks.

let previous =

core::mem::replace(&mut *list.borrow_vec_mut(), new_callbacks.clone());

for item in previous.into_iter().rev() {

drop(item);

}

}

drop(callbacks);

for item in new_callbacks.into_iter().rev() {

drop(item);

}

Ok(current_len as isize - remaining as isize + cleared_callback0)

}

#[pymethod]

fn cancel(zelf: PyRef<Self>, args: CancelArgs, vm: &VirtualMachine) -> PyResult<bool> {

if zelf.fut_loop.read().is_none() {

return Err(vm.new_runtime_error("Future object is not initialized."));

}

if zelf.fut_state.load() != FutureState::Pending {

// Clear log_tb even when cancel fails

zelf.fut_log_tb.store(false, Ordering::Relaxed);

return Ok(false);

}

*zelf.fut_cancel_msg.write() = args.msg;

zelf.fut_state.store(FutureState::Cancelled);

Self::schedule_callbacks(&zelf, vm)?;

Ok(true)

}

#[pymethod]

fn cancelled(zelf: &Py<Self>) -> bool {

zelf.fut_state.load() == FutureState::Cancelled

}

#[pymethod]

fn done(zelf: &Py<Self>) -> bool {

zelf.fut_state.load() != FutureState::Pending

}

#[pymethod]

fn get_loop(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {

zelf.fut_loop

.read()

.clone()

.ok_or_else(|| vm.new_runtime_error("Future object is not initialized."))

}

#[pymethod]

fn _make_cancelled_error(zelf: &Py<Self>, vm: &VirtualMachine) -> PyBaseExceptionRef {

zelf.make_cancelled_error_impl(vm)

}

fn make_cancelled_error_impl(&self, vm: &VirtualMachine) -> PyBaseExceptionRef {

// If a saved CancelledError exists, take it (clearing the stored reference)

if let Some(exc) = self.fut_cancelled_exc.write().take()

&& let Ok(exc) = exc.downcast::<PyBaseException>()

{

return exc;

}

let msg = self.fut_cancel_msg.read().clone();

let args = if let Some(m) = msg { vec![m] } else { vec![] };

match get_cancelled_error_type(vm) {

Ok(cancelled_error) => vm.new_exception(cancelled_error, args),

Err(_) => vm.new_runtime_error("cancelled"),

}

}

fn schedule_callbacks(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {

zelf.schedule_callbacks_impl(zelf.as_object(), vm)

}

fn schedule_callbacks_impl(&self, owner: &PyObject, vm: &VirtualMachine) -> PyResult<()> {

let callback = self.fut_callback0.write().take();

let context = self.fut_context0.write().take();

if let Some(callback) = callback {

// Keep the list visible until the first call_soon returns: it

// may remove callbacks, and may also fail while scheduling.

if let Err(err) = self.call_soon_with_context(owner, callback, context, vm) {

let callbacks = self.fut_callbacks.write().take();

drop(callbacks);

return Err(err);

}

}

let callbacks = self.fut_callbacks.write().take();

if let Some(callbacks) = callbacks {

let list = callbacks.downcast_ref::<PyList>().unwrap();

let items = list.borrow_vec().to_vec();

for item in &items {

let tuple = item.downcast_ref::<PyTuple>().unwrap();

self.call_soon_with_context(

owner,

tuple.as_slice()[0].clone(),

Some(tuple.as_slice()[1].clone()),

vm,

)?;

}

}

Ok(())

}

fn call_soon_with_context(

&self,

owner: &PyObject,

callback: PyObjectRef,

context: Option<PyObjectRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

let loop_obj = self.fut_loop.read().clone();

if let Some(loop_obj) = loop_obj {

// call_soon(callback, *args, context=context)

// callback receives the future as its argument

let future_arg: PyObjectRef = owner.to_owned();

let args = if let Some(ctx) = context {

FuncArgs::new(

vec![callback, future_arg],

KwArgs::new(core::iter::once((Wtf8Buf::from("context"), ctx)).collect()),

)

} else {

FuncArgs::new(vec![callback, future_arg], KwArgs::default())

};

vm.call_method(&loop_obj, "call_soon", args)?;

}

Ok(())

}

// Properties

#[pygetset]

fn _state(zelf: &Py<Self>) -> &'static str {

zelf.fut_state.load().as_str()

}

#[pygetset]

fn _asyncio_future_blocking(zelf: &Py<Self>) -> bool {

zelf.fut_blocking.load(Ordering::Relaxed)

}

#[pygetset(setter)]

fn set__asyncio_future_blocking(

zelf: &Py<Self>,

value: PySetterValue<bool>,

vm: &VirtualMachine,

) -> PyResult<()> {

match value {

PySetterValue::Assign(v) => {

zelf.fut_blocking.store(v, Ordering::Relaxed);

Ok(())

}

PySetterValue::Delete => Err(vm.new_attribute_error("cannot delete attribute")),

}

}

#[pygetset]

fn _loop(zelf: &Py<Self>, vm: &VirtualMachine) -> PyObjectRef {

zelf.fut_loop

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none())

}

#[pygetset]

fn _callbacks(zelf: &Py<Self>, vm: &VirtualMachine) -> PyObjectRef {

let mut result = Vec::new();

let fut_callback0 = zelf.fut_callback0.read().clone();

if let Some(cb0) = fut_callback0 {

let ctx0 = zelf

.fut_context0

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none());

result.push(vm.ctx.new_tuple(vec![cb0, ctx0]).into());

}

let fut_callbacks = zelf.fut_callbacks.read().clone();

if let Some(callbacks) = fut_callbacks {

let list: PyListRef = callbacks.downcast().unwrap();

for item in list.borrow_vec().iter() {

result.push(item.clone());

}

}

// Return None if no callbacks

if result.is_empty() {

vm.ctx.none()

} else {

vm.ctx.new_list(result).into()

}

}

#[pygetset]

fn _result(zelf: &Py<Self>, vm: &VirtualMachine) -> PyObjectRef {

zelf.fut_result

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none())

}

#[pygetset]

fn _exception(zelf: &Py<Self>, vm: &VirtualMachine) -> PyObjectRef {

zelf.fut_exception

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none())

}

#[pygetset]

fn _log_traceback(zelf: &Py<Self>) -> bool {

zelf.fut_log_tb.load(Ordering::Relaxed)

}

#[pygetset(setter)]

fn set__log_traceback(

zelf: &Py<Self>,

value: PySetterValue<bool>,

vm: &VirtualMachine,

) -> PyResult<()> {

match value {

PySetterValue::Assign(v) => {

if v {

return Err(vm.new_value_error("_log_traceback can only be set to False"));

}

zelf.fut_log_tb.store(false, Ordering::Relaxed);

Ok(())

}

PySetterValue::Delete => Err(vm.new_attribute_error("cannot delete attribute")),

}

}

#[pygetset]

fn _source_traceback(zelf: &Py<Self>, vm: &VirtualMachine) -> PyObjectRef {

zelf.fut_source_tb

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none())

}

#[pygetset]

fn _cancel_message(zelf: &Py<Self>, vm: &VirtualMachine) -> PyObjectRef {

zelf.fut_cancel_msg

.read()

.clone()

.unwrap_or_else(|| vm.ctx.none())

}

#[pygetset(setter)]

fn set__cancel_message(zelf: &Py<Self>, value: PySetterValue) {

match value {

PySetterValue::Assign(v) => *zelf.fut_cancel_msg.write() = Some(v),

PySetterValue::Delete => *zelf.fut_cancel_msg.write() = None,

}

}

#[pygetset]

fn _asyncio_awaited_by(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {

let awaited_by = zelf.fut_awaited_by.read().clone();

match awaited_by {

None => Ok(vm.ctx.none()),

Some(obj) => {

if zelf.fut_awaited_by_is_set.load(Ordering::Relaxed) {

// Already a Set

Ok(obj)

} else {

// Single object - create a Set for the return value

let new_set = PySet::default().into_ref(&vm.ctx);