tachyon-types 0.7.3

use std::collections::HashMap;

use myko::command::{CommandContext, CommandError, CommandHandler};
use myko_macros::myko_command;
use serde_json::Value;

use crate::{
    derive_resources, GetAllInventorys, GetPinsByIds, GetPlaybooksByIds, GetRunsByIds, HostStats,
    OutputLine, OutputLineId, OutputStream, Pin, PinId, Playbook, PlaybookId, Run, RunId, RunStatus,
};

/// Run an Ansible playbook. Creates a new `Run` entity with status `Queued`;
/// the admission step in tachyon-server promotes it to `Running` once it holds
/// a lease on every resource it targets, and the run watcher then executes it.
#[myko_command(RunId)]
pub struct RunPlaybook {
    pub playbook_id: PlaybookId,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub extra_vars: Option<Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub limit: Option<String>,
}

impl CommandHandler for RunPlaybook {
    fn execute(self, ctx: CommandContext) -> Result<RunId, CommandError> {
        let run = prepare_playbook_run(&ctx, self)?;
        let run_id = run.id.clone();
        ctx.emit_set(&run)?;
        Ok(run_id)
    }
}

/// Validate and construct a queued run without publishing it.
///
/// A domain command can use this to batch the run atomically with related state
/// transitions. Pulse-cluster uses it to record every cancellation request and
/// its queued teardown together, so a client crash cannot leave a half-abort.
pub fn prepare_playbook_run(
    ctx: &CommandContext,
    command: RunPlaybook,
) -> Result<Run, CommandError> {
    // Confirm the playbook exists.
    ctx.exec_query_first(GetPlaybooksByIds {
        ids: vec![command.playbook_id.clone()],
    })?
    .ok_or_else(|| CommandError {
        tx: ctx.tx().to_string(),
        command_id: ctx.command_id.to_string(),
        message: format!("Playbook {} not found", command.playbook_id.0),
    })?;

    // Derive the resource set from the limit against the current inventory,
    // so the queue can serialize on shared clusters. No inventory (e.g. it
    // failed to discover) means no derivable resources → the run leases
    // nothing and is admitted immediately.
    let resources = match ctx.exec_query(GetAllInventorys {})?.into_iter().next() {
        Some(inv) => derive_resources(command.limit.as_deref(), &inv),
        None => Vec::new(),
    };

    Ok(Run {
        playbook_id: command.playbook_id,
        status: RunStatus::Queued,
        started_at: chrono::Utc::now().to_rfc3339(),
        cancel_requested_at: None,
        finished_at: None,
        exit_code: None,
        extra_vars: command.extra_vars,
        limit: command.limit,
        resources,
        host_stats: Default::default(),
        // Attribution of the LAUNCHING connection. Stamped here from the command's
        // request context — #[myko_client_id] emit-stamping only covers direct client
        // SET events, never a handler's server-side emit (lv-67c8). None = an
        // internal dispatch (saga/startup) with no originating connection.
        client_id: ctx.client_id().map(|c| c.to_string()),
        claimed_by: None, // set by ClaimRun when a runner wins the run
        id: RunId::from(uuid::Uuid::new_v4().to_string()),
    })
}

/// Request cancellation of a playbook execution.
///
/// A queued or not-yet-claimed run can become terminal immediately because no
/// executor exists. A claimed `Running` run stays non-terminal and keeps its
/// lease until the runner has terminated and reaped the playbook process group;
/// `FinishRun` is the acknowledgement that completes it as `Cancelled`.
/// Repeated requests and requests for an already-terminal run are idempotent.
#[myko_command]
pub struct CancelRun {
    pub run_id: RunId,
}

impl CommandHandler for CancelRun {
    fn execute(self, ctx: CommandContext) -> Result<(), CommandError> {
        let current = ctx
            .exec_query_first(GetRunsByIds {
                ids: vec![self.run_id.clone()],
            })?
            .ok_or_else(|| CommandError {
                tx: ctx.tx().to_string(),
                command_id: ctx.command_id.to_string(),
                message: format!("Run {} not found", self.run_id.0),
            })?;

        let Some(updated) = request_cancel_transition(&current, chrono::Utc::now().to_rfc3339())
        else {
            return Ok(());
        };
        ctx.emit_set(&updated)?;
        Ok(())
    }
}

pub fn request_cancel_transition(current: &Run, requested_at: String) -> Option<Run> {
    // Only an in-flight run can be cancelled; a finished run stays as it is.
    if !matches!(current.status, RunStatus::Queued | RunStatus::Running) {
        return None;
    }

    let mut updated = current.clone();
    let requested_at = updated.cancel_requested_at.clone().unwrap_or(requested_at);
    updated.cancel_requested_at = Some(requested_at.clone());

    // Queued runs and Running runs that have not been claimed cannot have
    // spawned a playbook. They need no runner acknowledgement.
    if updated.status == RunStatus::Queued || updated.claimed_by.is_none() {
        updated.status = RunStatus::Cancelled;
        updated.finished_at = Some(requested_at);
        updated.exit_code = None;
    }
    Some(updated)
}

/// Claim a run for execution — the compare-and-swap that makes exactly ONE runner
/// own each run. Succeeds only while the run is `Running` and UNCLAIMED; commands
/// execute serialized in the cell, so the check-and-set is atomic and two racing
/// runners cannot both win. A loser treats the error as "not mine, skip" — the
/// claim is the execution gate, not a failure.
///
/// Idempotent for the winner: re-claiming a run you already hold succeeds (a
/// subscription re-fire after winning must not error).
#[myko_command(RunId)]
pub struct ClaimRun {
    pub run_id: RunId,
    /// Stable runner identity (one per runner instance/host). A restarted runner
    /// with the same id uses it to recognize — and fail out — runs it left behind.
    pub runner_id: String,
}

impl CommandHandler for ClaimRun {
    // Returns the RunId (not unit) so the CALLER CAN AWAIT the result — a unit
    // command's response never populates the client's result cell, which made the
    // runner's claim-await time out on every claim that actually landed.
    fn execute(self, ctx: CommandContext) -> Result<RunId, CommandError> {
        let current = ctx
            .exec_query_first(GetRunsByIds {
                ids: vec![self.run_id.clone()],
            })?
            .ok_or_else(|| CommandError {
                tx: ctx.tx().to_string(),
                command_id: ctx.command_id.to_string(),
                message: format!("Run {} not found", self.run_id.0),
            })?;

        if current.status != RunStatus::Running {
            return Err(CommandError {
                tx: ctx.tx().to_string(),
                command_id: ctx.command_id.to_string(),
                message: format!(
                    "ClaimRun: run {} is {:?}, not Running — nothing to execute",
                    self.run_id.0, current.status
                ),
            });
        }
        match &current.claimed_by {
            Some(holder) if *holder == self.runner_id => return Ok(self.run_id), // already ours
            Some(holder) => {
                return Err(CommandError {
                    tx: ctx.tx().to_string(),
                    command_id: ctx.command_id.to_string(),
                    message: format!(
                        "ClaimRun: run {} already claimed by '{holder}'",
                        self.run_id.0
                    ),
                });
            }
            None => {}
        }

        let mut updated = (*current).clone();
        updated.claimed_by = Some(self.runner_id);
        ctx.emit_set(&updated)?;
        Ok(self.run_id)
    }
}

/// Pin a playbook to the sidebar's "Pinned" section. Idempotent: re-pinning
/// an already-pinned playbook is a no-op. Refuses to overwrite a locked
/// (config-defined) pin.
#[myko_command]
pub struct PinPlaybook {
    pub playbook_id: PlaybookId,
}

impl CommandHandler for PinPlaybook {
    fn execute(self, ctx: CommandContext) -> Result<(), CommandError> {
        ctx.exec_query_first(GetPlaybooksByIds {
            ids: vec![self.playbook_id.clone()],
        })?
        .ok_or_else(|| CommandError {
            tx: ctx.tx().to_string(),
            command_id: ctx.command_id.to_string(),
            message: format!("Playbook {} not found", self.playbook_id.0),
        })?;

        let pin_id = PinId::from(self.playbook_id.0.clone());
        if let Some(existing) = ctx.exec_query_first(GetPinsByIds {
            ids: vec![pin_id.clone()],
        })? {
            if existing.locked {
                return Ok(());
            }
        }

        let pin = Pin {
            playbook_id: self.playbook_id,
            locked: false,
            id: pin_id,
        };
        ctx.emit_set(&pin)?;
        Ok(())
    }
}

/// Append one line of playbook output to a run. Issued by the out-of-process
/// runner as it streams `ansible-playbook` output back to the cell (the
/// in-process runner writes `OutputLine`s directly via `ctx`). `seq` is the
/// runner's monotonic per-run counter across both stdout+stderr, so the UI can
/// order lines and jump to a parsed failure's source line.
#[myko_command]
pub struct AppendRunOutput {
    pub run_id: RunId,
    pub stream: OutputStream,
    pub line: String,
    pub seq: u64,
}

impl CommandHandler for AppendRunOutput {
    fn execute(self, ctx: CommandContext) -> Result<(), CommandError> {
        let entity = OutputLine {
            id: OutputLineId::from(format!("{}:{}", self.run_id.0, self.seq)),
            run_id: self.run_id,
            stream: self.stream,
            line: self.line,
            seq: self.seq,
        };
        ctx.emit_set(&entity)?;
        Ok(())
    }
}

/// Mark a run finished, reported by the runner only after `ansible-playbook`
/// and its process group have exited and output has drained.
///
/// The cell decides the terminal status: a cancellation request wins;
/// otherwise exit code 0 -> `Completed`, anything else -> `Failed`.
/// `exit_code` is `None` when the process was killed by a signal or never
/// spawned.
#[myko_command]
pub struct FinishRun {
    pub run_id: RunId,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub exit_code: Option<i32>,
}

impl CommandHandler for FinishRun {
    fn execute(self, ctx: CommandContext) -> Result<(), CommandError> {
        let current = ctx
            .exec_query_first(GetRunsByIds {
                ids: vec![self.run_id.clone()],
            })?
            .ok_or_else(|| CommandError {
                tx: ctx.tx().to_string(),
                command_id: ctx.command_id.to_string(),
                message: format!("Run {} not found", self.run_id.0),
            })?;

        let updated = finish_transition(&current, self.exit_code, chrono::Utc::now().to_rfc3339());
        ctx.emit_set(&updated)?;
        Ok(())
    }
}

fn finish_transition(current: &Run, exit_code: Option<i32>, finished_at: String) -> Run {
    let mut updated = current.clone();
    // A concurrent cancel wins, but terminal Cancelled is written only now:
    // this transition is the executor's proof that nothing is still running.
    if updated.cancel_requested_at.is_some() || updated.status == RunStatus::Cancelled {
        updated.status = RunStatus::Cancelled;
    } else {
        updated.status = if exit_code == Some(0) {
            RunStatus::Completed
        } else {
            RunStatus::Failed
        };
    }
    updated.exit_code = exit_code;
    updated.finished_at = Some(finished_at);
    updated
}

/// Declare the playbooks a runner has available. The runner scans its
/// `playbook_dir` and reports each as a `Playbook` entity (upsert by id).
/// Discovery lives where the playbook files are - the runner - not the cell,
/// so `RunPlaybook`'s existence check has something to resolve against.
#[myko_command]
pub struct DeclarePlaybooks {
    pub playbooks: Vec<Playbook>,
}

impl CommandHandler for DeclarePlaybooks {
    fn execute(self, ctx: CommandContext) -> Result<(), CommandError> {
        for pb in self.playbooks {
            ctx.emit_set(&pb)?;
        }
        Ok(())
    }
}

/// Record an EXTERNALLY-executed (CLI) op as a terminal `Run`.
///
/// Mutating cluster ops are run via `ansible-playbook` on the control node, so they
/// never dispatch through the cell and the runs wall never sees them. `RecordRun`
/// lets the deploy-side callback self-report the op at completion: it creates a
/// `Run` ALREADY TERMINAL (Completed/Failed/Cancelled). The runner executes every
/// `Running` run (it keys purely on `status == Running`, with no executor/owner
/// field to tell a self-declared run from a dispatched one), so a terminal run is
/// invisible to it — no double-execution, no `RunStatus` change. The wall renders
/// it like any other run (it reads `GetAllRuns`, source-agnostic); captured stdout
/// streams on via `AppendRunOutput(run_id, …)`.
///
/// Terminal-only BY CONTRACT: `Running`/`Queued` are rejected fail-loud — a
/// non-terminal self-declared run WOULD be picked up and re-executed by the runner.
/// Records hold no lease (`resources: []`): a completed op is a record, not a work
/// item. Lives here (not the cell) so every myko consumer — the cell handler AND
/// the deploy-side recorder client — shares ONE definition; no hand-synced contract.
#[myko_command(RunId)]
pub struct RecordRun {
    /// The playbook that was run (e.g. `ue-cluster-redeploy`).
    pub playbook_id: String,
    /// Terminal outcome ONLY — `completed` | `failed` | `cancelled`. `running` /
    /// `queued` are rejected: a non-terminal self-declared run would be executed
    /// by the runner.
    pub status: RunStatus,
    /// RFC3339 — when the external op started.
    pub started_at: String,
    /// RFC3339 — when it finished (a record is created post-hoc, so always known).
    pub finished_at: String,
    /// Process exit code, when captured.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub exit_code: Option<i32>,
    /// The op's vars (put operator / cluster-def context here for the "ran via CLI"
    /// attribution; the actor's WS identity is separately server-stamped as client_id).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub extra_vars: Option<Value>,
    /// The ansible `--limit`, when scoped.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub limit: Option<String>,
    /// Per-host ok/changed/failed/… recap, when the caller parsed the PLAY RECAP.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub host_stats: Option<HashMap<String, HostStats>>,
    /// ATTACH mode: the id of an EXISTING dispatched Run to enrich with this
    /// execution report (host_stats, limit) instead of minting a new record.
    ///
    /// This is the recap-decoupling fix: the Run LIFECYCLE (status, exit_code,
    /// finished_at) is owned by the runner via `FinishRun`; the EXECUTION DETAIL
    /// (per-host recap) is owned by whoever ran ansible and holds the stats object —
    /// the deploy-side callback. Before this field the callback could only mint a
    /// duplicate record or stay silent, so every cell-dispatched run recorded an
    /// empty recap. The runner advertises the id to the playbook process as
    /// `TACHYON_RUN_ID`; the callback passes it back here.
    ///
    /// Attach does NOT touch lifecycle fields and is accepted while the Run is
    /// still `Running` (ansible stats fire before the process exits).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub run_id: Option<RunId>,
}

impl CommandHandler for RecordRun {
    fn execute(self, ctx: CommandContext) -> Result<RunId, CommandError> {
        // ATTACH mode: enrich the named EXISTING run with the execution report.
        // Lifecycle fields (status / exit_code / finished_at) belong to the runner's
        // `FinishRun` and are deliberately untouched, so the two writers never race
        // over the same fact. Accepted while the run is still Running — ansible's
        // stats callback fires before the process exits.
        if let Some(rid) = self.run_id {
            let current = ctx
                .exec_query_first(GetRunsByIds { ids: vec![rid.clone()] })?
                .ok_or_else(|| CommandError {
                    tx: ctx.tx().to_string(),
                    command_id: ctx.command_id.to_string(),
                    message: format!(
                        "RecordRun attach: run {} not found — TACHYON_RUN_ID names a run this cell does not have",
                        rid.0
                    ),
                })?;
            let mut updated = (*current).clone();
            if let Some(hs) = self.host_stats {
                updated.host_stats = hs;
            }
            if updated.limit.is_none() {
                updated.limit = self.limit;
            }
            ctx.emit_set(&updated)?;
            return Ok(rid);
        }

        // MINT mode (no run_id): a deploy-side self-report of an operator-run op.
        // Terminal-only. A Running/Queued self-declared run would be grabbed and
        // re-executed by the runner (it watches status == Running); refuse it so a
        // record can never trigger a real op.
        if !matches!(
            self.status,
            RunStatus::Completed | RunStatus::Failed | RunStatus::Cancelled
        ) {
            return Err(CommandError {
                tx: ctx.tx().to_string(),
                command_id: ctx.command_id.to_string(),
                message: format!(
                    "RecordRun records COMPLETED external ops — status must be completed/failed/cancelled, not {:?} (a non-terminal run would be double-executed by the runner)",
                    self.status
                ),
            });
        }

        let id = RunId::from(uuid::Uuid::new_v4().to_string());
        let run = Run {
            id: id.clone(),
            playbook_id: PlaybookId::from(self.playbook_id),
            status: self.status,
            started_at: self.started_at,
            cancel_requested_at: None,
            finished_at: Some(self.finished_at),
            exit_code: self.exit_code,
            extra_vars: self.extra_vars,
            limit: self.limit,
            // A completed record holds no lease — it is not a work item.
            resources: Vec::new(),
            host_stats: self.host_stats.unwrap_or_default(),
            // The RECORDING connection (callback/CLI), stamped from the request context —
            // handler emits are never #[myko_client_id]-stamped (lv-67c8).
            client_id: ctx.client_id().map(|c| c.to_string()),
            claimed_by: None, // a terminal record is not a work item; nothing claims it
        };
        ctx.emit_set(&run)?;
        Ok(id)
    }
}

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

    /// Wire compat: a pre-0.6 payload (no runId) must still deserialize with
    /// run_id None — the mint path — so old callers keep working unchanged.
    #[test]
    fn record_run_without_run_id_is_mint_shape() {
        let old: RecordRun = serde_json::from_str(
            r#"{"playbookId":"ue-cluster-up","status":"completed",
                "startedAt":"2026-07-24T00:00:00Z","finishedAt":"2026-07-24T00:01:00Z"}"#,
        )
        .unwrap();
        assert!(old.run_id.is_none());
        assert!(old.host_stats.is_none());
    }

    #[test]
    fn record_run_with_run_id_is_attach_shape() {
        let attach: RecordRun = serde_json::from_str(
            r#"{"playbookId":"ue-cluster-up","status":"completed",
                "startedAt":"2026-07-24T00:00:00Z","finishedAt":"2026-07-24T00:01:00Z",
                "runId":"abc-123",
                "hostStats":{"render-01":{"ok":5,"changed":1,"unreachable":0,"failed":0,"skipped":0,"rescued":0,"ignored":0}}}"#,
        )
        .unwrap();
        assert_eq!(attach.run_id.as_ref().map(|r| r.0.as_ref()), Some("abc-123"));
        assert_eq!(attach.host_stats.unwrap()["render-01"].ok, 5);
    }
}

#[cfg(test)]
mod cancellation_tests {
    use super::{finish_transition, request_cancel_transition};
    use crate::{PlaybookId, Run, RunId, RunStatus};

    fn run(status: RunStatus, claimed_by: Option<&str>) -> Run {
        Run {
            id: RunId::from("run-1".to_string()),
            playbook_id: PlaybookId::from("playbook".to_string()),
            status,
            started_at: "start".into(),
            cancel_requested_at: None,
            finished_at: None,
            exit_code: None,
            extra_vars: None,
            limit: None,
            resources: vec!["render-01".into()],
            host_stats: Default::default(),
            client_id: None,
            claimed_by: claimed_by.map(str::to_owned),
        }
    }

    #[test]
    fn claimed_running_cancel_stays_nonterminal_until_finish_ack() {
        let requested =
            request_cancel_transition(&run(RunStatus::Running, Some("runner")), "cancel-at".into())
                .expect("running run is cancellable");
        assert_eq!(requested.status, RunStatus::Running);
        assert_eq!(requested.cancel_requested_at.as_deref(), Some("cancel-at"));
        assert!(requested.finished_at.is_none());

        let finished = finish_transition(&requested, None, "finished-at".into());
        assert_eq!(finished.status, RunStatus::Cancelled);
        assert_eq!(finished.finished_at.as_deref(), Some("finished-at"));
    }

    #[test]
    fn queued_and_unclaimed_running_cancel_immediately() {
        for initial in [run(RunStatus::Queued, None), run(RunStatus::Running, None)] {
            let cancelled =
                request_cancel_transition(&initial, "cancel-at".into()).expect("in-flight run");
            assert_eq!(cancelled.status, RunStatus::Cancelled);
            assert_eq!(cancelled.finished_at.as_deref(), Some("cancel-at"));
        }
    }

    #[test]
    fn repeated_cancel_is_idempotent_and_terminal_cancel_is_untouched() {
        let first =
            request_cancel_transition(&run(RunStatus::Running, Some("runner")), "first".into())
                .unwrap();
        let second = request_cancel_transition(&first, "second".into()).unwrap();
        assert_eq!(second.cancel_requested_at.as_deref(), Some("first"));

        let terminal = finish_transition(&second, None, "done".into());
        assert!(request_cancel_transition(&terminal, "later".into()).is_none());
    }

    #[test]
    fn finish_without_cancel_uses_exit_code() {
        assert_eq!(
            finish_transition(
                &run(RunStatus::Running, Some("runner")),
                Some(0),
                "done".into()
            )
            .status,
            RunStatus::Completed
        );
        assert_eq!(
            finish_transition(
                &run(RunStatus::Running, Some("runner")),
                Some(2),
                "done".into()
            )
            .status,
            RunStatus::Failed
        );
    }
}