pdk-classy 1.8.0

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

use std::{
    cell::{Cell, RefCell},
    rc::Rc,
};

use crate::{
    context,
    extract::{
        context::FilterContext, extractability, AlreadyExtracted, Exclusive, FromContext,
        FromContextOnce,
    },
    BoxFuture,
};

use super::{
    dynamic_exchange::DynamicExchange, entity::EntityState, request::InvalidRequestState,
    response::InvalidResponseState, IntoBodyState, RequestBodyState, RequestData,
    RequestHeadersState, RequestState, ResponseBodyState, ResponseHeadersState, ResponseState,
};

use super::{IntoBodyStreamState, RequestBodyStreamState, ResponseBodyStreamState};

pub struct RequestContext {
    parent: Rc<FilterContext>,
    exchange: Rc<RefCell<DynamicExchange>>,
}

impl RequestContext {
    pub(super) fn new(parent: Rc<FilterContext>, exchange: Rc<RefCell<DynamicExchange>>) -> Self {
        Self { parent, exchange }
    }

    fn parent(&self) -> &FilterContext {
        &self.parent
    }
}

context!(RequestContext => FilterContext {RequestContext::parent});

impl FromContextOnce<RequestContext> for RequestState {
    type Error = InvalidRequestState;

    type Future<'c> = BoxFuture<'c, Result<Self, Self::Error>>;

    fn from_context_once(context: Exclusive<RequestContext>) -> Self::Future<'_> {
        Box::pin(RequestState::new(context.exchange.clone()))
    }
}

impl FromContextOnce<RequestContext> for RequestHeadersState {
    type Error = InvalidRequestState;

    type Future<'c> = BoxFuture<'c, Result<Self, Self::Error>>;

    fn from_context_once(context: Exclusive<RequestContext>) -> Self::Future<'_> {
        Box::pin(async {
            Ok(RequestState::from_context_once(context)
                .await?
                .into_headers_state()
                .await)
        })
    }
}

impl FromContextOnce<RequestContext> for RequestBodyState {
    type Error = InvalidRequestState;

    type Future<'c> = BoxFuture<'c, Result<Self, Self::Error>>;

    fn from_context_once(context: Exclusive<RequestContext>) -> Self::Future<'_> {
        Box::pin(async {
            Ok(RequestState::from_context_once(context)
                .await?
                .into_body_state()
                .await)
        })
    }
}

impl FromContextOnce<RequestContext> for RequestBodyStreamState {
    type Error = InvalidRequestState;

    type Future<'c> = BoxFuture<'c, Result<Self, Self::Error>>;

    fn from_context_once(context: Exclusive<RequestContext>) -> Self::Future<'_> {
        Box::pin(async {
            Ok(RequestState::from_context_once(context)
                .await?
                .into_body_stream_state()
                .await)
        })
    }
}

pub struct ResponseContext<D> {
    parent: Rc<FilterContext>,
    exchange: Rc<RefCell<DynamicExchange>>,
    request_data: Cell<Option<RequestData<D>>>,
}

impl<D> ResponseContext<D> {
    pub(super) fn new(
        parent: Rc<FilterContext>,
        exchange: Rc<RefCell<DynamicExchange>>,
        request_data: RequestData<D>,
    ) -> Self {
        Self {
            parent,
            exchange,
            request_data: Cell::new(Some(request_data)),
        }
    }

    fn parent(&self) -> &FilterContext {
        &self.parent
    }
}

context!(<D> ResponseContext<D> => FilterContext {ResponseContext::parent});

impl<D> FromContextOnce<ResponseContext<D>> for ResponseState {
    type Error = InvalidResponseState;

    type Future<'c>
        = BoxFuture<'c, Result<Self, Self::Error>>
    where
        D: 'c;

    fn from_context_once(context: Exclusive<ResponseContext<D>>) -> Self::Future<'_> {
        Box::pin(ResponseState::new(context.exchange.clone()))
    }
}

impl<D> FromContextOnce<ResponseContext<D>> for ResponseHeadersState {
    type Error = InvalidResponseState;

    type Future<'c>
        = BoxFuture<'c, Result<Self, Self::Error>>
    where
        D: 'c;

    fn from_context_once(context: Exclusive<ResponseContext<D>>) -> Self::Future<'_> {
        Box::pin(async {
            Ok(ResponseState::from_context_once(context)
                .await?
                .into_headers_state()
                .await)
        })
    }
}

impl<D> FromContextOnce<ResponseContext<D>> for ResponseBodyState {
    type Error = InvalidResponseState;

    type Future<'c>
        = BoxFuture<'c, Result<Self, Self::Error>>
    where
        D: 'c;

    fn from_context_once(context: Exclusive<ResponseContext<D>>) -> Self::Future<'_> {
        Box::pin(async {
            Ok(ResponseState::from_context_once(context)
                .await?
                .into_body_state()
                .await)
        })
    }
}

impl<D> FromContextOnce<ResponseContext<D>> for ResponseBodyStreamState {
    type Error = InvalidResponseState;

    type Future<'c>
        = BoxFuture<'c, Result<Self, Self::Error>>
    where
        D: 'c;

    fn from_context_once(context: Exclusive<ResponseContext<D>>) -> Self::Future<'_> {
        Box::pin(async {
            Ok(ResponseState::from_context_once(context)
                .await?
                .into_body_stream_state()
                .await)
        })
    }
}

impl<D> FromContext<ResponseContext<D>, extractability::Transitive> for RequestData<D>
where
    D: 'static,
{
    type Error = AlreadyExtracted<RequestData<D>>;

    fn from_context(context: &ResponseContext<D>) -> Result<Self, Self::Error> {
        context
            .request_data
            .take()
            .ok_or_else(AlreadyExtracted::default)
    }
}