dataplane-sdk 0.1.1

//  Copyright (c) 2026 Metaform Systems, Inc
//
//  This program and the accompanying materials are made available under the
//  terms of the Apache License, Version 2.0 which is available at
//  https://www.apache.org/licenses/LICENSE-2.0
//
//    SPDX-License-Identifier: Apache-2.0
//
//    Contributors:
//         Metaform Systems, Inc. - initial API and implementation
//

use crate::{
    core::{
        db::{
            control_plane::ControlPlaneRepo,
            data_flow::DataFlowRepo,
            tx::{Transaction, TransactionalContext},
        },
        error::{DbError, HandlerError},
        handler::DataFlowHandler,
        model::{
            data_address::DataAddress,
            data_flow::{DataFlow, DataFlowState, DataFlowType, TransitionError},
            messages::{
                DataFlowPrepareMessage, DataFlowResumeMessage, DataFlowStartMessage,
                DataFlowStartedNotificationMessage, DataFlowStatusMessage,
                DataFlowStatusResponseMessage,
            },
        },
    },
    error::{SdkError, SdkResult},
};

pub struct DataPlaneSdkInternal<C>
where
    C: TransactionalContext,
{
    pub(crate) ctx: C,
    pub(crate) repo: Box<dyn DataFlowRepo<Transaction = C::Transaction>>,
    pub(crate) control_plane_repo: Box<dyn ControlPlaneRepo<Transaction = C::Transaction>>,
    pub(crate) handler: Box<dyn DataFlowHandler<Transaction = C::Transaction>>,
    pub(crate) client: reqwest::Client,
}

impl<C> DataPlaneSdkInternal<C>
where
    C: TransactionalContext,
{
    pub async fn start(
        &self,
        participant_context_id: &str,
        control_plane_id: &str,
        req: DataFlowStartMessage,
    ) -> SdkResult<DataFlowStatusMessage> {
        let mut flow = DataFlow::builder()
            .id(req.process_id)
            .counter_party_id(req.counter_party_id)
            .maybe_data_address(req.data_address)
            .participant_context_id(participant_context_id)
            .state(DataFlowState::Initiating)
            .metadata(req.metadata)
            .claims(req.claims)
            .participant_id(req.participant_id)
            .dataspace_context(req.dataspace_context)
            .dataset_id(req.dataset_id)
            .agreement_id(req.agreement_id)
            .control_plane_id(control_plane_id)
            .labels(req.labels)
            .profile(req.profile)
            .kind(DataFlowType::Provider)
            .build();

        if self.handler.can_handle(&flow).await? {
            let mut tx = self.ctx.begin().await?;
            let response = self.handler.on_start(&mut tx, &flow).await?;

            match response.state {
                DataFlowState::Starting => flow.transition_to_starting()?,
                DataFlowState::Started => flow.transition_to_started()?,
                _ => {
                    return Err(SdkError::Handler(HandlerError::NotSupported(
                        "Handler can only transition to Starting or Started state".to_string(),
                    )));
                }
            }

            self.repo.create(&mut tx, &flow).await?;
            tx.commit().await?;

            Ok(response)
        } else {
            Err(SdkError::Handler(HandlerError::NotSupported(
                "Data flow handler cannot handle this flow".to_string(),
            )))
        }
    }

    pub async fn prepare(
        &self,
        participant_context_id: &str,
        control_plane_id: &str,
        req: DataFlowPrepareMessage,
    ) -> SdkResult<DataFlowStatusMessage> {
        let mut flow = DataFlow::builder()
            .id(req.process_id)
            .counter_party_id(req.counter_party_id)
            .participant_context_id(participant_context_id)
            .state(DataFlowState::Initiating)
            .metadata(req.metadata)
            .claims(req.claims)
            .participant_id(req.participant_id)
            .dataspace_context(req.dataspace_context)
            .dataset_id(req.dataset_id)
            .agreement_id(req.agreement_id)
            .control_plane_id(control_plane_id)
            .labels(req.labels)
            .profile(req.profile)
            .kind(DataFlowType::Consumer)
            .build();

        if self.handler.can_handle(&flow).await? {
            let mut tx = self.ctx.begin().await?;
            let response = self.handler.on_prepare(&mut tx, &flow).await?;

            match response.state {
                DataFlowState::Preparing => flow.transition_to_preparing()?,
                DataFlowState::Prepared => flow.transition_to_prepared()?,
                _ => {
                    return Err(SdkError::Handler(HandlerError::NotSupported(format!(
                        "Handler can only transition to Preparing or Prepared state: current state {:?}",
                        response.state
                    ))));
                }
            }
            self.repo.create(&mut tx, &flow).await?;
            tx.commit().await?;

            Ok(response)
        } else {
            Err(SdkError::Handler(HandlerError::NotSupported(
                "Data flow handler cannot handle this flow".to_string(),
            )))
        }
    }

    pub async fn terminate(
        &self,
        _ctx: &str,
        flow_id: &str,
        reason: Option<String>,
    ) -> SdkResult<()> {
        let mut tx = self.ctx.begin().await?;
        let mut flow = self
            .repo
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        flow.transition_to_terminated(reason)?;
        self.repo.update(&mut tx, &flow).await?;

        self.handler.on_terminate(&mut tx, &flow).await?;

        tx.commit().await?;
        Ok(())
    }

    pub async fn started(
        &self,
        _ctx: &str,
        flow_id: &str,
        msg: DataFlowStartedNotificationMessage,
    ) -> SdkResult<()> {
        let mut tx = self.ctx.begin().await?;
        let mut flow = self
            .repo
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        flow.data_address = msg.data_address;

        self.handler.on_started(&mut tx, &flow).await?;

        flow.transition_to_started()?;
        self.repo.update(&mut tx, &flow).await?;

        tx.commit().await?;
        Ok(())
    }

    pub async fn completed(&self, _ctx: &str, flow_id: &str) -> SdkResult<()> {
        let mut tx = self.ctx.begin().await?;
        let mut flow = self
            .repo
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        flow.transition_to_completed()?;

        self.handler.on_completed(&mut tx, &flow).await?;

        self.repo.update(&mut tx, &flow).await?;

        tx.commit().await?;
        Ok(())
    }

    pub async fn suspend(
        &self,
        _ctx: &str,
        flow_id: &str,
        reason: Option<String>,
    ) -> SdkResult<()> {
        let mut tx = self.ctx.begin().await?;

        let mut flow = self
            .repo
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        flow.transition_to_suspended(reason)?;
        self.repo.update(&mut tx, &flow).await?;

        self.handler.on_suspend(&mut tx, &flow).await?;

        tx.commit().await?;

        Ok(())
    }

    pub async fn resume(
        &self,
        _ctx: &str,
        flow_id: &str,
        msg: DataFlowResumeMessage,
    ) -> SdkResult<DataFlowStatusMessage> {
        let mut tx = self.ctx.begin().await?;

        let mut flow = self
            .repo
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        flow.data_address = msg.data_address;

        let response = self.handler.on_resume(&mut tx, &flow).await?;

        match response.state {
            DataFlowState::Starting => flow.transition_to_starting()?,
            DataFlowState::Started => flow.transition_to_started()?,
            _ => {
                return Err(SdkError::Handler(HandlerError::NotSupported(
                    "Handler can only transition to Starting or Started state".to_string(),
                )));
            }
        }

        self.repo.update(&mut tx, &flow).await?;

        tx.commit().await?;

        Ok(response)
    }
    pub async fn status(
        &self,
        _ctx: &str,
        flow_id: &str,
    ) -> SdkResult<DataFlowStatusResponseMessage> {
        let mut tx = self.ctx.begin().await?;

        let flow = self
            .repo
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        tx.commit().await?;

        Ok(DataFlowStatusResponseMessage::builder()
            .data_flow_id(flow.id)
            .state(flow.state)
            .build())
    }

    /// Notifies the control plane that the data flow has been prepared, by
    /// POSTing to the URL of the flow's referenced control plane (resolved from
    /// `control_plane_id`). See the Data Plane Signaling "Control Plane
    /// Endpoint" section.
    pub async fn notify_prepared(
        &self,
        ctx: &str,
        flow_id: &str,
        data_address: Option<DataAddress>,
    ) -> SdkResult<()> {
        self.send_callback(
            ctx,
            flow_id,
            "prepared",
            data_address.clone(),
            None,
            move |flow| {
                flow.data_address = data_address;
                flow.transition_to_prepared()
            },
        )
        .await
    }

    /// Notifies the control plane that the data flow has started.
    pub async fn notify_started(
        &self,
        ctx: &str,
        flow_id: &str,
        data_address: Option<DataAddress>,
    ) -> SdkResult<()> {
        self.send_callback(
            ctx,
            flow_id,
            "started",
            data_address.clone(),
            None,
            |flow| {
                flow.data_address = data_address;
                flow.transition_to_started()
            },
        )
        .await
    }

    /// Notifies the control plane that the data flow has completed.
    pub async fn notify_completed(&self, ctx: &str, flow_id: &str) -> SdkResult<()> {
        self.send_callback(ctx, flow_id, "completed", None, None, |flow| {
            flow.transition_to_completed()
        })
        .await
    }

    /// Notifies the control plane that the data flow has errored. The optional
    /// `error` is forwarded as the `error` field of the status message.
    pub async fn notify_errored(
        &self,
        ctx: &str,
        flow_id: &str,
        error: Option<String>,
    ) -> SdkResult<()> {
        self.send_callback(ctx, flow_id, "errored", None, error.clone(), move |flow| {
            flow.transition_to_terminated(error)
        })
        .await
    }

    async fn send_callback<CB>(
        &self,
        _ctx: &str,
        flow_id: &str,
        operation: &str,
        data_address: Option<DataAddress>,
        error: Option<String>,
        op: CB,
    ) -> SdkResult<()>
    where
        CB: FnOnce(&mut DataFlow) -> Result<(), TransitionError>,
    {
        let mut tx = self.ctx.begin().await?;

        let mut flow = self
            .fetch_by_id(&mut tx, flow_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;

        let control_plane = self
            .control_plane_repo
            .fetch_by_id(&mut tx, &flow.control_plane_id)
            .await?
            .ok_or_else(|| DbError::NotFound(flow.control_plane_id.clone()))?;

        op(&mut flow)?;

        let msg = DataFlowStatusMessage::builder()
            .data_flow_id(flow.id.clone())
            .maybe_data_address(data_address)
            .state(flow.state.clone())
            .maybe_error(error)
            .build();

        let url = format!(
            "{}/transfers/{}/dataflow/{}",
            control_plane.url.trim_end_matches('/'),
            flow.id,
            operation
        );

        let resp = self.client.post(&url).json(&msg).send().await?;
        if !resp.status().is_success() {
            let status = resp.status().as_u16();
            let body = resp.text().await.unwrap_or_default();
            return Err(SdkError::NotificationStatus { status, body });
        }

        self.repo.update(&mut tx, &flow).await?;
        tx.commit().await?;

        Ok(())
    }

    pub async fn fetch_by_id(
        &self,
        tx: &mut C::Transaction,
        flow_id: &str,
    ) -> SdkResult<Option<DataFlow>> {
        self.repo.fetch_by_id(tx, flow_id).await.map(Ok)?
    }

    pub fn ctx(&self) -> &C {
        &self.ctx
    }
}