pdk-classy 1.9.1-alpha.2

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

use std::{cell::RefCell, marker::PhantomData, rc::Rc};

use crate::event::{After, BodyEvent, EventData, HeadersAccessor, HeadersEvent};

#[cfg(feature = "enable_stop_iteration")]
use super::headers_body::HeadersBodyExchange;
use super::{
    body::BodyExchange,
    body_stream::BodyStreamExchange,
    dynamic_exchange::{DynamicExchange, ExchangeEvent},
    entity::HeadersHandler,
};

pub struct HeadersExchange<H, B> {
    pub(super) contains_body: bool,
    exchange: Rc<RefCell<DynamicExchange>>,
    _events: PhantomData<(H, B)>,
}

impl<H, B> HeadersExchange<H, B>
where
    H: HeadersEvent + ExchangeEvent,
    B: BodyEvent + After<H> + ExchangeEvent,
{
    pub(super) fn new(exchange: Rc<RefCell<DynamicExchange>>, contains_body: bool) -> Self {
        Self {
            contains_body,
            exchange,
            _events: PhantomData,
        }
    }
}

#[cfg(not(feature = "enable_stop_iteration"))]
macro_rules! with_event_data {
    ($self:expr, $method:ident $(, $arg:expr)*) => {
        $self.with_headers(|h| h.$method($($arg.clone()),*))
        .expect("Exchange should be in header or header event")
    };
}

#[cfg(feature = "enable_stop_iteration")]
macro_rules! with_body_fallback {
    ($self:expr, $method:ident $(, $arg:expr)*) => {
        if let Some(res) = $self.with_headers(|h| h.$method($($arg),*)) {
            res
        } else {
            $self.with_body(|h| h.$method($($arg),*))
                .expect("Exchange should be in header or body event")
        }
    };
}

#[cfg(feature = "enable_stop_iteration")]
macro_rules! with_body_fallback_cloned {
    ($self:expr, $method:ident $(, $arg:expr)*) => {
        if let Some(res) = $self.with_headers(|h| h.$method($($arg.clone()),*)) {
            res
        } else {
            $self.with_body(|h| h.$method($($arg),*))
                .expect("Exchange should be in header or body event")
        }
    };
}

impl<H, B> HeadersExchange<H, B>
where
    H: HeadersEvent + ExchangeEvent,
    B: BodyEvent + After<H> + ExchangeEvent,
{
    #[cfg(feature = "enable_stop_iteration")]
    pub fn with_body<F, T>(&self, op: F) -> Option<T>
    where
        for<'e> F: FnOnce(&crate::event::EventDataStream<'e, H>) -> T,
    {
        self.exchange
            .borrow()
            .get::<H>()
            .map(|exchange| op(&exchange.event_data_stream()))
    }

    pub fn with_headers<F, T>(&self, op: F) -> Option<T>
    where
        for<'f> F: FnOnce(&EventData<'f, H>) -> T,
    {
        self.exchange
            .borrow()
            .get::<H>()
            .and_then(|exchange| exchange.event_data())
            .map(|event_data| op(&event_data))
    }
}

#[cfg(feature = "enable_stop_iteration")]
impl<H, B> HeadersExchange<H, B>
where
    H: HeadersEvent + ExchangeEvent,
    B: BodyEvent + After<H> + ExchangeEvent + 'static,
    for<'e> EventData<'e, H>: HeadersAccessor,
    for<'f> crate::event::EventDataStream<'f, H>: HeadersAccessor,
{
    pub async fn into_headers_body_state(self) -> HeadersBodyExchange<B, H> {
        HeadersBodyExchange::new(self.exchange, self.contains_body).await
    }

    pub fn handler(&self) -> &dyn HeadersHandler {
        self
    }
}

impl<H, B> HeadersExchange<H, B>
where
    H: HeadersEvent + ExchangeEvent,
    B: BodyEvent + After<H> + ExchangeEvent + 'static,
    for<'e> EventData<'e, H>: HeadersAccessor,
{
    pub async fn into_body_state(self) -> BodyExchange<B> {
        BodyExchange::new(self.exchange, self.contains_body).await
    }

    pub async fn into_body_stream_state(self) -> BodyStreamExchange<B> {
        BodyStreamExchange::new(self.exchange, self.contains_body).await
    }

    #[cfg(not(feature = "enable_stop_iteration"))]
    pub fn handler(&self) -> &dyn HeadersHandler {
        self
    }

    pub fn send_response(&self, status_code: u32, headers: Vec<(&str, &str)>, body: Option<&[u8]>) {
        if let Some(reactor) = self.exchange.borrow().get_reactor().as_ref() {
            reactor.cancel_response()
        }

        self.exchange
            .borrow_mut()
            .send_response(status_code, headers, body);
    }
}

#[cfg(not(feature = "enable_stop_iteration"))]
impl<H, B> HeadersHandler for HeadersExchange<H, B>
where
    H: HeadersEvent + ExchangeEvent,
    B: BodyEvent + After<H> + ExchangeEvent,
    for<'e> EventData<'e, H>: HeadersAccessor,
{
    fn headers(&self) -> Vec<(String, String)> {
        with_event_data!(self, headers)
    }

    fn header(&self, name: &str) -> Option<String> {
        with_event_data!(self, header, name)
    }

    fn add_header(&self, name: &str, value: &str) {
        with_event_data!(self, add_header, name, value)
    }

    fn set_header(&self, name: &str, value: &str) {
        with_event_data!(self, set_header, name, value)
    }

    fn set_headers(&self, headers: Vec<(&str, &str)>) {
        with_event_data!(self, set_headers, headers)
    }

    fn remove_header(&self, name: &str) {
        with_event_data!(self, remove_header, name)
    }
}

#[cfg(feature = "enable_stop_iteration")]
impl<H, B> HeadersHandler for HeadersExchange<H, B>
where
    H: HeadersEvent + ExchangeEvent,
    B: BodyEvent + After<H> + ExchangeEvent,
    for<'e> EventData<'e, H>: HeadersAccessor,
    for<'f> crate::event::EventDataStream<'f, H>: HeadersAccessor,
{
    fn headers(&self) -> Vec<(String, String)> {
        with_body_fallback!(self, headers)
    }

    fn header(&self, name: &str) -> Option<String> {
        if let Some(res) = self.with_headers(|h| h.header(name)) {
            res
        } else {
            self.with_body(|h| h.header(name))
                .expect("Exchange should be in header or body event")
        }
    }

    fn add_header(&self, name: &str, value: &str) {
        with_body_fallback!(self, add_header, name, value)
    }

    fn set_header(&self, name: &str, value: &str) {
        with_body_fallback!(self, set_header, name, value)
    }

    fn set_headers(&self, headers: Vec<(&str, &str)>) {
        with_body_fallback_cloned!(self, set_headers, headers)
    }

    fn remove_header(&self, name: &str) {
        with_body_fallback!(self, remove_header, name)
    }
}