//! ASGI callbacks - Implements send() and receive() callables for ASGI apps.

use super::monotonic_now_ns;
use super::{AsgiContext, AsgiSendEvent};
use bytes::Bytes;
use parking_lot::Mutex;
use pyo3::exceptions::PyStopIteration;
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::sync::PyOnceLock;
use pyo3::types::{PyBytes, PyDict};
use std::sync::{Arc, OnceLock};
use tokio::sync::oneshot;

/// ASGI Scope type (just a PyDict wrapper for clarity)
pub type AsgiScope = Py<PyDict>;

#[pyclass]
struct ReadyAwaitable {
    result: Py<PyAny>,
    yielded: bool,
}

#[pymethods]
impl ReadyAwaitable {
    fn __await__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&mut self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        if self.yielded {
            return Err(PyStopIteration::new_err(py.None()));
        }
        self.yielded = true;
        Err(PyStopIteration::new_err((self.result.clone_ref(py),)))
    }
}

#[inline]
fn ready_awaitable<'py>(py: Python<'py>, value: Py<PyAny>) -> PyResult<Bound<'py, PyAny>> {
    let aw = Py::new(
        py,
        ReadyAwaitable {
            result: value,
            yielded: false,
        },
    )?;
    Ok(aw.into_bound(py).into_any())
}

#[pyclass(frozen, freelist = 256)]
struct PyEmptyAwaitable;

#[pymethods]
impl PyEmptyAwaitable {
    fn __await__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&self) -> Option<()> {
        None
    }
}

#[inline]
fn empty_awaitable<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
    static EMPTY_AWAITABLE: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
    let obj = EMPTY_AWAITABLE.get_or_try_init(py, || {
        let aw = Py::new(py, PyEmptyAwaitable)?;
        Ok::<Py<PyAny>, PyErr>(aw.into_any())
    })?;
    Ok(obj.bind(py).clone().into_any())
}

#[inline]
fn use_empty_awaitable() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| {
        matches!(
            std::env::var("FALCORN_ASGI_EMPTY_AWAITABLE")
                .ok()
                .as_deref(),
            Some("1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON")
        )
    })
}

#[derive(Default)]
struct ReceiveState {
    body: Arc<Bytes>,
    pos: usize,
    sent_request: bool,
    sent_disconnect: bool,
}

#[derive(Default)]
struct AsgiIoTimings {
    dispatch_enqueue_ns: u64,
    dispatch_dequeue_ns: u64,
    dispatch_schedule_ns: u64,
    first_receive_ns: u64,
    last_send_ns: u64,
    scope_build_ns: u64,
    coro_create_ns: u64,
}

/// ASGI Receive callable
#[pyclass(frozen)]
#[derive(Clone)]
pub struct AsgiReceive {
    state: Arc<Mutex<ReceiveState>>,
    timings: Option<Arc<Mutex<AsgiIoTimings>>>,
}

impl AsgiReceive {
    fn new(body: Arc<Bytes>, timings: Option<Arc<Mutex<AsgiIoTimings>>>) -> Self {
        Self {
            state: Arc::new(Mutex::new(ReceiveState {
                body,
                pos: 0,
                sent_request: false,
                sent_disconnect: false,
            })),
            timings,
        }
    }
}

#[pymethods]
impl AsgiReceive {
    /// Called by ASGI app to receive events
    fn __call__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        if let Some(timings) = &self.timings {
            let mut t = timings.lock();
            if t.first_receive_ns == 0 {
                t.first_receive_ns = monotonic_now_ns();
            }
        }

        let mut state = self.state.lock();

        let event = PyDict::new(py);
        if !state.sent_request && state.body.is_empty() {
            event.set_item(intern!(py, "type"), "http.request")?;
            event.set_item(intern!(py, "body"), PyBytes::new(py, &[]))?;
            event.set_item(intern!(py, "more_body"), false)?;
            state.sent_request = true;
        } else if state.pos < state.body.len() {
            let remaining = state.body.len() - state.pos;
            let chunk_size = 64 * 1024;
            let take = remaining.min(chunk_size);
            let start = state.pos;
            let end = start + take;
            state.pos = end;
            let more_body = state.pos < state.body.len();

            event.set_item(intern!(py, "type"), "http.request")?;
            event.set_item(
                intern!(py, "body"),
                PyBytes::new(py, &state.body[start..end]),
            )?;
            event.set_item(intern!(py, "more_body"), more_body)?;
            state.sent_request = true;
        } else if !state.sent_disconnect {
            event.set_item(intern!(py, "type"), "http.disconnect")?;
            state.sent_disconnect = true;
        } else {
            event.set_item(intern!(py, "type"), "http.disconnect")?;
        }

        ready_awaitable(py, event.into_any().unbind())
    }
}

#[derive(Default)]
struct SendState {
    events: Option<Vec<AsgiSendEvent>>,
    completion_tx: Option<oneshot::Sender<()>>,
    completed: bool,
    scope: Option<AsgiScope>,
}

pub struct AsgiSendEventPool {
    max_cached: usize,
    buffers: Mutex<Vec<Vec<AsgiSendEvent>>>,
}

impl AsgiSendEventPool {
    pub fn new(max_cached: usize) -> Arc<Self> {
        Arc::new(Self {
            max_cached: max_cached.max(1),
            buffers: Mutex::new(Vec::new()),
        })
    }

    fn acquire(&self) -> Vec<AsgiSendEvent> {
        self.buffers.lock().pop().unwrap_or_default()
    }

    fn release(&self, mut buf: Vec<AsgiSendEvent>) {
        buf.clear();
        let mut guard = self.buffers.lock();
        if guard.len() < self.max_cached {
            guard.push(buf);
        }
    }
}

/// ASGI Send callable
#[pyclass(frozen)]
#[derive(Clone)]
pub struct AsgiSend {
    state: Arc<Mutex<SendState>>,
    pool: Arc<AsgiSendEventPool>,
    timings: Option<Arc<Mutex<AsgiIoTimings>>>,
}

impl AsgiSend {
    fn new(
        pool: Arc<AsgiSendEventPool>,
        completion_tx: oneshot::Sender<()>,
        timings: Option<Arc<Mutex<AsgiIoTimings>>>,
    ) -> Self {
        let events = pool.acquire();
        Self {
            state: Arc::new(Mutex::new(SendState {
                events: Some(events),
                completion_tx: Some(completion_tx),
                completed: false,
                scope: None,
            })),
            pool,
            timings,
        }
    }

    pub fn take_events(&self) -> PyResult<Vec<AsgiSendEvent>> {
        let mut state = self.state.lock();
        Ok(state.events.take().unwrap_or_default())
    }

    pub fn recycle_events(&self, events: Vec<AsgiSendEvent>) {
        self.pool.release(events);
    }

    pub fn timings_snapshot(&self) -> Option<(u64, u64)> {
        self.timings.as_ref().map(|t| {
            let t = t.lock();
            (t.first_receive_ns, t.last_send_ns)
        })
    }

    pub fn dispatch_timings_snapshot(&self) -> Option<(u64, u64, u64)> {
        self.timings.as_ref().map(|t| {
            let t = t.lock();
            (
                t.dispatch_enqueue_ns,
                t.dispatch_dequeue_ns,
                t.dispatch_schedule_ns,
            )
        })
    }

    pub fn scope_coro_snapshot(&self) -> Option<(u64, u64)> {
        self.timings.as_ref().map(|t| {
            let t = t.lock();
            (t.scope_build_ns, t.coro_create_ns)
        })
    }

    pub fn mark_dispatch_enqueued(&self, now_ns: u64) {
        if let Some(timings) = &self.timings {
            let mut t = timings.lock();
            t.dispatch_enqueue_ns = now_ns;
        }
    }

    pub fn mark_dispatch_dequeued(&self, now_ns: u64) {
        if let Some(timings) = &self.timings {
            let mut t = timings.lock();
            t.dispatch_dequeue_ns = now_ns;
        }
    }

    pub fn mark_dispatch_scheduled(&self, now_ns: u64) {
        if let Some(timings) = &self.timings {
            let mut t = timings.lock();
            t.dispatch_schedule_ns = now_ns;
        }
    }

    pub fn mark_scope_build_ns(&self, dur_ns: u64) {
        if let Some(timings) = &self.timings {
            let mut t = timings.lock();
            t.scope_build_ns = dur_ns;
        }
    }

    pub fn mark_coro_create_ns(&self, dur_ns: u64) {
        if let Some(timings) = &self.timings {
            let mut t = timings.lock();
            t.coro_create_ns = dur_ns;
        }
    }

    pub fn is_completed(&self) -> bool {
        self.state.lock().completed
    }

    pub fn fail_request(&self) {
        let mut state = self.state.lock();
        if let Some(scope) = state.scope.take() {
            let _ = Python::attach(|py| AsgiContext::global().recycle_scope(py, scope));
        }
        if let Some(tx) = state.completion_tx.take() {
            let _ = tx.send(());
        }
        state.completed = true;
    }

    pub fn attach_scope(&self, scope: AsgiScope) {
        let mut state = self.state.lock();
        state.scope = Some(scope);
    }
}

#[pymethods]
impl AsgiSend {
    /// Called by ASGI app to send events
    fn __call__<'py>(&self, py: Python<'py>, message: Py<PyDict>) -> PyResult<Bound<'py, PyAny>> {
        let now_ns = if self.timings.is_some() {
            Some(monotonic_now_ns())
        } else {
            None
        };

        let msg = message.bind(py);
        let msg_type_any = msg
            .get_item(intern!(py, "type"))?
            .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err("Missing 'type' in message"))?;
        let msg_type: &str = msg_type_any.extract()?;

        let event = match msg_type {
            "http.response.start" => {
                let status = msg
                    .get_item(intern!(py, "status"))?
                    .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err("Missing 'status'"))?
                    .extract::<u16>()?;

                let headers = if let Some(headers_obj) = msg.get_item(intern!(py, "headers"))? {
                    Self::parse_headers(py, &headers_obj)?
                } else {
                    Vec::new()
                };

                AsgiSendEvent::HttpResponseStart { status, headers }
            }
            "http.response.body" => {
                let body = if let Some(body_obj) = msg.get_item(intern!(py, "body"))? {
                    if let Ok(py_bytes) = body_obj.cast::<PyBytes>() {
                        py_bytes.as_bytes().to_vec()
                    } else {
                        body_obj.extract::<Vec<u8>>()?
                    }
                } else {
                    Vec::new()
                };

                let more_body = msg
                    .get_item(intern!(py, "more_body"))?
                    .map(|v| v.extract::<bool>())
                    .transpose()?
                    .unwrap_or(false);

                AsgiSendEvent::HttpResponseBody { body, more_body }
            }
            _ => {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "Unknown message type: {}",
                    msg_type
                )));
            }
        };

        let mut state = self.state.lock();

        let maybe_completed = matches!(
            event,
            AsgiSendEvent::HttpResponseBody {
                more_body: false,
                ..
            }
        );

        if let Some(events) = state.events.as_mut() {
            events.push(event);
        } else {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "ASGI send buffer already consumed",
            ));
        }

        if maybe_completed && let Some(tx) = state.completion_tx.take() {
            let _ = tx.send(());
            state.completed = true;
            if let Some(scope) = state.scope.take() {
                let _ = AsgiContext::global().recycle_scope(py, scope);
            }
        }

        if let (Some(timings), Some(now_ns)) = (&self.timings, now_ns) {
            let mut t = timings.lock();
            t.last_send_ns = now_ns;
        }

        if use_empty_awaitable() {
            empty_awaitable(py)
        } else {
            ready_awaitable(py, py.None().into_any())
        }
    }
}

impl AsgiSend {
    /// Parse ASGI headers format (list of 2-tuples of bytes)
    fn parse_headers(
        _py: Python,
        headers_obj: &Bound<'_, PyAny>,
    ) -> PyResult<Vec<(Vec<u8>, Vec<u8>)>> {
        use pyo3::types::PyList;

        let headers_list = headers_obj.cast::<PyList>()?;
        let mut result = Vec::with_capacity(headers_list.len());

        for item in headers_list.iter() {
            let tuple = item.cast::<pyo3::types::PyTuple>()?;

            if tuple.len() != 2 {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "Headers must be 2-tuples",
                ));
            }

            let name_obj = tuple.get_item(0)?;
            let value_obj = tuple.get_item(1)?;

            let name = if let Ok(py_bytes) = name_obj.cast::<PyBytes>() {
                py_bytes.as_bytes().to_vec()
            } else {
                name_obj.extract::<Vec<u8>>()?
            };
            let value = if let Ok(py_bytes) = value_obj.cast::<PyBytes>() {
                py_bytes.as_bytes().to_vec()
            } else {
                value_obj.extract::<Vec<u8>>()?
            };

            result.push((name, value));
        }

        Ok(result)
    }
}

/// Create ASGI receive/send callables for a single request.
pub fn create_asgi_io(
    body: Arc<Bytes>,
    send_pool: Arc<AsgiSendEventPool>,
    completion_tx: oneshot::Sender<()>,
    track_timings: bool,
) -> (AsgiReceive, AsgiSend) {
    let timings = if track_timings {
        Some(Arc::new(Mutex::new(AsgiIoTimings::default())))
    } else {
        None
    };
    let receive = AsgiReceive::new(body, timings.clone());
    let send = AsgiSend::new(send_pool, completion_tx, timings);
    (receive, send)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_asgi_io() {
        let pool = AsgiSendEventPool::new(32);
        let (tx, _rx) = oneshot::channel();
        let (_receive, _send) = create_asgi_io(Arc::new(Bytes::new()), pool, tx, true);
        assert!(true);
    }
}
