pdk-classy 1.9.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use std::{cell::RefCell, task::Waker};

use crate::types::HttpCid;

/// Manages WebSocket connection state after HTTP upgrade.
///
/// Handles buffer offsets, pause state, and waker notifications for async frame processing.
pub struct WebSocketReactor {
    #[allow(dead_code)]
    context_id: HttpCid,
    state: RefCell<RawWebSocketReactor>,
}

/// Internal state for WebSocket connection management.
struct RawWebSocketReactor {
    /// Waker for upstream (client→server) handler
    upstream_waker: Option<Waker>,
    /// Waker for downstream (server→client) handler
    downstream_waker: Option<Waker>,
    /// True if upstream data is ready to be read
    upstream_data_ready: bool,
    /// True if downstream data is ready to be read
    downstream_data_ready: bool,
    /// True if upstream (request) flow is paused
    upstream_paused: bool,
    /// True if downstream (response) flow is paused
    downstream_paused: bool,
}

impl WebSocketReactor {
    /// Creates a new WebSocketReactor for the given HTTP context.
    pub fn new(context_id: HttpCid) -> Self {
        Self {
            context_id,
            state: RefCell::new(RawWebSocketReactor {
                upstream_waker: None,
                downstream_waker: None,
                upstream_data_ready: false,
                downstream_data_ready: false,
                upstream_paused: true,   // Start paused
                downstream_paused: true, // Start paused
            }),
        }
    }

    /// Registers a waker for upstream (client→server) state.next() calls.
    pub fn register_upstream_waker(&self, waker: Waker) {
        #[cfg(feature = "debug-logs")]
        log::debug!("WebSocketReactor::register_upstream_waker");
        let mut state = self.state.borrow_mut();
        state.upstream_waker = Some(waker);
    }

    /// Registers a waker for downstream (server→client) state.next() calls.
    pub fn register_downstream_waker(&self, waker: Waker) {
        #[cfg(feature = "debug-logs")]
        log::debug!("WebSocketReactor::register_downstream_waker");
        let mut state = self.state.borrow_mut();
        state.downstream_waker = Some(waker);
    }

    /// Checks if upstream data is ready to be consumed.
    /// Clears the flag after checking (one-shot).
    pub fn poll_upstream_data_ready(&self) -> bool {
        let mut state = self.state.borrow_mut();
        if state.upstream_data_ready {
            #[cfg(feature = "debug-logs")]
            log::debug!("WebSocketReactor::poll_upstream_data_ready: true, clearing flag");
            state.upstream_data_ready = false;
            true
        } else {
            #[cfg(feature = "debug-logs")]
            log::debug!("WebSocketReactor::poll_upstream_data_ready: false");
            false
        }
    }

    /// Sets the upstream data-ready flag.
    pub fn set_upstream_data_ready(&self, ready: bool) {
        let mut state = self.state.borrow_mut();
        state.upstream_data_ready = ready;
    }

    /// Checks if downstream data is ready to be consumed.
    /// Clears the flag after checking (one-shot).
    pub fn poll_downstream_data_ready(&self) -> bool {
        let mut state = self.state.borrow_mut();
        if state.downstream_data_ready {
            #[cfg(feature = "debug-logs")]
            log::debug!("WebSocketReactor::poll_downstream_data_ready: true, clearing flag");
            state.downstream_data_ready = false;
            true
        } else {
            #[cfg(feature = "debug-logs")]
            log::debug!("WebSocketReactor::poll_downstream_data_ready: false");
            false
        }
    }

    /// Sets downstream data ready flag.
    pub fn set_downstream_data_ready(&self, ready: bool) {
        let mut state = self.state.borrow_mut();
        state.downstream_data_ready = ready;
    }

    /// Wakes the upstream (client→server) handler.
    pub fn wake_upstream(&self) {
        #[cfg(feature = "debug-logs")]
        log::debug!("WebSocketReactor::wake_upstream: setting upstream_data_ready=true");
        let mut state = self.state.borrow_mut();
        state.upstream_data_ready = true;
        if let Some(waker) = state.upstream_waker.take() {
            #[cfg(feature = "debug-logs")]
            log::debug!("WebSocketReactor::wake_upstream: waker found, calling wake()");
            waker.wake();
        } else {
            log::debug!(
                "WebSocketReactor::wake_upstream: NO waker registered (handler terminated or not started)"
            );
        }
    }

    /// Wakes the downstream (server→client) handler.
    pub fn wake_downstream(&self) {
        #[cfg(feature = "debug-logs")]
        log::debug!("WebSocketReactor::wake_downstream: setting downstream_data_ready=true");
        let mut state = self.state.borrow_mut();
        state.downstream_data_ready = true;
        if let Some(waker) = state.downstream_waker.take() {
            #[cfg(feature = "debug-logs")]
            log::debug!("WebSocketReactor::wake_downstream: waker found, calling wake()");
            waker.wake();
        } else {
            log::debug!(
                "WebSocketReactor::wake_downstream: NO waker registered (handler terminated or not started)"
            );
        }
    }

    /// Returns true if upstream (request) flow is paused.
    pub fn upstream_paused(&self) -> bool {
        let state = self.state.borrow();
        state.upstream_paused
    }

    /// Sets upstream (request) pause state.
    pub fn set_upstream_paused(&self, paused: bool) {
        #[cfg(feature = "debug-logs")]
        log::debug!("WebSocketReactor::set_upstream_paused({paused})");
        self.state.borrow_mut().upstream_paused = paused;
    }

    /// Returns true if downstream (response) flow is paused.
    pub fn downstream_paused(&self) -> bool {
        let state = self.state.borrow();
        state.downstream_paused
    }

    /// Sets downstream (response) pause state.
    pub fn set_downstream_paused(&self, paused: bool) {
        #[cfg(feature = "debug-logs")]
        log::debug!("WebSocketReactor::set_downstream_paused({paused})");
        self.state.borrow_mut().downstream_paused = paused;
    }
}

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

    fn create_reactor() -> WebSocketReactor {
        WebSocketReactor::new(HttpCid::from(1))
    }

    #[test]
    fn new_reactor_starts_with_paused_state() {
        let reactor = create_reactor();
        assert!(reactor.upstream_paused(), "Upstream should start paused");
    }

    #[test]
    fn set_upstream_paused_changes_state() {
        let reactor = create_reactor();

        reactor.set_upstream_paused(false);
        assert!(
            !reactor.upstream_paused(),
            "Should be unpaused after set_upstream_paused(false)"
        );

        reactor.set_upstream_paused(true);
        assert!(
            reactor.upstream_paused(),
            "Should be paused after set_upstream_paused(true)"
        );
    }

    #[test]
    fn poll_upstream_data_ready_starts_false() {
        let reactor = create_reactor();
        assert!(
            !reactor.poll_upstream_data_ready(),
            "Should start with no data ready"
        );
    }

    #[test]
    fn poll_upstream_data_ready_clears_flag() {
        let reactor = create_reactor();

        // Set flag manually
        reactor.state.borrow_mut().upstream_data_ready = true;

        // First poll returns true and clears
        assert!(
            reactor.poll_upstream_data_ready(),
            "First poll should return true"
        );

        // Second poll returns false (flag cleared)
        assert!(
            !reactor.poll_upstream_data_ready(),
            "Second poll should return false"
        );
    }

    #[test]
    fn downstream_paused_starts_true() {
        let reactor = create_reactor();
        assert!(
            reactor.downstream_paused(),
            "Downstream should start paused"
        );
    }

    #[test]
    fn set_downstream_paused_changes_state() {
        let reactor = create_reactor();

        reactor.set_downstream_paused(false);
        assert!(
            !reactor.downstream_paused(),
            "Should be unpaused after set_downstream_paused(false)"
        );

        reactor.set_downstream_paused(true);
        assert!(
            reactor.downstream_paused(),
            "Should be paused after set_downstream_paused(true)"
        );
    }

    #[test]
    fn poll_downstream_data_ready_starts_false() {
        let reactor = create_reactor();
        assert!(
            !reactor.poll_downstream_data_ready(),
            "Should start with no data ready"
        );
    }

    #[test]
    fn poll_downstream_data_ready_clears_flag() {
        let reactor = create_reactor();

        // Set flag manually
        reactor.state.borrow_mut().downstream_data_ready = true;

        // First poll returns true and clears
        assert!(
            reactor.poll_downstream_data_ready(),
            "First poll should return true"
        );

        // Second poll returns false (flag cleared)
        assert!(
            !reactor.poll_downstream_data_ready(),
            "Second poll should return false"
        );
    }

    #[test]
    fn set_upstream_data_ready_updates_flag() {
        let reactor = create_reactor();

        reactor.set_upstream_data_ready(true);
        assert!(
            reactor.poll_upstream_data_ready(),
            "Flag should be set to true"
        );

        reactor.set_upstream_data_ready(false);
        assert!(
            !reactor.poll_upstream_data_ready(),
            "Flag should be set to false"
        );
    }

    #[test]
    fn set_downstream_data_ready_updates_flag() {
        let reactor = create_reactor();

        reactor.set_downstream_data_ready(true);
        assert!(
            reactor.poll_downstream_data_ready(),
            "Flag should be set to true"
        );

        reactor.set_downstream_data_ready(false);
        assert!(
            !reactor.poll_downstream_data_ready(),
            "Flag should be set to false"
        );
    }

    #[test]
    fn wake_upstream_sets_data_ready() {
        let reactor = create_reactor();

        // Initially false
        assert!(!reactor.poll_upstream_data_ready());

        // Wake should set the flag
        reactor.wake_upstream();

        // Now should be true
        assert!(
            reactor.poll_upstream_data_ready(),
            "wake_upstream should set data_ready flag"
        );
    }

    #[test]
    fn wake_downstream_sets_data_ready() {
        let reactor = create_reactor();

        // Initially false
        assert!(!reactor.poll_downstream_data_ready());

        // Wake should set the flag
        reactor.wake_downstream();

        // Now should be true
        assert!(
            reactor.poll_downstream_data_ready(),
            "wake_downstream should set data_ready flag"
        );
    }

    #[test]
    fn register_upstream_waker_stores_waker() {
        use std::sync::{Arc, Mutex};
        use std::task::Wake;

        struct TestWaker {
            woken: Arc<Mutex<bool>>,
        }

        impl Wake for TestWaker {
            fn wake(self: Arc<Self>) {
                *self.woken.lock().unwrap() = true;
            }
        }

        let reactor = create_reactor();
        let woken = Arc::new(Mutex::new(false));
        let test_waker = Arc::new(TestWaker {
            woken: woken.clone(),
        });
        let waker = test_waker.clone().into();

        // Register the waker
        reactor.register_upstream_waker(waker);

        // Wake should call the waker
        reactor.wake_upstream();

        // Check that waker was called
        assert!(*woken.lock().unwrap(), "Waker should have been called");
    }

    #[test]
    fn register_downstream_waker_stores_waker() {
        use std::sync::{Arc, Mutex};
        use std::task::Wake;

        struct TestWaker {
            woken: Arc<Mutex<bool>>,
        }

        impl Wake for TestWaker {
            fn wake(self: Arc<Self>) {
                *self.woken.lock().unwrap() = true;
            }
        }

        let reactor = create_reactor();
        let woken = Arc::new(Mutex::new(false));
        let test_waker = Arc::new(TestWaker {
            woken: woken.clone(),
        });
        let waker = test_waker.clone().into();

        // Register the waker
        reactor.register_downstream_waker(waker);

        // Wake should call the waker
        reactor.wake_downstream();

        // Check that waker was called
        assert!(*woken.lock().unwrap(), "Waker should have been called");
    }

    #[test]
    fn wake_without_waker_does_not_panic() {
        let reactor = create_reactor();

        // Should not panic even without a registered waker
        reactor.wake_upstream();
        reactor.wake_downstream();

        // Data ready flags should still be set
        assert!(reactor.poll_upstream_data_ready());
        assert!(reactor.poll_downstream_data_ready());
    }
}