meerkat-machine-schema 0.5.2

use indexmap::IndexMap;

use crate::{
    EffectDisposition, EffectDispositionRule, EffectEmit, EnumSchema, Expr, FieldInit, FieldSchema,
    Guard, InitSchema, InputMatch, InvariantSchema, MachineSchema, RustBinding, StateSchema,
    TransitionSchema, TypeRef, Update, VariantSchema,
};

pub fn comms_drain_lifecycle_machine() -> MachineSchema {
    MachineSchema {
        machine: "CommsDrainLifecycleMachine".into(),
        version: 1,
        rust: RustBinding {
            crate_name: "meerkat-runtime".into(),
            module: "generated::comms_drain_lifecycle".into(),
        },
        state: StateSchema {
            phase: EnumSchema {
                name: "CommsDrainPhase".into(),
                variants: vec![
                    variant("Inactive"),
                    variant("Starting"),
                    variant("Running"),
                    variant("ExitedRespawnable"),
                    variant("Stopped"),
                ],
            },
            fields: vec![field(
                "mode",
                TypeRef::Option(Box::new(TypeRef::Enum("CommsDrainMode".into()))),
            )],
            init: InitSchema {
                phase: "Inactive".into(),
                fields: vec![init("mode", Expr::None)],
            },
            terminal_phases: vec!["Stopped".into()],
        },
        inputs: EnumSchema {
            name: "CommsDrainLifecycleInput".into(),
            variants: vec![
                VariantSchema {
                    name: "EnsureRunning".into(),
                    fields: vec![field("mode", TypeRef::Enum("CommsDrainMode".into()))],
                },
                variant("TaskSpawned"),
                VariantSchema {
                    name: "TaskExited".into(),
                    fields: vec![field("reason", TypeRef::Enum("DrainExitReason".into()))],
                },
                variant("StopRequested"),
                variant("AbortObserved"),
            ],
        },
        effects: EnumSchema {
            name: "CommsDrainLifecycleEffect".into(),
            variants: vec![
                VariantSchema {
                    name: "SpawnDrainTask".into(),
                    fields: vec![field("mode", TypeRef::Enum("CommsDrainMode".into()))],
                },
                variant("AbortDrainTask"),
            ],
        },
        helpers: vec![],
        derived: vec![],
        invariants: vec![
            // Running or Starting implies mode is set
            InvariantSchema {
                name: "active_implies_mode_set".into(),
                expr: Expr::Or(vec![
                    Expr::And(vec![
                        Expr::Neq(
                            Box::new(Expr::CurrentPhase),
                            Box::new(Expr::Phase("Starting".into())),
                        ),
                        Expr::Neq(
                            Box::new(Expr::CurrentPhase),
                            Box::new(Expr::Phase("Running".into())),
                        ),
                    ]),
                    Expr::Neq(Box::new(Expr::Field("mode".into())), Box::new(Expr::None)),
                ]),
            },
        ],
        transitions: vec![
            // (Inactive, EnsureRunning{mode}) -> Starting
            TransitionSchema {
                name: "EnsureRunningFromInactive".into(),
                from: vec!["Inactive".into()],
                on: InputMatch {
                    variant: "EnsureRunning".into(),
                    bindings: vec!["mode".into()],
                },
                guards: vec![],
                updates: vec![Update::Assign {
                    field: "mode".into(),
                    expr: Expr::Some(Box::new(Expr::Binding("mode".into()))),
                }],
                to: "Starting".into(),
                emit: vec![EffectEmit {
                    variant: "SpawnDrainTask".into(),
                    fields: IndexMap::from([("mode".into(), Expr::Binding("mode".into()))]),
                }],
            },
            // (Starting, TaskSpawned) -> Running
            TransitionSchema {
                name: "TaskSpawnedFromStarting".into(),
                from: vec!["Starting".into()],
                on: InputMatch {
                    variant: "TaskSpawned".into(),
                    bindings: vec![],
                },
                guards: vec![],
                updates: vec![],
                to: "Running".into(),
                emit: vec![],
            },
            // (Starting, TaskExited{reason=Failed}) when mode=PersistentHost -> ExitedRespawnable
            TransitionSchema {
                name: "TaskExitedFromStartingRespawnable".into(),
                from: vec!["Starting".into()],
                on: InputMatch {
                    variant: "TaskExited".into(),
                    bindings: vec!["reason".into()],
                },
                guards: vec![
                    Guard {
                        name: "reason_is_failed".into(),
                        expr: Expr::Eq(
                            Box::new(Expr::Binding("reason".into())),
                            Box::new(drain_exit_reason(DrainExitReasonVariant::Failed)),
                        ),
                    },
                    Guard {
                        name: "mode_is_persistent_host".into(),
                        expr: Expr::Eq(
                            Box::new(Expr::Field("mode".into())),
                            Box::new(Expr::Some(Box::new(comms_drain_mode(
                                CommsDrainModeVariant::PersistentHost,
                            )))),
                        ),
                    },
                ],
                updates: vec![],
                to: "ExitedRespawnable".into(),
                emit: vec![],
            },
            // (Starting, TaskExited{reason}) -> Stopped (non-respawnable cases)
            TransitionSchema {
                name: "TaskExitedFromStartingStopped".into(),
                from: vec!["Starting".into()],
                on: InputMatch {
                    variant: "TaskExited".into(),
                    bindings: vec!["reason".into()],
                },
                guards: vec![Guard {
                    name: "not_respawnable".into(),
                    expr: Expr::Or(vec![
                        // reason is not Failed
                        Expr::Neq(
                            Box::new(Expr::Binding("reason".into())),
                            Box::new(drain_exit_reason(DrainExitReasonVariant::Failed)),
                        ),
                        // OR mode is not PersistentHost
                        Expr::Neq(
                            Box::new(Expr::Field("mode".into())),
                            Box::new(Expr::Some(Box::new(comms_drain_mode(
                                CommsDrainModeVariant::PersistentHost,
                            )))),
                        ),
                    ]),
                }],
                updates: vec![],
                to: "Stopped".into(),
                emit: vec![],
            },
            // (Running, TaskExited{reason=Failed}) when mode=PersistentHost -> ExitedRespawnable
            TransitionSchema {
                name: "TaskExitedFromRunningRespawnable".into(),
                from: vec!["Running".into()],
                on: InputMatch {
                    variant: "TaskExited".into(),
                    bindings: vec!["reason".into()],
                },
                guards: vec![
                    Guard {
                        name: "reason_is_failed".into(),
                        expr: Expr::Eq(
                            Box::new(Expr::Binding("reason".into())),
                            Box::new(drain_exit_reason(DrainExitReasonVariant::Failed)),
                        ),
                    },
                    Guard {
                        name: "mode_is_persistent_host".into(),
                        expr: Expr::Eq(
                            Box::new(Expr::Field("mode".into())),
                            Box::new(Expr::Some(Box::new(comms_drain_mode(
                                CommsDrainModeVariant::PersistentHost,
                            )))),
                        ),
                    },
                ],
                updates: vec![],
                to: "ExitedRespawnable".into(),
                emit: vec![],
            },
            // (Running, TaskExited{reason}) -> Stopped (non-respawnable cases)
            TransitionSchema {
                name: "TaskExitedFromRunningStopped".into(),
                from: vec!["Running".into()],
                on: InputMatch {
                    variant: "TaskExited".into(),
                    bindings: vec!["reason".into()],
                },
                guards: vec![Guard {
                    name: "not_respawnable".into(),
                    expr: Expr::Or(vec![
                        Expr::Neq(
                            Box::new(Expr::Binding("reason".into())),
                            Box::new(drain_exit_reason(DrainExitReasonVariant::Failed)),
                        ),
                        Expr::Neq(
                            Box::new(Expr::Field("mode".into())),
                            Box::new(Expr::Some(Box::new(comms_drain_mode(
                                CommsDrainModeVariant::PersistentHost,
                            )))),
                        ),
                    ]),
                }],
                updates: vec![],
                to: "Stopped".into(),
                emit: vec![],
            },
            // (Running, StopRequested) -> Stopped
            TransitionSchema {
                name: "StopRequestedFromRunning".into(),
                from: vec!["Running".into()],
                on: InputMatch {
                    variant: "StopRequested".into(),
                    bindings: vec![],
                },
                guards: vec![],
                updates: vec![],
                to: "Stopped".into(),
                emit: vec![EffectEmit {
                    variant: "AbortDrainTask".into(),
                    fields: IndexMap::new(),
                }],
            },
            // (Starting, StopRequested) -> Stopped
            TransitionSchema {
                name: "StopRequestedFromStarting".into(),
                from: vec!["Starting".into()],
                on: InputMatch {
                    variant: "StopRequested".into(),
                    bindings: vec![],
                },
                guards: vec![],
                updates: vec![],
                to: "Stopped".into(),
                emit: vec![EffectEmit {
                    variant: "AbortDrainTask".into(),
                    fields: IndexMap::new(),
                }],
            },
            // (ExitedRespawnable, EnsureRunning{mode}) -> Starting
            TransitionSchema {
                name: "EnsureRunningFromExitedRespawnable".into(),
                from: vec!["ExitedRespawnable".into()],
                on: InputMatch {
                    variant: "EnsureRunning".into(),
                    bindings: vec!["mode".into()],
                },
                guards: vec![],
                updates: vec![Update::Assign {
                    field: "mode".into(),
                    expr: Expr::Some(Box::new(Expr::Binding("mode".into()))),
                }],
                to: "Starting".into(),
                emit: vec![EffectEmit {
                    variant: "SpawnDrainTask".into(),
                    fields: IndexMap::from([("mode".into(), Expr::Binding("mode".into()))]),
                }],
            },
            // (ExitedRespawnable, StopRequested) -> Stopped
            TransitionSchema {
                name: "StopRequestedFromExitedRespawnable".into(),
                from: vec!["ExitedRespawnable".into()],
                on: InputMatch {
                    variant: "StopRequested".into(),
                    bindings: vec![],
                },
                guards: vec![],
                updates: vec![],
                to: "Stopped".into(),
                emit: vec![],
            },
            // (Stopped, EnsureRunning{mode}) -> Starting
            TransitionSchema {
                name: "EnsureRunningFromStopped".into(),
                from: vec!["Stopped".into()],
                on: InputMatch {
                    variant: "EnsureRunning".into(),
                    bindings: vec!["mode".into()],
                },
                guards: vec![],
                updates: vec![Update::Assign {
                    field: "mode".into(),
                    expr: Expr::Some(Box::new(Expr::Binding("mode".into()))),
                }],
                to: "Starting".into(),
                emit: vec![EffectEmit {
                    variant: "SpawnDrainTask".into(),
                    fields: IndexMap::from([("mode".into(), Expr::Binding("mode".into()))]),
                }],
            },
            // (Running|Starting, AbortObserved) -> Stopped
            TransitionSchema {
                name: "AbortObservedFromActive".into(),
                from: vec!["Running".into(), "Starting".into()],
                on: InputMatch {
                    variant: "AbortObserved".into(),
                    bindings: vec![],
                },
                guards: vec![],
                updates: vec![],
                to: "Stopped".into(),
                emit: vec![],
            },
        ],
        ci_step_limit: None,
        effect_dispositions: vec![
            EffectDispositionRule {
                effect_variant: "SpawnDrainTask".into(),
                disposition: EffectDisposition::Local,
                handoff_protocol: Some("comms_drain_spawn".into()),
            },
            EffectDispositionRule {
                effect_variant: "AbortDrainTask".into(),
                disposition: EffectDisposition::Local,
                handoff_protocol: Some("comms_drain_abort".into()),
            },
        ],
    }
}

fn field(name: &str, ty: TypeRef) -> FieldSchema {
    FieldSchema {
        name: name.into(),
        ty,
    }
}

fn init(field: &str, expr: Expr) -> FieldInit {
    FieldInit {
        field: field.into(),
        expr,
    }
}

fn variant(name: &str) -> VariantSchema {
    VariantSchema {
        name: name.into(),
        fields: vec![],
    }
}

// ---------------------------------------------------------------------------
// Local enum variant helpers
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy)]
enum CommsDrainModeVariant {
    #[allow(dead_code)]
    AttachedSession,
    PersistentHost,
}

#[derive(Debug, Clone, Copy)]
enum DrainExitReasonVariant {
    Failed,
    #[allow(dead_code)]
    Dismissed,
    #[allow(dead_code)]
    Aborted,
    #[allow(dead_code)]
    SessionShutdown,
}

fn comms_drain_mode(variant: CommsDrainModeVariant) -> Expr {
    Expr::NamedVariant {
        enum_name: "CommsDrainMode".into(),
        variant: match variant {
            CommsDrainModeVariant::AttachedSession => "AttachedSession".into(),
            CommsDrainModeVariant::PersistentHost => "PersistentHost".into(),
        },
    }
}

fn drain_exit_reason(variant: DrainExitReasonVariant) -> Expr {
    Expr::NamedVariant {
        enum_name: "DrainExitReason".into(),
        variant: match variant {
            DrainExitReasonVariant::Dismissed => "Dismissed".into(),
            DrainExitReasonVariant::Failed => "Failed".into(),
            DrainExitReasonVariant::Aborted => "Aborted".into(),
            DrainExitReasonVariant::SessionShutdown => "SessionShutdown".into(),
        },
    }
}