//! Callback-based ASGI scheduler (CPython-focused).

use pyo3::exceptions::PyStopIteration;
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::sync::Arc;

pub(crate) struct CallbackScheduler {
    call_soon_fn: Py<PyAny>,
    pynone: Py<PyAny>,
    use_fast_capi: bool,
}

impl CallbackScheduler {
    pub(crate) fn new(py: Python<'_>, loop_obj: &Py<PyAny>) -> PyResult<Self> {
        let call_soon_fn = loop_obj
            .bind(py)
            .getattr(intern!(py, "call_soon"))?
            .unbind();
        Ok(Self {
            call_soon_fn,
            pynone: py.None(),
            use_fast_capi: default_fast_capi(),
        })
    }

    pub(crate) fn run(self: &Arc<Self>, py: Python<'_>, coro: Py<PyAny>, ctx: Py<PyAny>) {
        let ctx_kwargs = PyDict::new(py);
        let _ = ctx_kwargs.set_item(intern!(py, "context"), ctx.clone_ref(py));
        let state = Arc::new(CallbackSchedulerState {
            sched: self.clone(),
            coro,
            ctx,
            ctx_kwargs: ctx_kwargs.unbind(),
        });

        let step = Py::new(
            py,
            CallbackSchedulerStep {
                state: state.clone(),
            },
        );
        if let Ok(step) = step {
            let _ = state.ctx.call_method1(py, "run", (step,));
        }
    }

    fn step(&self, py: Python<'_>, state: Arc<CallbackSchedulerState>) {
        #[cfg(not(PyPy))]
        let res = if self.use_fast_capi {
            unsafe {
                let mut pres = std::ptr::null_mut::<pyo3::ffi::PyObject>();
                let gres =
                    pyo3::ffi::PyIter_Send(state.coro.as_ptr(), self.pynone.as_ptr(), &mut pres);
                match gres {
                    pyo3::ffi::PySendResult::PYGEN_NEXT => Bound::from_owned_ptr_or_err(py, pres),
                    pyo3::ffi::PySendResult::PYGEN_RETURN => {
                        // Completed
                        return;
                    }
                    pyo3::ffi::PySendResult::PYGEN_ERROR => Err(PyErr::fetch(py)),
                }
            }
        } else {
            state
                .coro
                .bind(py)
                .call_method1("send", (self.pynone.clone_ref(py),))
        };
        #[cfg(PyPy)]
        let res = state
            .coro
            .bind(py)
            .call_method1("send", (self.pynone.clone_ref(py),));
        match res {
            Ok(v) => {
                if let Ok(blocking) = v
                    .getattr("_asyncio_future_blocking")
                    .and_then(|b| b.extract::<bool>())
                {
                    if blocking {
                        let _ = v.setattr("_asyncio_future_blocking", false);
                        let waker = Py::new(
                            py,
                            CallbackSchedulerWaker {
                                state: state.clone(),
                            },
                        );
                        if let Ok(waker) = waker {
                            let _ = v.call_method(
                                "add_done_callback",
                                (waker,),
                                Some(state.ctx_kwargs.bind(py)),
                            );
                        }
                        return;
                    }
                }
                self.reschedule(py, state);
            }
            Err(err) => {
                if err.is_instance_of::<PyStopIteration>(py) {
                    // Coroutine finished. Completion is signaled by ASGI send.
                    let _ = err;
                    return;
                }
                // Propagate error into coroutine via throw.
                self.throw(py, state, err.value(py).clone().into());
            }
        }
    }

    fn reschedule(&self, py: Python<'_>, state: Arc<CallbackSchedulerState>) {
        let step = Py::new(
            py,
            CallbackSchedulerStep {
                state: state.clone(),
            },
        );
        if let Ok(step) = step {
            let call_soon = self.call_soon_fn.bind(py);
            let _ = call_soon.call((step,), Some(state.ctx_kwargs.bind(py)));
        }
    }

    fn throw(&self, py: Python<'_>, state: Arc<CallbackSchedulerState>, err: Py<PyAny>) {
        #[cfg(not(PyPy))]
        if self.use_fast_capi {
            unsafe {
                let _ = pyo3::ffi::PyObject_CallMethodOneArg(
                    state.coro.as_ptr(),
                    pyo3::intern!(py, "throw").as_ptr(),
                    err.as_ptr(),
                );
                pyo3::ffi::PyErr_Clear();
                return;
            }
        }
        let _ = state.coro.bind(py).call_method1("throw", (err,));
        unsafe {
            pyo3::ffi::PyErr_Clear();
        }
    }
}

#[inline]
fn default_fast_capi() -> bool {
    #[cfg(PyPy)]
    {
        return false;
    }
    #[cfg(not(PyPy))]
    {
        !matches!(
            std::env::var("FALCORN_ASGI_FAST_CAPI").ok().as_deref(),
            Some("0" | "false" | "FALSE" | "no" | "NO" | "off" | "OFF")
        )
    }
}

struct CallbackSchedulerState {
    sched: Arc<CallbackScheduler>,
    coro: Py<PyAny>,
    ctx: Py<PyAny>,
    ctx_kwargs: Py<PyDict>,
}

#[pyclass(frozen, module = "falcorn._falcorn")]
struct CallbackSchedulerStep {
    state: Arc<CallbackSchedulerState>,
}

#[pymethods]
impl CallbackSchedulerStep {
    fn __call__(&self, py: Python<'_>) {
        self.state.sched.step(py, self.state.clone());
    }
}

#[pyclass(frozen, module = "falcorn._falcorn")]
struct CallbackSchedulerWaker {
    state: Arc<CallbackSchedulerState>,
}

#[pymethods]
impl CallbackSchedulerWaker {
    fn __call__(&self, py: Python<'_>, fut: Py<PyAny>) {
        match fut.call_method0(py, "result") {
            Ok(_) => self.state.sched.step(py, self.state.clone()),
            Err(err) => {
                self.state
                    .sched
                    .throw(py, self.state.clone(), err.value(py).clone().into())
            }
        }
    }
}
