staircase 0.0.7

use std::{borrow::Cow, future::Future, marker::PhantomData, sync::Arc};

use k8s_openapi::jiff::Timestamp;
use kube::{Client, core::object::HasStatus};

/// Shared, immutable facts for one reconcile run.
pub struct RunContext<R> {
    // Note: evaluate restricting this to a &R
    /// Resource currently being reconciled.
    pub resource:  Arc<R>,
    /// Kubernetes client used by this run.
    pub client:    Client,
    /// Timestamp captured before the first step starts.
    ///
    /// Use this instead of reading the current time in individual steps when
    /// all decisions in one run should share the same time reference.
    pub run_start: Timestamp,
}

#[expect(type_alias_bounds)]
pub type MutableStepResult<R: HasStatus, D, E> = Result<MutableStepOutcome<<R as HasStatus>::Status, D>, E>;

pub type StepResult<D, E> = Result<StepOutcome<D>, E>;

pub trait StepMode<R: HasStatus, D>: Send + Sync + 'static {
    type Outcome: Send + 'static;

    fn into_mutable_outcome(outcome: Self::Outcome) -> MutableStepOutcome<R::Status, D>;
}

pub struct Immutable;

impl<R, D> StepMode<R, D> for Immutable
where
    R: HasStatus,
    D: Send + 'static,
{
    type Outcome = StepOutcome<D>;

    fn into_mutable_outcome(outcome: Self::Outcome) -> MutableStepOutcome<R::Status, D> {
        match outcome {
            StepOutcome::Stop => MutableStepOutcome::Stop,
            StepOutcome::NoModification { data } => MutableStepOutcome::NoModification { data },
        }
    }
}

pub struct Mutable;

impl<R, D> StepMode<R, D> for Mutable
where
    R: HasStatus<Status: Send + 'static>,
    D: Send + 'static,
{
    type Outcome = MutableStepOutcome<R::Status, D>;

    fn into_mutable_outcome(outcome: Self::Outcome) -> MutableStepOutcome<R::Status, D> { outcome }
}

/// Main trait of the step-based operator approach.
///
/// A [`Step`] should check one fix that may bring current state closer to the
/// desired state. If the fix is missing, the step may apply that one meaningful
/// modification and return [`MutableStepOutcome::Modified`]. The current run
/// then stops so the next run can start from freshly observed state.
///
/// - If this mutable step applied a change, return
///   [`MutableStepOutcome::Modified`] with an optional status update.
/// - If this step did not make any change, return
///   [`StepOutcome::NoModification`] and pass data to the next step.
/// - If the step does not yet have all the information to make a reasonable
///   decision, return [`StepOutcome::Stop`].
///
/// Each step specifies the data it consumes via [`Step::InData`] and the data
/// it produces via [`Step::OutData`]. Use this to pass cluster state through
/// the **Run** without fetching the same external resource twice.
///
/// Immutable steps are the default read-only shape and return [`StepOutcome`].
/// They cannot return [`MutableStepOutcome::Modified`], so they should never
/// modify Kubernetes or external systems. Immutable steps can be used in
/// parallel to gather data or evaluate conditions.
///
/// See the [crate-level documentation](crate) for the broader reconciliation
/// model.
pub trait Step {
    /// Whether this step is immutable or may modify external state.
    type Mode: StepMode<Self::Resource, Self::OutData>;
    /// Data consumed by this step.
    type InData: Send + 'static;
    /// Data produced for the next step when this step does not modify anything.
    type OutData: Send + 'static;
    /// Resource being reconciled. It must expose a status type.
    type Resource: HasStatus;
    /// Step-specific error type.
    type Error: Send + 'static;

    /// Executes this step.
    #[expect(clippy::type_complexity)]
    fn run<'a>(
        &'a self,
        context: &'a RunContext<Self::Resource>,
        data: Self::InData,
    ) -> impl Future<Output = Result<<Self::Mode as StepMode<Self::Resource, Self::OutData>>::Outcome, Self::Error>>
    + Send
    + 'a;

    /// Override this to specify a name for tracing and metrics.
    fn traced_name(&self) -> Cow<'static, str> { Cow::Borrowed(std::any::type_name::<Self>()) }
}

pub trait ImmutableStep: Step<Mode = Immutable> {}

impl<S> ImmutableStep for S where S: Step<Mode = Immutable> {}

pub trait MutableStep: Step<Mode = Mutable> {}

impl<S> MutableStep for S where S: Step<Mode = Mutable> {}

type StepFnMarker<R, InData, OutData, Error, Mode> = fn() -> (R, InData, OutData, Error, Mode);

fn owned_context<R>(context: &RunContext<R>) -> RunContext<R> {
    RunContext {
        resource:  context.resource.clone(),
        client:    context.client.clone(),
        run_start: context.run_start,
    }
}

/// Closure-backed [`Step`] returned by [`immutable_step`] and [`mutable_step`].
pub struct StepFn<R, InData, OutData, Error, Mode, F> {
    name:    Cow<'static, str>,
    f:       F,
    _marker: PhantomData<StepFnMarker<R, InData, OutData, Error, Mode>>,
}

/// Builds an immutable step from an async closure.
///
/// The closure receives a cloned [`RunContext`] by value. Cloning the context
/// clones the resource [`Arc`] and Kubernetes client handle, so closure futures
/// can use `async move` without borrowing from the executor.
pub fn immutable_step<R, InData, OutData, Error, F, Fut>(
    name: impl Into<Cow<'static, str>>,
    f: F,
) -> StepFn<R, InData, OutData, Error, Immutable, F>
where
    R: HasStatus,
    InData: Send + 'static,
    OutData: Send + 'static,
    Error: Send + 'static,
    F: Fn(RunContext<R>, InData) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = StepResult<OutData, Error>> + Send + 'static,
{
    StepFn {
        name: name.into(),
        f,
        _marker: PhantomData,
    }
}

/// Builds a mutable step from an async closure.
///
/// The closure receives a cloned [`RunContext`] by value. Cloning the context
/// clones the resource [`Arc`] and Kubernetes client handle, so closure futures
/// can use `async move` without borrowing from the executor.
pub fn mutable_step<R, InData, OutData, Error, F, Fut>(
    name: impl Into<Cow<'static, str>>,
    f: F,
) -> StepFn<R, InData, OutData, Error, Mutable, F>
where
    R: HasStatus<Status: Send + 'static>,
    InData: Send + 'static,
    OutData: Send + 'static,
    Error: Send + 'static,
    F: Fn(RunContext<R>, InData) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = MutableStepResult<R, OutData, Error>> + Send + 'static,
{
    StepFn {
        name: name.into(),
        f,
        _marker: PhantomData,
    }
}

impl<R, InData, OutData, Error, F, Fut> Step for StepFn<R, InData, OutData, Error, Immutable, F>
where
    R: HasStatus,
    InData: Send + 'static,
    OutData: Send + 'static,
    Error: Send + 'static,
    F: Fn(RunContext<R>, InData) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = StepResult<OutData, Error>> + Send + 'static,
{
    type Error = Error;
    type InData = InData;
    type Mode = Immutable;
    type OutData = OutData;
    type Resource = R;

    fn run<'a>(
        &'a self,
        context: &'a RunContext<Self::Resource>,
        data: Self::InData,
    ) -> impl Future<Output = StepResult<Self::OutData, Self::Error>> + Send + 'a {
        let context = owned_context(context);
        (self.f)(context, data)
    }

    fn traced_name(&self) -> Cow<'static, str> { self.name.clone() }
}

impl<R, InData, OutData, Error, F, Fut> Step for StepFn<R, InData, OutData, Error, Mutable, F>
where
    R: HasStatus<Status: Send + 'static>,
    InData: Send + 'static,
    OutData: Send + 'static,
    Error: Send + 'static,
    F: Fn(RunContext<R>, InData) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = MutableStepResult<R, OutData, Error>> + Send + 'static,
{
    type Error = Error;
    type InData = InData;
    type Mode = Mutable;
    type OutData = OutData;
    type Resource = R;

    fn run<'a>(
        &'a self,
        context: &'a RunContext<Self::Resource>,
        data: Self::InData,
    ) -> impl Future<Output = MutableStepResult<Self::Resource, Self::OutData, Self::Error>> + Send + 'a {
        let context = owned_context(context);
        (self.f)(context, data)
    }

    fn traced_name(&self) -> Cow<'static, str> { self.name.clone() }
}

/// Outcome for immutable steps.
pub enum StepOutcome<D> {
    /// Stop this run without reporting a modification.
    ///
    /// Use this when the controller is waiting for external state or does not
    /// yet have enough information to make a useful decision.
    Stop,
    /// Continue with the next step.
    ///
    /// No modification was made. `data` is transferred to the next step.
    NoModification { data: D },
}

/// Outcome for steps that may modify external state.
pub enum MutableStepOutcome<S, D> {
    /// Stop this run without reporting a modification.
    ///
    /// Use this when the controller is waiting for external state or does not
    /// yet have enough information to make a useful decision.
    Stop,
    /// Report that this step made one meaningful modification.
    ///
    /// The current run stops immediately. If `status` is present, staircase
    /// will try to write it as a best-effort status update. Reconcile
    /// correctness must not depend on that status update succeeding in the same
    /// run.
    Modified {
        status: Option<S>,
        /* , events: Option<Event> */
    },
    /// Continue with the next step.
    ///
    /// No modification was made. `data` is transferred to the next step.
    NoModification { data: D },
}