datum-agent 0.10.0

use std::{
    collections::BTreeMap,
    fmt,
    sync::Arc,
    time::{Duration, Instant, SystemTime},
};

use datum::{
    Flow, KillSwitches, NotUsed, RestartSettings, RunnableGraph, SharedKillSwitch,
    StreamCompletion, StreamError, StreamInstrumentationRegistry, StreamResult,
};

/// Stable identifier assigned by the local registry when a job is submitted.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct JobId(pub u64);

/// Factory used by [`JobSpec`] to build a fresh job blueprint for one generation.
///
/// The factory receives a [`JobContext`] containing the generation identity and an agent-owned
/// `SharedKillSwitch`. The factory should build a `RunnableGraph<JobMat>` and should not
/// materialize it.
pub type JobGraphFactory =
    dyn Fn(JobContext) -> StreamResult<RunnableGraph<JobMat>> + Send + Sync + 'static;

/// A named, supervised stream program.
#[derive(Clone)]
pub struct JobSpec {
    pub name: String,
    pub factory: Arc<JobGraphFactory>,
    pub restart_policy: JobRestartPolicy,
    pub drain_behavior: JobDrainBehavior,
    pub cluster: Option<ClusterJobMetadata>,
}

impl JobSpec {
    #[must_use]
    pub fn new<F>(name: impl Into<String>, factory: F) -> Self
    where
        F: Fn(JobContext) -> StreamResult<RunnableGraph<JobMat>> + Send + Sync + 'static,
    {
        Self {
            name: name.into(),
            factory: Arc::new(factory),
            restart_policy: JobRestartPolicy::Never,
            drain_behavior: JobDrainBehavior::default(),
            cluster: None,
        }
    }

    #[must_use]
    pub fn with_restart_policy(mut self, restart_policy: JobRestartPolicy) -> Self {
        self.restart_policy = restart_policy;
        self
    }

    #[must_use]
    pub fn with_drain_behavior(mut self, drain_behavior: JobDrainBehavior) -> Self {
        self.drain_behavior = drain_behavior;
        self
    }

    #[must_use]
    pub fn with_cluster_metadata(mut self, metadata: ClusterJobMetadata) -> Self {
        self.cluster = Some(metadata);
        self
    }
}

impl fmt::Debug for JobSpec {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("JobSpec")
            .field("name", &self.name)
            .field("restart_policy", &self.restart_policy)
            .field("drain_behavior", &self.drain_behavior)
            .field("cluster", &self.cluster)
            .finish_non_exhaustive()
    }
}

/// Placement strategy selected for a cluster-submitted job.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PlacementStrategy {
    /// Pick the eligible node with the fewest cluster jobs, then tie-break by node id.
    LeastJobs,
    /// Place the job on a specific node id.
    Pinned { node_id: String },
}

/// Placement constraints attached to a cluster-submitted job.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PlacementSpec {
    pub role_constraint: Option<String>,
    pub strategy: PlacementStrategy,
}

impl PlacementSpec {
    #[must_use]
    pub fn least_jobs(role_constraint: Option<String>) -> Self {
        Self {
            role_constraint,
            strategy: PlacementStrategy::LeastJobs,
        }
    }

    #[must_use]
    pub fn pinned(node_id: impl Into<String>, role_constraint: Option<String>) -> Self {
        Self {
            role_constraint,
            strategy: PlacementStrategy::Pinned {
                node_id: node_id.into(),
            },
        }
    }
}

/// One placement-history entry for a cluster-submitted job.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ClusterPlacementHistory {
    pub generation: u64,
    pub from_node: Option<String>,
    pub to_node: String,
    pub reason: String,
    pub timestamp: SystemTime,
}

/// Cluster ownership metadata persisted in the local registry entry for jobs
/// started by the placement coordinator.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ClusterJobMetadata {
    pub factory_name: String,
    pub params: BTreeMap<String, String>,
    pub placement: PlacementSpec,
    pub coordinator_node: String,
    pub assigned_node: String,
    pub placement_generation: u64,
    pub history: Vec<ClusterPlacementHistory>,
}

/// Context passed to a job blueprint factory for one materialized generation.
#[derive(Clone, Debug)]
pub struct JobContext {
    name: Arc<str>,
    job_id: JobId,
    generation: u64,
    kill_switch: SharedKillSwitch,
    instrumentation: StreamInstrumentationRegistry,
}

impl JobContext {
    pub(crate) fn new(
        name: impl Into<Arc<str>>,
        job_id: JobId,
        generation: u64,
        kill_switch: SharedKillSwitch,
        instrumentation: StreamInstrumentationRegistry,
    ) -> Self {
        Self {
            name: name.into(),
            job_id,
            generation,
            kill_switch,
            instrumentation,
        }
    }

    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    #[must_use]
    pub fn job_id(&self) -> JobId {
        self.job_id
    }

    #[must_use]
    pub fn generation(&self) -> u64 {
        self.generation
    }

    #[must_use]
    pub fn kill_switch(&self) -> SharedKillSwitch {
        self.kill_switch.clone()
    }

    #[must_use]
    pub fn instrumentation_registry(&self) -> &StreamInstrumentationRegistry {
        &self.instrumentation
    }

    /// Return a flow backed by the agent-owned shared kill switch.
    ///
    /// Graph factories should wire this flow into the job path when graceful drain is supported.
    #[must_use]
    pub fn drain_flow<T: Send + 'static>(&self) -> Flow<T, T, SharedKillSwitch> {
        self.kill_switch.flow()
    }

    #[must_use]
    pub fn control(&self) -> JobControl {
        JobControl::graceful(self.kill_switch.clone())
    }
}

/// Standardized materialized value for daemon-managed jobs.
pub struct JobMat {
    completion: StreamCompletion<NotUsed>,
    control: JobControl,
}

impl JobMat {
    #[must_use]
    pub fn new(completion: StreamCompletion<NotUsed>, control: JobControl) -> Self {
        Self {
            completion,
            control,
        }
    }

    #[must_use]
    pub fn graceful(completion: StreamCompletion<NotUsed>, kill_switch: SharedKillSwitch) -> Self {
        Self::new(completion, JobControl::graceful(kill_switch))
    }

    #[must_use]
    pub fn cancel_only(completion: StreamCompletion<NotUsed>) -> Self {
        Self::new(completion, JobControl::cancel_only())
    }

    pub(crate) fn into_parts(self) -> (StreamCompletion<NotUsed>, JobControl) {
        (self.completion, self.control)
    }
}

impl fmt::Debug for JobMat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("JobMat")
            .field("control", &self.control)
            .finish_non_exhaustive()
    }
}

/// Control handles associated with a running job generation.
#[derive(Clone, Debug)]
pub struct JobControl {
    kill_switch: Option<SharedKillSwitch>,
}

impl JobControl {
    #[must_use]
    pub fn graceful(kill_switch: SharedKillSwitch) -> Self {
        Self {
            kill_switch: Some(kill_switch),
        }
    }

    #[must_use]
    pub fn cancel_only() -> Self {
        Self { kill_switch: None }
    }

    #[must_use]
    pub fn drain_supported(&self) -> bool {
        self.kill_switch.is_some()
    }

    #[must_use]
    pub fn kill_switch(&self) -> Option<SharedKillSwitch> {
        self.kill_switch.clone()
    }

    pub(crate) fn shutdown(&self) -> bool {
        if let Some(kill_switch) = &self.kill_switch {
            kill_switch.shutdown();
            true
        } else {
            false
        }
    }

    pub(crate) fn abort(&self, error: StreamError) -> bool {
        if let Some(kill_switch) = &self.kill_switch {
            kill_switch.abort(error);
            true
        } else {
            false
        }
    }
}

/// Registry-visible restart policy for a job.
#[derive(Clone, Debug)]
pub enum JobRestartPolicy {
    Never,
    OnFailure(RestartSettings),
    Always(RestartSettings),
    Manual,
}

impl JobRestartPolicy {
    #[must_use]
    pub fn never() -> Self {
        Self::Never
    }

    #[must_use]
    pub fn on_failure(settings: RestartSettings) -> Self {
        Self::OnFailure(settings)
    }

    #[must_use]
    pub fn always(settings: RestartSettings) -> Self {
        Self::Always(settings)
    }

    #[must_use]
    pub fn manual() -> Self {
        Self::Manual
    }

    pub(crate) fn settings_for(&self, cause: RestartCause) -> Option<RestartSettings> {
        match (self, cause) {
            (Self::Always(settings), RestartCause::Failure | RestartCause::Completion)
            | (Self::OnFailure(settings), RestartCause::Failure) => Some(settings.clone()),
            (Self::Never | Self::Manual | Self::OnFailure(_), _) => None,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum RestartCause {
    Failure,
    Completion,
}

/// Drain behavior declared by a job spec.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum JobDrainBehavior {
    Graceful { timeout: Duration },
    CancelOnly,
}

impl JobDrainBehavior {
    #[must_use]
    pub fn graceful(timeout: Duration) -> Self {
        Self::Graceful { timeout }
    }

    #[must_use]
    pub fn cancel_only() -> Self {
        Self::CancelOnly
    }

    #[must_use]
    pub fn timeout(&self) -> Option<Duration> {
        match self {
            Self::Graceful { timeout } => Some(*timeout),
            Self::CancelOnly => None,
        }
    }
}

impl Default for JobDrainBehavior {
    fn default() -> Self {
        Self::Graceful {
            timeout: Duration::from_secs(30),
        }
    }
}

/// Desired state recorded by the registry.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DesiredJobState {
    Running,
    Draining,
    Stopped,
}

/// Observed lifecycle state of a job.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobState {
    Submitted,
    Starting,
    Running,
    Draining,
    BackingOff,
    Completed,
    Drained,
    Stopped,
    Failed,
}

/// Why the latest generation exited.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JobExitReason {
    Completed,
    Failed(StreamError),
    Drained,
    Stopped,
    DrainTimedOut,
}

/// Reserved integration point for WP-A1 stream instrumentation.
///
/// WP-A1 is not present in the current main branch, so WP-A2 intentionally exposes no per-element
/// counters yet. Future instrumentation handles should fill this snapshot without changing the
/// registry control-plane API.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct JobInstrumentationSnapshot {}

/// Point-in-time registry snapshot for one job.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JobStatus {
    pub name: String,
    pub job_id: JobId,
    pub state: JobState,
    pub desired_state: DesiredJobState,
    pub generation: u64,
    pub starts_total: u64,
    pub restarts_total: u64,
    pub last_start_at: Option<SystemTime>,
    pub last_exit_at: Option<SystemTime>,
    pub last_exit_reason: Option<JobExitReason>,
    pub backoff_until: Option<Instant>,
    pub drain_deadline: Option<Instant>,
    pub drain_supported: bool,
    pub active_streams: Option<usize>,
    pub instrumentation: Option<JobInstrumentationSnapshot>,
    pub cluster: Option<ClusterJobMetadata>,
}

/// Lifecycle event emitted by the registry.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JobEvent {
    pub sequence: u64,
    pub timestamp: SystemTime,
    pub name: String,
    pub job_id: JobId,
    pub generation: u64,
    pub kind: JobEventKind,
}

/// Kind-specific lifecycle event details.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JobEventKind {
    Submitted,
    Started,
    Failed { reason: JobExitReason },
    RestartScheduled { delay: Duration },
    Restarted { previous_generation: u64 },
    Draining,
    Drained,
    Stopped { reason: JobExitReason },
    Completed,
}

pub(crate) fn new_generation_kill_switch(name: &str, generation: u64) -> SharedKillSwitch {
    KillSwitches::shared(format!("job:{name}:{generation}"))
}