pdk-classy 1.9.0-alpha.3

// 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<S = ()> {
    parent: Rc<FilterContext>,
    exchange: Rc<RefCell<DynamicExchange>>,
    state: Rc<S>,
}

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

    pub(super) fn state(&self) -> &Rc<S> {
        &self.state
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    pub(super) fn state(&self) -> &Rc<S> {
        &self.state
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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