pdk-classy 1.3.0-alpha.0

// Copyright 2023 Salesforce, Inc. All rights reserved.
use std::{
    cell::RefCell,
    future::{ready, Ready},
    rc::Rc,
};

use crate::{
    event::{EventKind, ResponseBody, ResponseHeaders},
    BoxFuture,
};

use super::{
    body::BodyExchange, dynamic_exchange::DynamicExchange, entity::IntoBodyState,
    headers::HeadersExchange, BodyState, EntityState, HeadersState,
};

/// Initial state of the Response state-machine.
/// This type is intended to be injected as a parameter in Response filters.
pub struct ResponseState {
    contains_body: bool,
    exchange: Rc<RefCell<DynamicExchange>>,
}

/// Error that represents an invalid state in the Response state-machine.
#[derive(thiserror::Error, Debug)]
pub enum InvalidResponseState {
    /// Unexpected low-level event.
    #[error("Invalid response state: {0:?}")]
    Event(EventKind),

    /// Unexpected early response.
    #[error("Early response")]
    EarlyResponse,
}

impl ResponseState {
    #[allow(clippy::await_holding_refcell_ref)]
    pub(crate) async fn new(
        exchange: Rc<RefCell<DynamicExchange>>,
    ) -> Result<Self, InvalidResponseState> {
        let contains_body;
        {
            let mut e = exchange.borrow_mut();
            let Some(exchange) = e.wait_for_event::<ResponseHeaders>().await else {
                let error = exchange
                    .borrow()
                    .current_event()
                    .map(InvalidResponseState::Event)
                    .unwrap_or(InvalidResponseState::EarlyResponse);
                return Err(error);
            };
            contains_body = !exchange
                .event_data()
                .expect("Must contain headers")
                .event
                .end_of_stream;
        }

        Ok(Self {
            contains_body,
            exchange,
        })
    }
}

impl IntoBodyState for ResponseState {
    type BodyState = ResponseBodyState;
    type BodyStateFuture = BoxFuture<'static, Self::BodyState>;

    fn into_body_state(self) -> Self::BodyStateFuture {
        Box::pin(async {
            ResponseBodyState {
                inner: BodyExchange::new(self.exchange, self.contains_body).await,
            }
        })
    }
}

impl EntityState for ResponseState {
    type HeadersState = ResponseHeadersState;
    type HeadersStateFuture = Ready<Self::HeadersState>;

    fn into_headers_state(self) -> Self::HeadersStateFuture {
        ready(ResponseHeadersState {
            inner: HeadersExchange::new(self.exchange, self.contains_body),
        })
    }
}

/// Headers state of the Response state-machine.
pub struct ResponseHeadersState {
    inner: HeadersExchange<ResponseHeaders, ResponseBody>,
}

impl ResponseHeadersState {
    pub fn status_code(&self) -> u32 {
        self.inner.with_event_data(|e| e.status_code())
    }
}

impl IntoBodyState for ResponseHeadersState {
    type BodyState = ResponseBodyState;

    type BodyStateFuture = BoxFuture<'static, Self::BodyState>;

    fn into_body_state(self) -> Self::BodyStateFuture {
        Box::pin(async {
            ResponseBodyState {
                inner: self.inner.into_body_state().await,
            }
        })
    }
}

impl HeadersState for ResponseHeadersState {
    fn handler(&self) -> &dyn super::HeadersHandler {
        self.inner.handler()
    }

    fn contains_body(&self) -> bool {
        self.inner.contains_body
    }
}

/// Body state of the Response state-machine.
pub struct ResponseBodyState {
    inner: BodyExchange<ResponseBody>,
}

impl BodyState for ResponseBodyState {
    fn handler(&self) -> &dyn super::BodyHandler {
        &self.inner
    }

    fn contains_body(&self) -> bool {
        self.inner.contains_body()
    }
}