//! P4.3 — a sharded `SagaWorker` fleet splits the timeout scan across shards
//! while still driving every instance to terminal.
//!
//! Two workers over `num_shards = 8` must (a) converge to a disjoint shard
//! split covering all shards (the distribution mechanism actually engaged,
//! not one worker doing everything) and (b) collectively time out every
//! suspended instance — correctness is unchanged by sharding because the
//! per-instance lease, not shard ownership, is the fence.
//!
//! Gated on `DATABASE_URL`; skips silently when unset.

#![allow(dead_code)]

use std::collections::BTreeSet;
use std::time::Duration;

use mire::{Aggregate, EventData};
use mire_sagas::{
    Emit, PgSagaRunner, Saga, SagaBuilder, SagaInstanceStatus, SagaWorker, Start, State, task,
};
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;

mod common;

#[derive(Default, Clone)]
struct Trigger;

#[derive(Serialize, Deserialize, Clone, EventData)]
#[mire(event_type = "sws.requested")]
struct Requested {
    id: String,
}

impl Aggregate for Trigger {
    type Event = Requested;
    fn stream_category() -> &'static str {
        "sws-trigger"
    }
    fn apply(&mut self, _: &Self::Event) {}
}

#[derive(Default, Clone, Serialize, Deserialize)]
struct St {
    id: String,
}

// One step that emits to an external worker with a short deadline; nobody
// answers, so it must time out and fail terminal — driven only by a worker.
#[derive(Debug, Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "sws.pay-req")]
struct PayReq {
    id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "sws.paid")]
struct Paid;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct PayDone;

fn pay(State(s): State<St>) -> Emit {
    Emit::event(PayReq { id: s.id.clone() }).timeout(Duration::from_millis(150))
}

struct ShardedSaga;
impl Saga for ShardedSaga {
    type State = St;
    fn name() -> &'static str {
        "sharded-worker-test"
    }
    fn build(b: SagaBuilder<Self>) -> SagaBuilder<Self> {
        b.start_on::<Trigger, _>(|e: &Requested| {
            Some(Start {
                instance_id: e.id.clone(),
                state: St { id: e.id.clone() },
            })
        })
        .step(
            "pay",
            task(pay).on_response(|_: Paid| PayDone).terminal_success(),
        )
    }
}

const N_SHARDS: u32 = 8;

fn runner(store: &mire::EventStore) -> PgSagaRunner<ShardedSaga> {
    PgSagaRunner::new(
        ShardedSaga::build(SagaBuilder::<ShardedSaga>::for_saga()),
        store.clone(),
    )
    .with_num_shards(N_SHARDS)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn sharded_fleet_splits_scan_and_drives_all_to_terminal() {
    let Some(store) = common::store().await else {
        return;
    };

    // Trigger a spread of instances; each suspends awaiting a response that
    // never comes (150ms deadline).
    const M: usize = 24;
    let trig = runner(&store);
    let mut ids = Vec::new();
    for _ in 0..M {
        let id = common::unique("sws");
        trig.trigger::<Trigger>(&Requested { id: id.clone() })
            .await
            .unwrap()
            .unwrap();
        ids.push(id);
    }

    // Two workers, distinct identities (default per-runner UUID), short shard
    // lease TTL so they converge fast.
    let cancel = CancellationToken::new();
    let mut handles = Vec::new();
    for _ in 0..2 {
        let w = SagaWorker::new(runner(&store))
            .poll_interval(Duration::from_millis(50))
            .shard_lease_ttl(Duration::from_secs(2));
        let c = cancel.clone();
        handles.push(tokio::spawn(async move { w.run(c).await }));
    }

    // Wait for every instance to reach terminal, driven only by the workers.
    let mut all_terminal = false;
    for _ in 0..120 {
        let mut done = true;
        for id in &ids {
            if !trig.is_terminal(id).await.unwrap() {
                done = false;
                break;
            }
        }
        if done {
            all_terminal = true;
            break;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }

    // Sample shard ownership while the fleet is still running: both workers
    // must hold a non-empty, disjoint share and together cover all shards —
    // proof the scan was actually distributed, not done by one worker.
    let leases: Vec<(String, i32)> = sqlx::query_as(
        "SELECT worker_id, shard_id FROM es_saga_shard_leases \
         WHERE saga_name = $1 AND leased_until > now() ORDER BY shard_id",
    )
    .bind(ShardedSaga::name())
    .fetch_all(store.pool())
    .await
    .unwrap();

    cancel.cancel();
    for h in handles {
        h.await.unwrap().unwrap();
    }

    assert!(
        all_terminal,
        "sharded worker fleet should drive every instance terminal"
    );

    let workers: BTreeSet<&str> = leases.iter().map(|(w, _)| w.as_str()).collect();
    let shards: BTreeSet<i32> = leases.iter().map(|(_, s)| *s).collect();
    assert_eq!(
        workers.len(),
        2,
        "both workers should hold shards (distribution engaged), got {workers:?}"
    );
    assert_eq!(
        shards,
        (0..N_SHARDS as i32).collect(),
        "the fleet must own every shard between them, got {shards:?}"
    );
    // Disjoint: no shard appears for two workers (PK guarantees one row per
    // shard, so == covers it, but assert the count to be explicit).
    assert_eq!(
        leases.len(),
        N_SHARDS as usize,
        "each shard owned exactly once"
    );

    // Every instance ended terminal-failed (its deadline expired uncomp'd).
    for id in &ids {
        assert_eq!(
            trig.status(id).await.unwrap(),
            Some(SagaInstanceStatus::FailedTerminal),
            "instance {id} should have timed out to FailedTerminal"
        );
    }
}
