//! travel_booking_service — orchestrating **real third-party resources**.
//!
//! This is the "what it's actually like to build a service on mire" example.
//! Unlike a topology-only demo, the saga here talks to genuine external
//! systems: a set of mock third-party providers running as a real **axum**
//! HTTP backend, called over the wire with a real **reqwest** client.
//!
//! ## The model it demonstrates
//!
//! Every step is a pure, *synchronous* **`decide`** `fn` that folds
//! state/outputs into a request event and suspends — it cannot perform IO (it
//! isn't `async`), so a crash-unsafe inline effect won't even compile. The
//! *effects* (the provider IO) live entirely off the saga's hot path: each
//! forward request is fulfilled **durably via the outbox** by `HttpPublisher`,
//! which calls the provider and `deliver`s the response back, mapped to the
//! step's output by `.on_response(...)`. This is the right shape for a sharded
//! distributed service — external IO survives crashes, runs off the
//! trigger/deliver path, and never executes under recovery contention.
//!
//! Idempotency is the contract that makes at-least-once orchestration safe:
//! every provider call carries a stable per-step key, and each mock service
//! dedupes on it — so a republish after a crash, or a re-drive, hits the
//! *same* key and the real-world effect happens once.
//!
//! Topology: `flight` ∥ `hotel` → `payment` → `confirm`. The off-hot-path,
//! callback-shaped external IO — `flight` / `hotel` / `payment` — emits a
//! request with **no registered effect**, so the runner routes it to the
//! outbox; the `HttpPublisher` (driven by the worker's outbox drain) fulfils it
//! — `deliver`ing the response on success, and leaving a never-settling payment
//! to time out (→ cascade). `confirm` and the **compensations** are fulfilled
//! IN-PROCESS via co-located `.effect(...)` handlers: `confirm` so a delivered
//! payment tickets the carrier inline (a permanent rejection panics → cascade),
//! and compensations because the cascade dispatches them directly (the runner
//! requires a registered handler). The failure-mirror of `deliver` —
//! [`PgSagaRunner::deliver_failure`] — lets an external publisher fail an
//! awaiting step too, should a forward request hit a permanent error.
//!
//! The accompanying test (`tests/travel_booking_service.rs`) stress-tests the
//! assumptions: flaky providers (transient 503s), concurrent bookings, a
//! sharded worker pool, a compensating timeout, repeated worker crashes, and a
//! cancellation that lands *while a payment capture is in flight* (the
//! publish-time TOCTOU — the publisher refunds the orphan charge) — asserting
//! every booking converges to its correct terminal, money-safe, **exactly
//! once**.
//!
//! Run: `DATABASE_URL=… cargo run -p mire-sagas --example travel_booking_service`

#![allow(dead_code)]

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};

use axum::{Json, Router, extract::State as AxState, http::StatusCode, routing::post};
use mire::{Aggregate, EventData, EventStore};
use mire_sagas::{
    Dep, Emit, EmitPublisher, IdempotencyKey, OutboxMessage, Output, PgSagaRunner, Req, Saga,
    SagaBuilder, SagaInstanceStatus, SagaWorker, Start, State, comp, task,
};
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;

// =====================================================================
// Domain: the triggering event + the saga's state + per-step outputs.
// =====================================================================

#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "booking.requested")]
pub struct BookingRequested {
    pub booking_id: String,
    pub flight_id: String,
    pub hotel_id: String,
    pub amount_cents: i64,
}

#[derive(Default)]
pub struct Booking;
impl Aggregate for Booking {
    type Event = BookingRequested;
    fn stream_category() -> &'static str {
        "tbs-booking"
    }
    fn apply(&mut self, _: &Self::Event) {}
}

#[derive(Clone, Default, Serialize, Deserialize)]
pub struct TravelBooking {
    pub booking_id: String,
    pub flight_id: String,
    pub hotel_id: String,
    pub amount_cents: i64,
}

#[derive(Clone, Serialize, Deserialize)]
pub struct FlightReserved {
    pub hold_id: String,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct HotelReserved {
    pub hold_id: String,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct PaymentSettled {
    pub payment_id: String,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct Confirmed {
    pub ticket: String,
}

// Payment is ASYNC: the step emits a capture *request* and suspends; the
// provider's result arrives later as a `PaymentCaptured` event delivered back
// into the saga (see `HttpPublisher` + the deliver loop). These two events are
// the request/response pair for that round-trip.
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.payment.capture-request")]
pub struct RequestPaymentCapture {
    pub booking_id: String,
    pub amount_cents: i64,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.payment.captured")]
pub struct PaymentCaptured {
    pub payment_id: String,
}

// A PUSH event: an unsolicited signal the saga never requested. A customer (or
// an upstream system) cancels the booking mid-flight; whatever step is in
// progress, the saga rolls back — `triggers_compensation` runs the cascade
// over the completed steps. This is the event/push direction (vs the
// request/response pull of the provider calls).
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.booking.cancel")]
pub struct CancelBooking {
    pub reason: String,
}

// The request/response pairs for the IN-PROCESS effects. Each `decide` step
// emits a `Request*` event; its co-located `.effect(...)` handler does the
// provider IO and returns the matching response, which `.on_response(...)` maps
// to the step's typed output. This is the same emit→fulfil→fold round-trip as
// payment — the only difference is *where* the fulfilment runs: these handlers
// are registered on the saga (`.effect`) and the runner auto-invokes them
// in-process, whereas payment is fulfilled off the hot path by `HttpPublisher`.
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.flight.hold-request")]
pub struct RequestFlightHold {
    pub booking_id: String,
    pub flight_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.flight.held")]
pub struct FlightHeld {
    pub hold_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.hotel.hold-request")]
pub struct RequestHotelHold {
    pub booking_id: String,
    pub hotel_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.hotel.held")]
pub struct HotelHeld {
    pub hold_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.ticket.issue-request")]
pub struct IssueCarrierTicket {
    pub booking_id: String,
    pub payment_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.ticket.issued")]
pub struct TicketIssued {
    pub ticket: String,
}

// Compensation request/response pairs — the "undo" effects (release a hold,
// refund a charge) are full effects too, fulfilled by `.effect(...)` handlers.
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.flight.release-request")]
pub struct ReleaseFlightHold {
    pub hold_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.flight.released")]
pub struct FlightReleased {}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.hotel.release-request")]
pub struct ReleaseHotelHold {
    pub hold_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.hotel.released")]
pub struct HotelReleased {}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.payment.refund-request")]
pub struct RequestPaymentRefund {
    pub payment_id: String,
}
#[derive(Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "tbs.payment.refunded")]
pub struct PaymentRefunded {}

// =====================================================================
// The provider client. A single facade injected into effects as `Dep`.
// In a real service this wraps your HTTP/queue SDKs; here it's reqwest
// against the mock backend. Every call passes the saga's idempotency
// token so the provider can dedupe.
// =====================================================================

#[derive(Clone)]
pub struct Providers {
    http: reqwest::Client,
    base: String,
}

#[derive(Serialize)]
struct HoldFlightReq<'a> {
    idem: &'a str,
    flight_id: &'a str,
}
#[derive(Serialize)]
struct HoldHotelReq<'a> {
    idem: &'a str,
    hotel_id: &'a str,
}
#[derive(Serialize)]
struct CaptureReq<'a> {
    idem: &'a str,
    booking_id: &'a str,
    amount_cents: i64,
}
#[derive(Serialize)]
struct TicketReq<'a> {
    idem: &'a str,
    booking_id: &'a str,
    payment_id: &'a str,
}
#[derive(Deserialize)]
struct HoldResp {
    hold_id: String,
}
#[derive(Deserialize)]
struct CaptureResp {
    /// `false` models a capture the provider accepted but hasn't settled yet
    /// (it would call back later) — for a doomed booking it never does, so the
    /// saga's payment step times out and compensates.
    captured: bool,
    payment_id: Option<String>,
}
#[derive(Deserialize)]
struct TicketResp {
    ticket: String,
}

impl Providers {
    pub fn new(base: impl Into<String>) -> Self {
        Self {
            http: reqwest::Client::new(),
            base: base.into(),
        }
    }

    async fn post_json<B: Serialize, R: for<'de> Deserialize<'de>>(
        &self,
        path: &str,
        body: &B,
    ) -> anyhow::Result<R> {
        // Retry transient 503s — providers are flaky in the stress test. The
        // idempotency key means a retried request that the provider already
        // processed returns the same result, so retrying never double-books.
        let url = format!("{}{path}", self.base);
        let mut attempt = 0u32;
        loop {
            attempt += 1;
            let resp = self.http.post(&url).json(body).send().await?;
            if resp.status() == reqwest::StatusCode::SERVICE_UNAVAILABLE && attempt < 12 {
                tokio::time::sleep(std::time::Duration::from_millis(15)).await;
                continue;
            }
            return Ok(resp.error_for_status()?.json::<R>().await?);
        }
    }

    pub async fn hold_flight(&self, flight_id: &str, idem: &str) -> anyhow::Result<String> {
        let r: HoldResp = self
            .post_json("/flight/hold", &HoldFlightReq { idem, flight_id })
            .await?;
        Ok(r.hold_id)
    }
    pub async fn hold_hotel(&self, hotel_id: &str, idem: &str) -> anyhow::Result<String> {
        let r: HoldResp = self
            .post_json("/hotel/hold", &HoldHotelReq { idem, hotel_id })
            .await?;
        Ok(r.hold_id)
    }
    /// Submit a capture. `Some(payment_id)` = settled now; `None` = accepted
    /// but unsettled (the saga will time out waiting and compensate).
    pub async fn capture_payment(
        &self,
        booking_id: &str,
        amount_cents: i64,
        idem: &str,
    ) -> anyhow::Result<Option<String>> {
        let r: CaptureResp = self
            .post_json(
                "/payment/capture",
                &CaptureReq {
                    idem,
                    booking_id,
                    amount_cents,
                },
            )
            .await?;
        Ok(if r.captured { r.payment_id } else { None })
    }
    pub async fn issue_ticket(
        &self,
        booking_id: &str,
        payment_id: &str,
        idem: &str,
    ) -> anyhow::Result<String> {
        let r: TicketResp = self
            .post_json(
                "/carrier/ticket",
                &TicketReq {
                    idem,
                    booking_id,
                    payment_id,
                },
            )
            .await?;
        Ok(r.ticket)
    }

    pub async fn release_flight(&self, hold_id: &str, idem: &str) -> anyhow::Result<()> {
        let _: serde_json::Value = self
            .post_json("/flight/release", &ReleaseReq { idem, hold_id })
            .await?;
        Ok(())
    }
    pub async fn release_hotel(&self, hold_id: &str, idem: &str) -> anyhow::Result<()> {
        let _: serde_json::Value = self
            .post_json("/hotel/release", &ReleaseReq { idem, hold_id })
            .await?;
        Ok(())
    }
    pub async fn refund_payment(&self, payment_id: &str, idem: &str) -> anyhow::Result<()> {
        let _: serde_json::Value = self
            .post_json("/payment/refund", &RefundReq { idem, payment_id })
            .await?;
        Ok(())
    }
}

#[derive(Serialize)]
struct ReleaseReq<'a> {
    idem: &'a str,
    hold_id: &'a str,
}
#[derive(Serialize)]
struct RefundReq<'a> {
    idem: &'a str,
    payment_id: &'a str,
}

// =====================================================================
// The saga: each step is a pure, sync `decide` that EMITS a request; the
// provider IO lives in an async `effect` handler co-located on the step via
// `.effect(...)`. The runner auto-invokes the effect in-process and folds its
// response back through `.on_response(...)`. Orchestration (order, deps,
// terminal) is in `build()`; effects (the IO) are the `.effect` handlers; the
// providers live behind HTTP. A `decide` cannot accidentally perform IO — it
// is not `async`, so a crash-unsafe inline effect simply won't compile.
// =====================================================================

// --- decide steps: sync, pure, emit a request -------------------------

fn reserve_flight(State(s): State<TravelBooking>) -> Emit {
    Emit::event(RequestFlightHold {
        booking_id: s.booking_id.clone(),
        flight_id: s.flight_id.clone(),
    })
}

fn reserve_hotel(State(s): State<TravelBooking>) -> Emit {
    Emit::event(RequestHotelHold {
        booking_id: s.booking_id.clone(),
        hotel_id: s.hotel_id.clone(),
    })
}

// Payment is the ONE step fulfilled OFF the hot path: it emits a capture
// request and suspends with a deadline, and the `HttpPublisher` (driven by the
// outbox drain) performs the provider call and delivers `PaymentCaptured` back
// later — the "submit, park, resume on the result" (callback) shape. Because no
// `.effect(...)` is registered for `RequestPaymentCapture`, the runner routes
// it to the outbox instead of dispatching in-process. The timeout turns a
// never-settling capture into a compensating failure.
fn request_payment(State(s): State<TravelBooking>) -> Emit {
    Emit::event(RequestPaymentCapture {
        booking_id: s.booking_id.clone(),
        amount_cents: s.amount_cents,
    })
    .timeout(std::time::Duration::from_secs(2))
}

// `confirm` reads the settled payment id from its upstream output and emits the
// ticket-issue request carrying it; the effect handler only needs the request.
fn request_ticket(State(s): State<TravelBooking>, Output(payment): Output<PaymentSettled>) -> Emit {
    Emit::event(IssueCarrierTicket {
        booking_id: s.booking_id.clone(),
        payment_id: payment.payment_id,
    })
}

// --- compensation decides: sync, pure, emit an undo request -----------

fn release_flight(Output(r): Output<FlightReserved>) -> Emit {
    Emit::event(ReleaseFlightHold { hold_id: r.hold_id })
}

fn release_hotel(Output(r): Output<HotelReserved>) -> Emit {
    Emit::event(ReleaseHotelHold { hold_id: r.hold_id })
}

// Compensating a COMPLETED async step: if a later step fails after payment
// settled (e.g. the carrier rejects the ticket), the cascade refunds the
// captured payment. Reads the payment id from the step's output.
fn refund_payment(Output(p): Output<PaymentSettled>) -> Emit {
    Emit::event(RequestPaymentRefund {
        payment_id: p.payment_id,
    })
}

// --- in-process effect handlers: async, do the IO on the hot path -----
//
// The flight/hotel holds and the payment capture are fulfilled OFF the hot path
// by `HttpPublisher` (below) via the durable outbox. The remaining in-process
// effects are `confirm` (so a delivered payment tickets the carrier inline) and
// the compensations (which the runner requires to be in-process). Each carries
// its step's `IdempotencyKey` into the provider call so the mock backend
// dedupes; a failed call panics → the step/compensation fails → cascade (a
// failed *compensation* ends `FailedTerminalNeedsIntervention`).

async fn handle_issue_ticket(
    Req(r): Req<IssueCarrierTicket>,
    Dep(p): Dep<Providers>,
    key: IdempotencyKey,
) -> TicketIssued {
    let ticket = p
        .issue_ticket(&r.booking_id, &r.payment_id, &key.token())
        .await
        .expect("carrier provider call failed");
    TicketIssued { ticket }
}

async fn handle_release_flight(
    Req(r): Req<ReleaseFlightHold>,
    Dep(p): Dep<Providers>,
    key: IdempotencyKey,
) -> FlightReleased {
    p.release_flight(&r.hold_id, &key.token())
        .await
        .expect("flight release failed");
    FlightReleased {}
}

async fn handle_release_hotel(
    Req(r): Req<ReleaseHotelHold>,
    Dep(p): Dep<Providers>,
    key: IdempotencyKey,
) -> HotelReleased {
    p.release_hotel(&r.hold_id, &key.token())
        .await
        .expect("hotel release failed");
    HotelReleased {}
}

async fn handle_refund_payment(
    Req(r): Req<RequestPaymentRefund>,
    Dep(p): Dep<Providers>,
    key: IdempotencyKey,
) -> PaymentRefunded {
    p.refund_payment(&r.payment_id, &key.token())
        .await
        .expect("payment refund failed");
    PaymentRefunded {}
}

/// The travel-booking topology, **generic over the saga marker type** so tests
/// can instantiate it under distinct saga *names*. Recovery (and timeout/outbox
/// scans) are saga-WIDE, so two heavy concurrent tests sharing one name would
/// recover each other's instances; a distinct name per test isolates them.
/// `TravelBookingSaga` below is the canonical instance the example/main use.
pub fn build_travel_topology<S: Saga<State = TravelBooking>>(b: SagaBuilder<S>) -> SagaBuilder<S> {
    b.start_on::<Booking, _>(|e: &BookingRequested| {
        Some(Start {
            instance_id: e.booking_id.clone(),
            state: TravelBooking {
                booking_id: e.booking_id.clone(),
                flight_id: e.flight_id.clone(),
                hotel_id: e.hotel_id.clone(),
                amount_cents: e.amount_cents,
            },
        })
    })
    // Every step is a sync `decide` that EMITS a request and suspends; the
    // request is fulfilled OFF the hot path by `HttpPublisher` (via the durable
    // outbox), which calls the provider and `deliver`s the response back —
    // mapped to the step's output by `.on_response(...)`. No forward step
    // registers an `.effect(...)`, so the runner routes each emit to the outbox
    // rather than dispatching in-process. This is the right shape for a sharded
    // distributed service: external IO is durable and off the trigger/deliver
    // hot path, so it survives crashes and never runs under recovery contention.
    //
    // COMPENSATIONS are the exception: the cascade dispatches compensation
    // effects IN-PROCESS (the runner requires a registered handler for a
    // compensation emit), so each rollback IO lives in a co-located
    // `.effect(...)` and reports its own success/failure inline.
    //
    // flight ∥ hotel: no dependency, both suspend at once; their holds publish
    // and complete concurrently.
    .step(
        "flight",
        task(reserve_flight)
            .on_response(|h: FlightHeld| FlightReserved { hold_id: h.hold_id })
            .compensate(
                comp(release_flight)
                    .effect(handle_release_flight)
                    .on_response::<FlightReleased>(),
            ),
    )
    .step(
        "hotel",
        task(reserve_hotel)
            .on_response(|h: HotelHeld| HotelReserved { hold_id: h.hold_id })
            .compensate(
                comp(release_hotel)
                    .effect(handle_release_hotel)
                    .on_response::<HotelReleased>(),
            ),
    )
    // payment needs both holds; it suspends awaiting `PaymentCaptured`, mapped
    // to `PaymentSettled`. Its compensation refunds via an in-process effect.
    .step(
        "payment",
        task(request_payment)
            .after(["flight", "hotel"])
            .on_response(|c: PaymentCaptured| PaymentSettled {
                payment_id: c.payment_id,
            })
            .compensate(
                comp(refund_payment)
                    .effect(handle_refund_payment)
                    .on_response::<PaymentRefunded>(),
            ),
    )
    // confirm is the ONE forward step fulfilled IN-PROCESS (`.effect`): it
    // completes inline when payment is delivered, so a settled payment tickets
    // the carrier in the same step transition (no extra outbox round-trip). A
    // permanent carrier rejection panics the effect → confirm fails → cascade
    // (refund + release). Terminal success (endpoint, no comp).
    .step(
        "confirm",
        task(request_ticket)
            .effect(handle_issue_ticket)
            .after("payment")
            .on_response(|t: TicketIssued| Confirmed { ticket: t.ticket })
            .terminal_success(),
    )
    // A pushed cancellation cancels the in-flight saga and compensates
    // whatever completed.
    .signal::<CancelBooking>()
    .triggers_compensation()
}

pub struct TravelBookingSaga;
impl Saga for TravelBookingSaga {
    type State = TravelBooking;
    fn name() -> &'static str {
        "travel-booking-service"
    }
    fn build(b: SagaBuilder<Self>) -> SagaBuilder<Self> {
        build_travel_topology(b)
    }
}

// =====================================================================
// HttpPublisher — drains the saga outbox and performs the async provider
// call OFF the hot path, then delivers the result back into the saga.
//
// Crucially it does both the provider call AND the deliver INSIDE publish:
// `process_outbox` only marks the outbox row published once publish returns
// Ok. So a worker that crashes after calling the provider but before the row
// is marked leaves it unpublished — a fresh worker re-publishes, which
// re-POSTs (the provider dedupes on the idempotency key → no second charge)
// and re-delivers (idempotent: the await is already consumed → a no-op). The
// durable outbox, not an in-memory queue, is what survives a crash.
// =====================================================================

pub struct HttpPublisher<S: Saga<State = TravelBooking> = TravelBookingSaga> {
    providers: Providers,
    deliver: Arc<PgSagaRunner<S>>,
}

// Manual `Clone` — the saga marker `S` is a zero-sized type that doesn't derive
// `Clone`, and we only hold a `Providers` (Clone) and an `Arc` (Clone), so the
// derive's spurious `S: Clone` bound is avoided.
impl<S: Saga<State = TravelBooking>> Clone for HttpPublisher<S> {
    fn clone(&self) -> Self {
        Self {
            providers: self.providers.clone(),
            deliver: self.deliver.clone(),
        }
    }
}

impl<S: Saga<State = TravelBooking>> HttpPublisher<S> {
    /// Build a publisher over `providers` that delivers results back through
    /// `deliver`. `spawn_worker` wires this up automatically; tests use it to
    /// drive `publish` directly and control timing.
    pub fn new(providers: Providers, deliver: Arc<PgSagaRunner<S>>) -> Self {
        Self { providers, deliver }
    }
}

impl<S: Saga<State = TravelBooking>> EmitPublisher for HttpPublisher<S> {
    fn publish(&self, msg: OutboxMessage<'_>) -> impl Future<Output = anyhow::Result<()>> + Send {
        let providers = self.providers.clone();
        let deliver = self.deliver.clone();
        let event_type = msg.event_type.to_string();
        let payload = msg.payload.clone();
        let instance = msg.instance_id.to_string();
        // Stable per-step dedup key (the IdempotencyKey shape) so a republish
        // after a crash hits the same provider key — capture happens once.
        let idem = format!("{}/{}", msg.instance_id, msg.step);
        // Skip fulfilling a request whose booking already went terminal (e.g. a
        // pushed CancelBooking rolled it back while the request sat in the
        // outbox) — don't acquire a hold / issue a ticket for a dead booking.
        async fn booking_terminal<S: Saga<State = TravelBooking>>(
            deliver: &PgSagaRunner<S>,
            instance: &str,
        ) -> bool {
            matches!(deliver.load(instance).await, Ok(Some(env)) if env.status.is_terminal())
        }
        async move {
            match event_type.as_str() {
                "tbs.flight.hold-request" => {
                    if booking_terminal(&deliver, &instance).await {
                        return Ok(());
                    }
                    let req: RequestFlightHold = serde_json::from_value(payload)?;
                    let hold_id = providers.hold_flight(&req.flight_id, &idem).await?;
                    deliver.deliver(&instance, &FlightHeld { hold_id }).await?;
                    Ok(())
                }
                "tbs.hotel.hold-request" => {
                    if booking_terminal(&deliver, &instance).await {
                        return Ok(());
                    }
                    let req: RequestHotelHold = serde_json::from_value(payload)?;
                    let hold_id = providers.hold_hotel(&req.hotel_id, &idem).await?;
                    deliver.deliver(&instance, &HotelHeld { hold_id }).await?;
                    Ok(())
                }
                "tbs.payment.capture-request" => {
                    // Don't charge a booking that's already terminal — e.g. one
                    // a pushed CancelBooking rolled back while its capture was
                    // still pending in the outbox. (A small TOCTOU window
                    // remains for a cancel landing *during* publish; a real
                    // system pairs this with payment compensation/refund.)
                    if let Ok(Some(env)) = deliver.load(&instance).await
                        && env.status.is_terminal()
                    {
                        return Ok(());
                    }
                    let req: RequestPaymentCapture = serde_json::from_value(payload)?;
                    if let Some(payment_id) = providers
                        .capture_payment(&req.booking_id, req.amount_cents, &idem)
                        .await?
                    {
                        // Close the publish-time TOCTOU: the booking may have gone
                        // terminal (a pushed CancelBooking) *while we were charging*
                        // — after the pre-charge check above passed. The payment
                        // step's await was never consumed, so its compensation
                        // (refund_payment) can't run on its own. Re-check, and if
                        // the booking is now terminal, refund the charge we just
                        // made so it isn't orphaned. (A residual micro-window
                        // between this check and `deliver` is a lease race, tracked
                        // as review H-2/H-3; for the deterministic mid-capture
                        // cancel this is the correct, money-safe resolution.)
                        if let Ok(Some(env)) = deliver.load(&instance).await
                            && env.status.is_terminal()
                        {
                            providers
                                .refund_payment(&payment_id, &format!("{idem}/toctou-refund"))
                                .await?;
                            return Ok(());
                        }
                        // Settled and still live → resume the saga. (Unsettled →
                        // no deliver → the payment step's deadline lapses →
                        // compensation.) Errors propagate so the row stays
                        // unpublished and the next drain retries — at-least-once.
                        deliver
                            .deliver(&instance, &PaymentCaptured { payment_id })
                            .await?;
                    }
                    Ok(())
                }
                other => anyhow::bail!("no publisher route for event `{other}`"),
            }
        }
    }
}

// =====================================================================
// The mock third-party backend (axum). Each endpoint is idempotent:
// keyed by the request's `idem` token, it does the effect once and
// returns the stored result on any repeat. `effects` counts distinct
// effects; `requests` counts every call — so a test can prove
// at-least-once delivery with exactly-once effect.
// =====================================================================

#[derive(Default)]
struct ServiceState {
    holds: HashMap<String, String>, // idem -> result id
    requests: u64,
    effects: u64,
    released: std::collections::HashSet<String>, // hold ids released (idempotent)
    refunded: std::collections::HashSet<String>, // payment ids refunded (idempotent)
}

#[derive(Clone, Default)]
pub struct Backend {
    flight: Arc<Mutex<ServiceState>>,
    hotel: Arc<Mutex<ServiceState>>,
    payment: Arc<Mutex<ServiceState>>,
    carrier: Arc<Mutex<ServiceState>>,
    /// 0 = never fail; N = return 503 on every Nth request (across all
    /// endpoints) to simulate a flaky provider the client must retry.
    fail_every: u64,
    seq: Arc<AtomicU64>,
    failures: Arc<AtomicU64>,
}

impl Backend {
    fn flaky(fail_every: u64) -> Self {
        Self {
            fail_every,
            ..Default::default()
        }
    }
    /// Deterministically transient: fails every Nth request, so the retry
    /// (the next request) succeeds — stresses the retry path without making
    /// the test itself flaky.
    fn should_fail(&self) -> bool {
        if self.fail_every == 0 {
            return false;
        }
        let n = self.seq.fetch_add(1, Ordering::Relaxed) + 1;
        if n.is_multiple_of(self.fail_every) {
            self.failures.fetch_add(1, Ordering::Relaxed);
            true
        } else {
            false
        }
    }
    /// How many requests were served a transient 503 (the client retried).
    pub fn failures(&self) -> u64 {
        self.failures.load(Ordering::Relaxed)
    }
}

impl Backend {
    /// (flight, hotel, payment, carrier) distinct-effect counts — each
    /// increments once per unique idempotency key, however many times the
    /// endpoint is called.
    pub fn effects(&self) -> (u64, u64, u64, u64) {
        (
            self.flight.lock().unwrap().effects,
            self.hotel.lock().unwrap().effects,
            self.payment.lock().unwrap().effects,
            self.carrier.lock().unwrap().effects,
        )
    }
    /// (flight, hotel, payment, carrier) total request counts.
    pub fn requests(&self) -> (u64, u64, u64, u64) {
        (
            self.flight.lock().unwrap().requests,
            self.hotel.lock().unwrap().requests,
            self.payment.lock().unwrap().requests,
            self.carrier.lock().unwrap().requests,
        )
    }
    /// The distinct idempotency keys the flight + hotel hold endpoints saw —
    /// for diagnosing a count miss (two keys for one booking ⇒ a double hold).
    pub fn hold_keys(&self) -> (Vec<String>, Vec<String>) {
        let mut f: Vec<_> = self.flight.lock().unwrap().holds.keys().cloned().collect();
        let mut h: Vec<_> = self.hotel.lock().unwrap().holds.keys().cloned().collect();
        f.sort();
        h.sort();
        (f, h)
    }
    /// (flight, hotel) distinct released-hold counts.
    pub fn releases(&self) -> (usize, usize) {
        (
            self.flight.lock().unwrap().released.len(),
            self.hotel.lock().unwrap().released.len(),
        )
    }
    /// Distinct refunded-payment count.
    pub fn refunds(&self) -> usize {
        self.payment.lock().unwrap().refunded.len()
    }
}

fn release(state: &Mutex<ServiceState>, hold_id: &str) {
    state.lock().unwrap().released.insert(hold_id.to_string());
}

fn idempotent(state: &Mutex<ServiceState>, idem: &str, make: impl FnOnce() -> String) -> String {
    let mut st = state.lock().unwrap();
    st.requests += 1;
    if let Some(existing) = st.holds.get(idem) {
        return existing.clone();
    }
    let result = make();
    st.effects += 1;
    st.holds.insert(idem.to_string(), result.clone());
    result
}

pub fn backend_router(backend: Backend) -> Router {
    Router::new()
        .route(
            "/flight/hold",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    let idem = r["idem"].as_str().unwrap_or_default().to_string();
                    let id = idempotent(&b.flight, &idem, || format!("FH-{idem}"));
                    Ok::<_, StatusCode>(Json(serde_json::json!({ "hold_id": id })))
                },
            ),
        )
        .route(
            "/hotel/hold",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    let idem = r["idem"].as_str().unwrap_or_default().to_string();
                    let id = idempotent(&b.hotel, &idem, || format!("HH-{idem}"));
                    Ok::<_, StatusCode>(Json(serde_json::json!({ "hold_id": id })))
                },
            ),
        )
        .route(
            "/payment/capture",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    let idem = r["idem"].as_str().unwrap_or_default().to_string();
                    let booking = r["booking_id"].as_str().unwrap_or_default();
                    // A "doomed-" booking is accepted but never settles — the saga
                    // waits, times out, and compensates the holds.
                    if booking.starts_with("doomed-") {
                        return Ok(Json(
                            serde_json::json!({ "captured": false, "payment_id": null }),
                        ));
                    }
                    // A "slowcap-" booking takes a beat to settle — the charge is
                    // recorded only after the delay. This opens a deterministic
                    // window: a caller can fire a cancellation *while the capture
                    // is in flight* (after the publisher's pre-charge terminal
                    // check, before the money moves) to exercise the publish-time
                    // TOCTOU and the orphan-charge refund that closes it.
                    if booking.starts_with("slowcap-") {
                        tokio::time::sleep(std::time::Duration::from_millis(300)).await;
                    }
                    let id = idempotent(&b.payment, &idem, || format!("PAY-{idem}"));
                    Ok::<_, StatusCode>(Json(
                        serde_json::json!({ "captured": true, "payment_id": id }),
                    ))
                },
            ),
        )
        .route(
            "/carrier/ticket",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    // A "noticket-" booking is permanently rejected (422, not a
                    // retryable 503) — the inline confirm effect fails, and the
                    // cascade refunds the (already-captured) payment + releases.
                    let booking = r["booking_id"].as_str().unwrap_or_default();
                    if booking.starts_with("noticket-") {
                        return Err(StatusCode::UNPROCESSABLE_ENTITY);
                    }
                    let idem = r["idem"].as_str().unwrap_or_default().to_string();
                    let id = idempotent(&b.carrier, &idem, || format!("TKT-{idem}"));
                    Ok::<_, StatusCode>(Json(serde_json::json!({ "ticket": id })))
                },
            ),
        )
        .route(
            "/flight/release",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    let hold = r["hold_id"].as_str().unwrap_or_default();
                    // A "stuckhold-" booking's flight hold can't be released (422,
                    // permanent) — the compensation fails, leaving a durable lock
                    // the saga surfaces as FailedTerminalNeedsIntervention.
                    if hold.contains("stuckhold") {
                        return Err(StatusCode::UNPROCESSABLE_ENTITY);
                    }
                    release(&b.flight, hold);
                    Ok::<_, StatusCode>(Json(serde_json::json!({ "released": true })))
                },
            ),
        )
        .route(
            "/hotel/release",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    release(&b.hotel, r["hold_id"].as_str().unwrap_or_default());
                    Ok::<_, StatusCode>(Json(serde_json::json!({ "released": true })))
                },
            ),
        )
        .route(
            "/payment/refund",
            post(
                |AxState(b): AxState<Backend>, Json(r): Json<serde_json::Value>| async move {
                    if b.should_fail() {
                        return Err(StatusCode::SERVICE_UNAVAILABLE);
                    }
                    let pid = r["payment_id"].as_str().unwrap_or_default().to_string();
                    b.payment.lock().unwrap().refunded.insert(pid);
                    Ok::<_, StatusCode>(Json(serde_json::json!({ "refunded": true })))
                },
            ),
        )
        .with_state(backend)
}

/// Bind the mock backend on an ephemeral port; returns its base URL, a
/// handle to inspect call counts, and the server task.
pub async fn spawn_backend() -> (String, Backend, tokio::task::JoinHandle<()>) {
    spawn_backend_with(0).await
}

/// Like [`spawn_backend`] but the providers fail every `fail_every`th request
/// with 503 (0 = never) — to stress the client/outbox retry paths.
pub async fn spawn_backend_with(fail_every: u64) -> (String, Backend, tokio::task::JoinHandle<()>) {
    let backend = Backend::flaky(fail_every);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let base = format!("http://{}", listener.local_addr().unwrap());
    let app = backend_router(backend.clone());
    let handle = tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    });
    (base, backend, handle)
}

// =====================================================================
// Harness — the running system: the mock backend and a SagaWorker driving
// outbox-drain (via the HttpPublisher, which delivers results back),
// timeouts, and recovery. `trigger` is the application's entry point. For
// crash tests, `spawn_worker` lets you abort + restart the worker directly.
// =====================================================================

pub struct Harness<S: Saga<State = TravelBooking> = TravelBookingSaga> {
    pub trigger: PgSagaRunner<S>,
    pub backend: Backend,
    cancel: CancellationToken,
    tasks: Vec<tokio::task::JoinHandle<()>>,
}

fn topology<S: Saga<State = TravelBooking>>(providers: &Providers) -> SagaBuilder<S> {
    build_travel_topology(SagaBuilder::<S>::for_saga()).dep(providers.clone())
}

/// A runner bound to the store + providers (its `trigger`/`deliver`/inline
/// effects use the providers). Cheap — runners share the store.
pub fn new_runner(store: EventStore, providers: &Providers) -> PgSagaRunner<TravelBookingSaga> {
    new_runner_sharded::<TravelBookingSaga>(store, providers, 1)
}

/// Like [`new_runner`] but generic over the saga marker `S` and partitioning
/// the saga across `num_shards` shards. The marker lets tests run the same
/// topology under a distinct saga *name* (recovery is saga-wide, so isolation
/// matters under concurrency). The trigger that *stamps* a row and every worker
/// that *scans* it must agree on `num_shards` (the runner hashes `instance_id`
/// into `0..num_shards`), so a sharded deployment builds all of its runners —
/// trigger, deliver, worker — through this with the same `S` and count.
pub fn new_runner_sharded<S: Saga<State = TravelBooking>>(
    store: EventStore,
    providers: &Providers,
    num_shards: u32,
) -> PgSagaRunner<S> {
    // A short instance-lease TTL: with the phase-7 fence making steals safe,
    // a crashed worker's lease lapses fast and recovery reclaims the instance
    // in a few seconds instead of tens. Heartbeat (ttl/3 ≈ 1s) keeps a slow
    // provider call from losing the lease, and a TTL this size keeps lease-row
    // churn (every heartbeat is an UPDATE) modest under the stress suite.
    // Production would use a larger TTL still.
    PgSagaRunner::new(topology::<S>(providers), store)
        .with_num_shards(num_shards)
        .with_lease_ttl(std::time::Duration::from_secs(3))
}

/// Spawn an autonomous worker (outbox drain + publisher, timeouts, recovery)
/// driven by `cancel`. Returns its task handle — abort it to simulate a hard
/// crash, or cancel the token for graceful shutdown.
pub fn spawn_worker(
    store: EventStore,
    providers: &Providers,
    cancel: CancellationToken,
) -> tokio::task::JoinHandle<()> {
    spawn_worker_sharded::<TravelBookingSaga>(store, providers, cancel, 1)
}

/// Like [`spawn_worker`] but generic over the saga marker `S` and owning a
/// converging subset of `num_shards` shards (via shard leases) — the horizontal
/// scale-out path. A short shard-lease TTL keeps fleet convergence fast for
/// tests; production uses a longer one. At `num_shards == 1` this is byte-for-
/// byte the single-worker path (the worker takes the whole-table scan and never
/// touches a shard lease), so [`spawn_worker`] delegates here.
pub fn spawn_worker_sharded<S: Saga<State = TravelBooking>>(
    store: EventStore,
    providers: &Providers,
    cancel: CancellationToken,
    num_shards: u32,
) -> tokio::task::JoinHandle<()> {
    let publisher = HttpPublisher::<S> {
        providers: providers.clone(),
        deliver: Arc::new(new_runner_sharded::<S>(
            store.clone(),
            providers,
            num_shards,
        )),
    };
    let worker = SagaWorker::new(
        new_runner_sharded::<S>(store, providers, num_shards).with_publisher(publisher),
    )
    .poll_interval(std::time::Duration::from_millis(50))
    .drain_interval(std::time::Duration::from_millis(50))
    // Phase 7: recovery is lease-expiry-driven — a crashed worker's
    // lease lapses (the short TTL set in `new_runner_sharded`) and the
    // next scan steals (fenced CAS) and re-drives. No `recovery_after`
    // staleness clock; the fence makes stealing a still-running
    // instance safe, so a slow-but-alive holder is never raced into a
    // duplicate effect. Scan often so takeover is prompt.
    .recovery_interval(std::time::Duration::from_millis(250))
    .shard_lease_ttl(std::time::Duration::from_secs(2));
    tokio::spawn(async move {
        let _ = worker.run(cancel).await;
    })
}

impl<S: Saga<State = TravelBooking>> Harness<S> {
    pub async fn start(store: EventStore) -> Self {
        Self::start_with(store, 0).await
    }

    /// Start with a flaky backend (providers fail every `fail_every`th request).
    pub async fn start_with(store: EventStore, fail_every: u64) -> Self {
        let (base, backend, server) = spawn_backend_with(fail_every).await;
        let providers = Providers::new(&base);
        let cancel = CancellationToken::new();
        let worker_task = spawn_worker_sharded::<S>(store.clone(), &providers, cancel.clone(), 1);
        let trigger = new_runner_sharded::<S>(store, &providers, 1);
        Harness {
            trigger,
            backend,
            cancel,
            tasks: vec![worker_task, server],
        }
    }

    /// Poll until `booking_id` reaches a terminal state or `budget` elapses.
    pub async fn await_terminal(
        &self,
        booking_id: &str,
        budget: std::time::Duration,
    ) -> Option<SagaInstanceStatus> {
        let start = std::time::Instant::now();
        loop {
            if let Ok(Some(env)) = self.trigger.load(booking_id).await
                && env.status.is_terminal()
            {
                return Some(env.status);
            }
            if start.elapsed() > budget {
                return None;
            }
            tokio::time::sleep(std::time::Duration::from_millis(25)).await;
        }
    }

    pub async fn shutdown(self) {
        self.cancel.cancel();
        // Abort AND join: a worker's recovery scan is saga-WIDE, so a worker
        // that outlives this harness would re-drive another (concurrently or
        // subsequently running) instance of the same saga — under this dropped
        // harness's now-dead backend. Joining guarantees it's stopped before we
        // return. (Tests serialize on a shared saga name; this is the
        // saga-test-isolation rule for live workers.)
        for t in self.tasks {
            t.abort();
            let _ = t.await;
        }
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    let pool = sqlx::PgPool::connect(&url).await?;
    let store = EventStore::new(pool);
    store.migrate().await?;
    mire_sagas::migrate(&store).await?;

    let h = Harness::<TravelBookingSaga>::start(store).await;
    let booking_id = format!("bk-{}", uuid::Uuid::new_v4().simple());
    h.trigger
        .trigger::<Booking>(&BookingRequested {
            booking_id: booking_id.clone(),
            flight_id: "LH-100".into(),
            hotel_id: "HILTON-OSL".into(),
            amount_cents: 150_000,
        })
        .await?;

    // flight+hotel ran inline during trigger; payment is async — the worker
    // publishes it and the deliver loop resumes the saga.
    let status = h
        .await_terminal(&booking_id, std::time::Duration::from_secs(10))
        .await;
    println!("booking {booking_id} → {status:?}");
    println!(
        "provider effects (flight,hotel,payment,carrier): {:?}",
        h.backend.effects()
    );
    assert_eq!(status, Some(SagaInstanceStatus::SucceededTerminal));
    h.shutdown().await;
    Ok(())
}
