ontoref-code/crates/ontoref-reflection/src/executor.rs

use std::{collections::HashMap, sync::Arc, time::Duration};

use anyhow::{anyhow, Context, Result};
use regex::Regex;
use stratum_graph::types::NodeId;
use stratum_state::{PipelineRun, PipelineRunId, PipelineStatus, StateTracker, StepRecord};
use tokio::task::JoinSet;
use tracing::{info, warn};

use crate::{
    dag,
    error::ReflectionError,
    mode::{ActionStep, Actor, ErrorStrategy, ReflectionMode},
};

/// Context provided to a mode execution run.
pub struct RunContext {
    /// Project identifier — used in NATS trigger subject.
    pub project: String,
    /// Parameter values substituted into step `cmd` fields via `{key}`
    /// placeholders.
    pub params: HashMap<String, String>,
    /// State tracker for recording pipeline run and step records.
    pub state: Arc<dyn StateTracker>,
    /// Optional NATS stream for publishing step completion events.
    #[cfg(feature = "nats")]
    pub nats: Option<Arc<platform_nats::EventStream>>,
}

/// Result of executing a reflection mode.
#[derive(Debug)]
pub struct ModeRun {
    pub mode_id: String,
    pub run_id: PipelineRunId,
    pub final_status: PipelineStatus,
}

impl ReflectionMode {
    /// Validate and execute this mode against the provided context.
    /// Steps within each topological layer run concurrently via `JoinSet`.
    /// Step failure behaviour is governed by `on_error.strategy`.
    pub async fn execute(&self, ctx: &RunContext) -> Result<ModeRun> {
        self.validate()
            .with_context(|| format!("mode '{}' failed pre-execution DAG validation", self.id))?;

        let trigger_subject = format!("ecosystem.reflection.{}.{}", self.id, ctx.project);
        let trigger_payload = serde_json::to_value(&ctx.params)
            .context("serializing RunContext.params as trigger payload")?;

        let run = PipelineRun::new(trigger_subject, trigger_payload);
        let run_id = run.id.clone();
        ctx.state
            .create_run(&run)
            .await
            .context("creating PipelineRun in state tracker")?;

        let layers = dag::topological_layers(&self.steps)
            .with_context(|| format!("computing execution layers for mode '{}'", self.id))?;

        let step_index: HashMap<&str, &ActionStep> =
            self.steps.iter().map(|s| (s.id.as_str(), s)).collect();

        for layer in &layers {
            let failure = run_layer(layer, &step_index, ctx, &run_id, &self.id).await?;

            if let Some(e) = failure {
                ctx.state
                    .update_status(&run_id, PipelineStatus::Failed)
                    .await
                    .context("updating pipeline status to Failed")?;
                return Err(e);
            }
        }

        ctx.state
            .update_status(&run_id, PipelineStatus::Success)
            .await
            .context("updating pipeline status to Success")?;

        info!(mode = %self.id, run = %run_id, "mode completed successfully");

        Ok(ModeRun {
            mode_id: self.id.clone(),
            run_id,
            final_status: PipelineStatus::Success,
        })
    }
}

/// Execute all steps in a layer concurrently. Returns the first fatal error, if
/// any.
async fn run_layer(
    layer: &[String],
    step_index: &HashMap<&str, &ActionStep>,
    ctx: &RunContext,
    run_id: &PipelineRunId,
    mode_id: &str,
) -> Result<Option<anyhow::Error>> {
    let mut set: JoinSet<(String, Result<()>)> = JoinSet::new();

    for step_id in layer {
        let step = (*step_index
            .get(step_id.as_str())
            .expect("layer ids are derived from the validated step list"))
        .clone();
        let owned_step_id = step.id.clone();
        let params = ctx.params.clone();
        let state = Arc::clone(&ctx.state);
        let owned_run_id = run_id.clone();
        let owned_mode_id = mode_id.to_owned();
        let owned_project = ctx.project.clone();
        #[cfg(feature = "nats")]
        let nats = ctx.nats.clone();

        set.spawn(async move {
            #[cfg(feature = "nats")]
            let result = execute_step(
                step,
                params,
                state,
                owned_run_id,
                owned_mode_id,
                owned_project,
                nats,
            )
            .await;
            #[cfg(not(feature = "nats"))]
            let result = execute_step(
                step,
                params,
                state,
                owned_run_id,
                owned_mode_id,
                owned_project,
            )
            .await;
            (owned_step_id, result)
        });
    }

    let mut fatal: Option<anyhow::Error> = None;

    while let Some(join_result) = set.join_next().await {
        let (step_id, result) =
            join_result.map_err(|e| anyhow!(ReflectionError::TaskPanic(e.to_string())))?;

        if let Err(e) = result {
            let step = step_index[step_id.as_str()];
            match step.on_error.strategy {
                ErrorStrategy::Stop | ErrorStrategy::Retry => {
                    // Retry is exhausted inside execute_step; treat the final failure as Stop.
                    set.abort_all();
                    fatal = Some(e);
                    break;
                }
                ErrorStrategy::Continue | ErrorStrategy::Fallback | ErrorStrategy::Branch => {
                    warn!(
                        step = %step_id,
                        strategy = ?step.on_error.strategy,
                        "step failed, continuing per on_error strategy: {e}"
                    );
                }
            }
        }
    }

    Ok(fatal)
}

/// Execute a single step: record start → run cmd (with retry) → record outcome.
/// All parameters are owned so this future is `'static` and can be spawned.
#[cfg_attr(not(feature = "nats"), allow(unused_variables))]
async fn execute_step(
    step: ActionStep,
    params: HashMap<String, String>,
    state: Arc<dyn StateTracker>,
    run_id: PipelineRunId,
    mode_id: String,
    project: String,
    #[cfg(feature = "nats")] nats: Option<Arc<platform_nats::EventStream>>,
) -> Result<()> {
    let start_record = StepRecord::start(NodeId(step.id.clone()));
    state
        .record_step(&run_id, &start_record)
        .await
        .with_context(|| format!("recording step '{}' start", step.id))?;

    info!(step = %step.id, action = %step.action, actor = ?step.actor, "executing step");

    let outcome = match step.actor {
        Actor::Human => {
            info!(
                step = %step.id,
                "step requires human action — skipping automated execution: {}",
                step.note.as_deref().unwrap_or(&step.action)
            );
            Ok(())
        }
        Actor::Agent | Actor::Both => match &step.cmd {
            None => {
                info!(step = %step.id, "no cmd — step is documentation-only");
                Ok(())
            }
            Some(cmd) => {
                let resolved = substitute_params(cmd, &params)
                    .with_context(|| format!("substituting params in step '{}' cmd", step.id))?;
                run_with_retry(&resolved, &step.on_error).await
            }
        },
    };

    match outcome {
        Ok(()) => {
            state
                .record_step(&run_id, &start_record.succeed(vec![]))
                .await
                .with_context(|| format!("recording step '{}' success", step.id))?;

            #[cfg(feature = "nats")]
            if let Some(ref nats) = nats {
                publish_step_event(nats, &run_id, &step.id, true).await;
                publish_kogral_capture(nats, &run_id, &mode_id, &project, &step.id, "success")
                    .await;
            }

            Ok(())
        }
        Err(e) => {
            let err_str = e.to_string();
            state
                .record_step(&run_id, &start_record.fail(err_str.clone()))
                .await
                .with_context(|| format!("recording step '{}' failure", step.id))?;

            #[cfg(feature = "nats")]
            if let Some(ref nats) = nats {
                publish_step_event(nats, &run_id, &step.id, false).await;
                publish_kogral_capture(nats, &run_id, &mode_id, &project, &step.id, "failed").await;
            }

            Err(anyhow!("step '{}' failed: {}", step.id, err_str))
        }
    }
}

/// Replace `{key}` placeholders in `cmd` with values from `params`.
/// Returns `Err(UnknownParams)` if any placeholder key is absent from `params`.
fn substitute_params(
    cmd: &str,
    params: &HashMap<String, String>,
) -> Result<String, ReflectionError> {
    let re = Regex::new(r"\{([^}]+)\}").expect("static regex is valid");

    let unknown: Vec<String> = re
        .captures_iter(cmd)
        .filter_map(|cap| {
            let key = cap[1].to_string();
            if params.contains_key(&key) {
                None
            } else {
                Some(key)
            }
        })
        .collect();

    if !unknown.is_empty() {
        return Err(ReflectionError::UnknownParams(unknown));
    }

    let mut result = cmd.to_owned();
    for (key, value) in params {
        result = result.replace(&format!("{{{key}}}"), value);
    }

    Ok(result)
}

/// Execute `cmd` via `sh -c`, retrying on failure when `on_error.strategy ==
/// Retry`.
async fn run_with_retry(cmd: &str, on_error: &crate::mode::OnError) -> Result<()> {
    let max_attempts = if on_error.strategy == ErrorStrategy::Retry {
        on_error.max.max(1)
    } else {
        1
    };

    for attempt in 0..max_attempts {
        match run_cmd(cmd).await {
            Ok(()) => return Ok(()),
            Err(e) => {
                if attempt + 1 < max_attempts {
                    let delay = Duration::from_secs(on_error.backoff_s);
                    warn!(
                        attempt = attempt + 1,
                        max = max_attempts,
                        "cmd failed, retrying in {delay:?}: {e}"
                    );
                    tokio::time::sleep(delay).await;
                } else {
                    return Err(e);
                }
            }
        }
    }

    // Defensive: reached only when max_attempts == 0, which is prevented by .max(1)
    // above.
    Err(anyhow!("retry loop exited without result"))
}

async fn run_cmd(cmd: &str) -> Result<()> {
    let output = tokio::process::Command::new("sh")
        .arg("-c")
        .arg(cmd)
        .output()
        .await
        .with_context(|| format!("spawning sh -c: {cmd}"))?;

    if output.status.success() {
        Ok(())
    } else {
        let stderr = String::from_utf8_lossy(&output.stderr);
        let stdout = String::from_utf8_lossy(&output.stdout);
        Err(anyhow!(
            "command exited {}: stderr='{}' stdout='{}'",
            output.status,
            stderr.trim(),
            stdout.trim()
        ))
    }
}

#[cfg(feature = "nats")]
async fn publish_step_event(
    nats: &platform_nats::EventStream,
    run_id: &PipelineRunId,
    step_id: &str,
    success: bool,
) {
    let subject = if success {
        "ecosystem.reflection.step.completed"
    } else {
        "ecosystem.reflection.step.failed"
    };

    let payload = serde_json::json!({
        "run_id":  run_id.0.to_string(),
        "step_id": step_id,
        "success": success,
    });

    match serde_json::to_vec(&payload) {
        Ok(bytes) => {
            if let Err(e) = nats.publish(subject, bytes::Bytes::from(bytes)).await {
                warn!(step = %step_id, "failed to publish step event to '{subject}': {e}");
            }
        }
        Err(e) => {
            warn!(step = %step_id, "failed to serialize step event: {e}");
        }
    }
}

/// Publish to `ecosystem.kogral.capture` so Kogral can record this step as an
/// Execution node in the shared knowledge graph. Matches the KogralCapture
/// payload contract defined in `nats/subjects.ncl`.
#[cfg(feature = "nats")]
async fn publish_kogral_capture(
    nats: &platform_nats::EventStream,
    run_id: &PipelineRunId,
    mode_id: &str,
    project: &str,
    step_id: &str,
    status: &str,
) {
    const SUBJECT: &str = "ecosystem.kogral.capture";

    let payload = serde_json::json!({
        "project":  project,
        "mode_id":  mode_id,
        "step_id":  step_id,
        "run_id":   run_id.0.to_string(),
        "action":   step_id,
        "status":   status,
        "context":  {},
    });

    match serde_json::to_vec(&payload) {
        Ok(bytes) => {
            if let Err(e) = nats.publish(SUBJECT, bytes::Bytes::from(bytes)).await {
                warn!(step = %step_id, "failed to publish kogral capture to '{SUBJECT}': {e}");
            }
        }
        Err(e) => {
            warn!(step = %step_id, "failed to serialize kogral capture payload: {e}");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn substitute_replaces_known_params() {
        let mut params = HashMap::new();
        params.insert("project_name".to_string(), "my-service".to_string());
        params.insert("project_dir".to_string(), "/tmp/my-service".to_string());

        let result =
            substitute_params("git -C {project_dir} init && echo {project_name}", &params).unwrap();
        assert_eq!(result, "git -C /tmp/my-service init && echo my-service");
    }

    #[test]
    fn substitute_rejects_unknown_params() {
        let params = HashMap::new();
        let err = substitute_params("echo {unknown_key}", &params).unwrap_err();
        assert!(matches!(
            err,
            ReflectionError::UnknownParams(keys) if keys.contains(&"unknown_key".to_string())
        ));
    }

    #[test]
    fn substitute_no_placeholders_is_identity() {
        let result = substitute_params("ls -la /tmp", &HashMap::new()).unwrap();
        assert_eq!(result, "ls -la /tmp");
    }

    #[tokio::test]
    async fn run_cmd_success() {
        run_cmd("true").await.unwrap();
    }

    #[tokio::test]
    async fn run_cmd_failure_returns_err() {
        let err = run_cmd("false").await.unwrap_err();
        assert!(err.to_string().contains("command exited"));
    }
}