taktwerk 0.6.0

//! `taktwerk check` and `taktwerk run`.

use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::time::Duration;

use anyhow::{Context as _, bail};
use taktwerk_core::clock::Pace;
use taktwerk_core::connector::ConnectorError;
use taktwerk_core::image::Direction;
use taktwerk_core::image::{ImageError, ImageHandle, image_with_records};
use taktwerk_core::overrides::{self, Initial, Override};
use taktwerk_core::plan::{Plan, param_buffer};
use taktwerk_core::project::{Project, Sink};
use taktwerk_core::record::{Invocation, ScenarioInfo};
use taktwerk_core::scenario::{self, EventSpec, ScenarioFile};
use taktwerk_core::schedule::{Engine, EngineHandle, EngineOptions, instantiate};
use taktwerk_core::stop::StopCondition;
use taktwerk_record::{RecordService, ServiceReport};
use tokio::runtime::Runtime;
use tokio::signal::unix::{SignalKind, signal};
use tokio::sync::watch;
use tokio::task::JoinSet;
use tracing::{error, info, warn};

use crate::setup::{self, Ready, Stage};
use crate::{recording, summary};

/// How long a stopped engine keeps serving remote pulls until they have every row and event.
const PULL_DRAIN: Duration = Duration::from_secs(5);
/// How often that is checked.
const PULL_CHECK: Duration = Duration::from_millis(20);

/// How long connectors get to stop after shutdown is signalled.
const CONNECTOR_STOP: Duration = Duration::from_secs(5);

fn runtime() -> anyhow::Result<Runtime> {
    tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .thread_name("taktwerk-io")
        .enable_all()
        .build()
        .context("cannot start the I/O runtime")
}

/// Load, resolve and verify `project_file`, print the plan; no cycle is started and no listen
/// socket opened, so it runs beside an engine serving the same endpoint.
pub fn check(project_file: &Path) -> anyhow::Result<ExitCode> {
    let rt = runtime()?;
    match rt.block_on(setup::prepare(project_file, Stage::Check)) {
        Ok(ready) => {
            print!("{}", summary::plan(&ready.plan));
            println!(
                "\nok: {} instance(s), {} signal(s), {} connector(s) verified (servers not \
                 listening)",
                ready.plan.instances.len(),
                ready.plan.layout.len(),
                ready.connectors.len()
            );
            Ok(ExitCode::SUCCESS)
        }
        Err(problems) => {
            eprintln!("{}: {problems}", project_file.display());
            Ok(ExitCode::FAILURE)
        }
    }
}

/// Write every tunable's start value into the image, so connectors show it from the start.
/// The first write of a tunable logs no change event, so these precede the run in the log.
fn write_start_values(plan: &Plan, image: &ImageHandle) -> Result<(), ImageError> {
    for inst in &plan.instances {
        for port in &inst.tunables {
            if let Some(value) = inst.spec.params.get(&port.name) {
                image.write(port.signal, value)?;
            }
        }
    }
    Ok(())
}

/// Write the initial value of every external input an override set.
fn write_initial_values(
    plan: &Plan,
    image: &ImageHandle,
    initial: &[Initial],
) -> anyhow::Result<()> {
    for i in initial {
        let fail = |e: String| anyhow::anyhow!("--set {}: {e}", i.signal);
        let id = plan
            .layout
            .id(&i.signal)
            .ok_or_else(|| fail("no such signal, instance or engine setting".into()))?;
        let spec = plan
            .layout
            .spec(id)
            .ok_or_else(|| fail("no such signal".into()))?;
        if spec.direction != Direction::Input {
            return Err(fail(format!(
                "`{}` is not an external input; an override sets external inputs, parameters, \
                 tunables and dimensions",
                i.signal
            )));
        }
        let value = param_buffer(&i.value, spec.ty, &spec.shape, spec.layout).map_err(fail)?;
        image.write(id, &value).map_err(|e| fail(e.to_string()))?;
    }
    Ok(())
}

/// Why the run ends.
#[derive(Debug)]
enum End {
    Signal(&'static str),
    EngineStopped,
    Connector(String, Result<(), ConnectorError>),
}

type Tasks = JoinSet<(String, Result<(), ConnectorError>)>;

/// How `taktwerk run` runs a project.
#[derive(Debug, Clone)]
pub struct RunOptions {
    /// Stop after this many ticks.
    pub cycles: Option<u64>,
    /// How ticks follow the wall clock.
    pub pace: Pace,
    /// Start the project's connectors.
    pub connectors: bool,
    /// Stop before the first tick at or after this engine time, seconds.
    pub until_time: Option<f64>,
    /// Stop conditions as written, parsed against the plan.
    pub stop_when: Vec<String>,
    /// Overrides, applied in order once the models are loaded.
    pub set: Vec<Override>,
    /// Model packages swapped in: model id → path.
    pub models: Vec<(String, PathBuf)>,
    /// The directory every file sink writes to instead of its `dir`.
    pub record_dir: Option<PathBuf>,
    /// The run seed.
    pub seed: Option<u64>,
    /// A scenario file: overrides, stops and events; the other options win over it.
    pub scenario: Option<PathBuf>,
}

/// A scenario file as read for a run.
struct Loaded {
    events: Vec<EventSpec>,
    info: ScenarioInfo,
}

/// `options` with the scenario file's overrides, models, stops and seed merged in under the
/// command line's, and the file's events.
fn with_scenario(options: &RunOptions) -> anyhow::Result<(RunOptions, Option<Loaded>)> {
    let mut merged = options.clone();
    let Some(path) = &options.scenario else {
        return Ok((merged, None));
    };
    let text = std::fs::read_to_string(path)
        .with_context(|| format!("read the scenario {}", path.display()))?;
    let file: ScenarioFile =
        toml::from_str(&text).with_context(|| format!("scenario {}", path.display()))?;
    let dir = path.parent().unwrap_or_else(|| Path::new(""));
    merged.seed = options.seed.or(file.seed);
    merged.cycles = options.cycles.or(file.stop.cycles);
    merged.until_time = options.until_time.or(file.stop.time_s);
    merged.stop_when = file
        .stop
        .when
        .iter()
        .chain(&options.stop_when)
        .cloned()
        .collect();
    merged.set = file
        .set
        .into_iter()
        .map(|(key, value)| Override { key, value })
        .chain(options.set.iter().cloned())
        .collect();
    merged.models = file
        .models
        .into_iter()
        .map(|(id, p)| (id, dir.join(p)))
        .chain(options.models.iter().cloned())
        .collect();
    let info = ScenarioInfo {
        path: path.display().to_string(),
        hash: taktwerk_record::util::sha256(text.as_bytes()),
        text,
    };
    Ok((
        merged,
        Some(Loaded {
            events: file.events,
            info,
        }),
    ))
}

impl RunOptions {
    /// `--speed` as given.
    fn speed(&self) -> String {
        match self.pace {
            Pace::Realtime => "1".to_owned(),
            Pace::Scaled(f) => f.to_string(),
            Pace::Unpaced => "max".to_owned(),
        }
    }
}

/// A path given on the command line, made absolute against the working directory.
fn absolute(path: &Path) -> anyhow::Result<PathBuf> {
    Ok(std::path::absolute(path)?)
}

/// Run `project_file` until SIGINT or SIGTERM, a model error, a connector failure or a stop
/// condition of `options`.
pub fn run(project_file: &Path, options: &RunOptions) -> anyhow::Result<ExitCode> {
    let (options, scenario) = with_scenario(options)?;
    let options = &options;
    let rt = runtime()?;
    let mut project = Project::load(project_file)?;
    if !options.connectors {
        project.connectors.clear();
    }
    if let Some(dir) = &options.record_dir {
        let dir = absolute(dir)?;
        for r in project
            .recorders
            .iter_mut()
            .filter(|r| r.sink == Sink::File)
        {
            r.dir.clone_from(&dir);
        }
    }
    let mut swapped = Vec::with_capacity(options.models.len());
    for (id, path) in &options.models {
        swapped.push(
            overrides::swap_model(&mut project, id, absolute(path)?).map_err(anyhow::Error::msg)?,
        );
    }
    if !options.pace.is_realtime()
        && let Some(c) = project.connectors.iter().find(|c| c.kind == "opcua-client")
    {
        bail!(
            "connector `{}`: an opcua-client exchanges values in wall time; run it in real time \
             (--speed 1) or without connectors (--no-connectors)",
            c.id
        );
    }
    let prepared = setup::prepare_loaded(
        project_file,
        project,
        &options.set,
        options.seed,
        Stage::Run,
    );
    let ready = match rt.block_on(prepared) {
        Ok(ready) => ready,
        Err(problems) => {
            eprintln!("{}: {problems}", project_file.display());
            return Ok(ExitCode::FAILURE);
        }
    };
    let Ready {
        project,
        adapters,
        connectors,
        plan,
        applied,
        initial,
        seeds,
    } = ready;
    info!(
        instances = plan.instances.len(),
        signals = plan.layout.len(),
        tick_ms = summary::ms(plan.tick),
        "plan resolved"
    );

    let (mut sigint, mut sigterm) = rt.block_on(async {
        Ok::<_, std::io::Error>((
            signal(SignalKind::interrupt())?,
            signal(SignalKind::terminate())?,
        ))
    })?;

    let mut stop_when = Vec::with_capacity(options.stop_when.len());
    for text in &options.stop_when {
        stop_when.push(StopCondition::parse(text, &plan.layout).map_err(anyhow::Error::msg)?);
    }
    if let Some(t) = options.until_time
        && !t.is_finite()
    {
        bail!("--until-time must be a finite number of seconds, not {t}");
    }
    let mut events = match &scenario {
        Some(l) => scenario::resolve(
            &l.events,
            &plan.layout,
            plan.start_time,
            plan.tick.as_secs_f64(),
        )
        .map_err(|e| anyhow::anyhow!("scenario {}: {e}", l.info.path))?,
        None => Vec::new(),
    };
    // The order they apply in, which the event log's indices refer to.
    events.sort_by_key(|e| e.cycle);
    let invocation = Invocation {
        speed: options.speed(),
        connectors: options.connectors,
        cycles: options.cycles,
        until_time: options.until_time,
        stop_when: options.stop_when.clone(),
        overrides: swapped.into_iter().chain(applied).collect(),
        seed: options.seed,
        seeds,
        scenario: scenario.map(|l| l.info),
    };
    let recorders = recording::records(project_file, &project, &adapters, &plan, &invocation)?;
    let (records, service_config) = recorders.unzip();
    let service_config = service_config.map(|mut c| {
        c.conditions.clone_from(&options.stop_when);
        c.scenario = events
            .iter()
            .map(|e| (e.source.clone(), e.detail.clone()))
            .collect();
        c
    });
    let (cycle, handle) = image_with_records(plan.layout.clone(), records);
    write_start_values(&plan, &handle).context("tunable start values")?;
    write_initial_values(&plan, &handle, &initial)?;
    let service = service_config
        .map(|config| RecordService::start(&handle, config))
        .transpose()
        .context("cannot start the recorders")?;
    let (shutdown, shutdown_rx) = watch::channel(false);
    let mut tasks = Tasks::new();
    for c in connectors {
        let id = c.id().to_owned();
        let fut = c.run(handle.clone(), shutdown_rx.clone());
        tasks.spawn_on(async move { (id, fut.await) }, rt.handle());
    }

    let engine_options = EngineOptions {
        realtime: project.engine.realtime.clone(),
        cycles: options.cycles,
        pace: options.pace,
        until_time: options.until_time,
        stop_when,
        events,
        seed: options.seed.unwrap_or(0),
    };
    let started =
        instantiate(&plan, &adapters).and_then(|i| Engine::start(plan, i, cycle, engine_options));
    let engine = match started {
        Ok(engine) => engine,
        Err(e) => {
            shutdown.send_replace(true);
            let _ = rt.block_on(stop_connectors(&mut tasks));
            if let Some(service) = service {
                report(&service.stop());
            }
            return Err(anyhow::Error::new(e).context("the engine did not start"));
        }
    };
    info!("running; SIGINT or SIGTERM stops");

    let end = rt.block_on(wait(&engine, &mut tasks, &mut sigint, &mut sigterm));
    match &end {
        End::Signal(name) => info!("{name} received, stopping"),
        End::EngineStopped => {}
        End::Connector(id, result) => match result {
            Ok(()) => warn!(connector = %id, "connector ended, stopping"),
            Err(e) => error!(connector = %id, "connector failed, stopping: {e}"),
        },
    }
    engine.stop();
    let result = engine.join();
    let mut failed = service.map(|s| report(&s.stop())).unwrap_or_default();
    if !matches!(end, End::Connector(..)) {
        rt.block_on(wait_for_pullers(&handle));
    }
    shutdown.send_replace(true);
    failed.extend(rt.block_on(stop_connectors(&mut tasks)));
    if let End::Connector(id, Err(e)) = end {
        failed.push(format!("connector `{id}`: {e}"));
    }

    let mut code = ExitCode::SUCCESS;
    match result {
        Ok(s) => {
            let reason = match s.condition.and_then(|k| options.stop_when.get(k)) {
                Some(text) => format!("condition {text}"),
                None => s.stop.as_str().to_owned(),
            };
            println!("{}  stop {reason}", summary::run(&s));
        }
        Err(e) => {
            eprintln!("error: {e}");
            code = ExitCode::FAILURE;
        }
    }
    for f in failed {
        eprintln!("error: {f}");
        code = ExitCode::FAILURE;
    }
    Ok(code)
}

/// Log what the record service did; its failures, as text.
fn report(report: &ServiceReport) -> Vec<String> {
    for sink in &report.sinks {
        info!(
            recorder = %sink.recorder,
            rows = sink.counts.rows_written,
            dropped = sink.counts.rows_dropped,
            parts = sink.stats.part,
            dir = %sink.dir.display(),
            "run recorded"
        );
    }
    if report.events_lost > 0 {
        warn!(
            lost = report.events_lost,
            "cycle events were lost before they were logged"
        );
    }
    report.errors.clone()
}

async fn wait(
    engine: &EngineHandle,
    tasks: &mut Tasks,
    sigint: &mut tokio::signal::unix::Signal,
    sigterm: &mut tokio::signal::unix::Signal,
) -> End {
    let mut finished = engine.finished();
    let connector = async {
        match tasks.join_next().await {
            Some(Ok((id, result))) => End::Connector(id, result),
            Some(Err(e)) => End::Connector("?".into(), Err(ConnectorError::Io(e.to_string()))),
            None => std::future::pending().await,
        }
    };
    tokio::select! {
        _ = sigint.recv() => End::Signal("SIGINT"),
        _ = sigterm.recv() => End::Signal("SIGTERM"),
        _ = finished.wait_for(|f| *f) => End::EngineStopped,
        end = connector => end,
    }
}

/// After the engine stopped and its last events are logged, keep serving until every recorder
/// pulled during the run has been pulled to its ring's head and the log's end, at most
/// [`PULL_DRAIN`]. No wait when nobody pulled.
async fn wait_for_pullers(image: &ImageHandle) {
    let Some(records) = image.records() else {
        return;
    };
    if !records.pulls_behind() {
        return;
    }
    info!("waiting for remote recorders to pull the end of the run");
    let drained = tokio::time::timeout(PULL_DRAIN, async {
        while records.pulls_behind() {
            tokio::time::sleep(PULL_CHECK).await;
        }
    })
    .await;
    if drained.is_err() {
        warn!("remote recorders did not pull the end of the run within {PULL_DRAIN:?}");
    }
}

/// Wait for every connector after shutdown; the failures, as text.
async fn stop_connectors(tasks: &mut Tasks) -> Vec<String> {
    let mut failed = Vec::new();
    let all = async {
        while let Some(joined) = tasks.join_next().await {
            match joined {
                Ok((_, Ok(()))) => {}
                Ok((id, Err(e))) => failed.push(format!("connector `{id}`: {e}")),
                Err(e) => failed.push(format!("connector task: {e}")),
            }
        }
    };
    if tokio::time::timeout(CONNECTOR_STOP, all).await.is_err() {
        tasks.abort_all();
        warn!("connectors did not stop within {CONNECTOR_STOP:?}; aborted");
    }
    failed
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    reason = "tests"
)]
mod tests {
    use std::collections::BTreeMap;

    use taktwerk_core::model::{Causality, Instances, ModelInterface, Variable};
    use taktwerk_core::project::Project;
    use taktwerk_core::value::{Buffer, Layout, ScalarType};

    use super::*;

    #[tokio::test]
    async fn tunable_start_values_reach_the_image() {
        let iface = ModelInterface {
            name: "g".into(),
            dimensions: vec![],
            variables: vec![Variable {
                name: "k".into(),
                causality: Causality::Tunable,
                ty: ScalarType::F64,
                shape: vec![],
                layout: Layout::RowMajor,
                unit: None,
                description: None,
                text: false,
            }],
            instances: Instances::Multiple,
        };
        let project: Project =
            "[engine]\ntick_ms = 1.0\n[models.g]\nkind = \"fmi\"\npath = \"g\"\n\
             [[instance]]\nid = \"a\"\nmodel = \"g\"\nparameters = { k = 2.5 }\n"
                .parse()
                .unwrap();
        let interfaces = BTreeMap::from([("g".to_owned(), iface)]);
        let plan = Plan::resolve(&project, &interfaces, &mut []).await.unwrap();
        let (_cycle, handle) = image_with_records(plan.layout.clone(), None);
        write_start_values(&plan, &handle).unwrap();
        let id = plan.layout.id("a.k").unwrap();
        let mut out = Buffer::zeroed(ScalarType::F64, 1);
        assert!(handle.read(id, &mut out).unwrap().is_some());
        assert_eq!(out, Buffer::F64(vec![2.5]));
    }
}