//! P4.2 — shard-ownership lease convergence.
//!
//! Drives `shard_lease::rebalance` directly (no Saga type needed — the
//! functions take `saga_name` as a string) against local Postgres, asserting
//! the fleet converges to an even, disjoint split and recovers from a dead
//! worker. Each test uses a unique `saga_name` so parallel binaries/reruns
//! never collide.
//!
//! Gated on `DATABASE_URL`; skips silently when unset.

#![allow(dead_code)]

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

use mire_sagas::shard_lease;

mod common;

// Generous so nothing lapses mid-test except where a test forces expiry.
const TTL: Duration = Duration::from_secs(60);
const WINDOW: Duration = Duration::from_secs(60);

fn saga_name() -> String {
    // shard_lease takes &str saga names; a unique one isolates each test.
    common::unique("shardlease")
}

async fn rebalance(pool: &sqlx::PgPool, saga: &str, worker: &str, n: u32) -> Vec<i32> {
    shard_lease::rebalance(pool, saga, worker, n, TTL, WINDOW, true)
        .await
        .unwrap()
}

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

    let owned = rebalance(pool, &saga, "w1", 8).await;
    assert_eq!(
        owned,
        (0..8).collect::<Vec<_>>(),
        "lone worker owns all shards"
    );

    // Idempotent: a second tick keeps the same set.
    let again = rebalance(pool, &saga, "w1", 8).await;
    assert_eq!(again, owned);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn two_workers_split_evenly() {
    let Some(store) = common::store().await else {
        return;
    };
    let pool = store.pool();
    let saga = saga_name();
    const N: u32 = 8;

    // Interleave ticks until stable.
    let mut a = Vec::new();
    let mut b = Vec::new();
    for _ in 0..6 {
        a = rebalance(pool, &saga, "wa", N).await;
        b = rebalance(pool, &saga, "wb", N).await;
    }

    let sa: BTreeSet<i32> = a.iter().copied().collect();
    let sb: BTreeSet<i32> = b.iter().copied().collect();

    // Disjoint and complete: every shard owned by exactly one worker.
    assert!(
        sa.is_disjoint(&sb),
        "shards must not be double-owned: {sa:?} {sb:?}"
    );
    let union: BTreeSet<i32> = sa.union(&sb).copied().collect();
    assert_eq!(
        union,
        (0..N as i32).collect(),
        "every shard owned: {union:?}"
    );

    // Even within one: ceil(8/2) = 4 each.
    assert_eq!(sa.len(), 4, "wa share {sa:?}");
    assert_eq!(sb.len(), 4, "wb share {sb:?}");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn three_workers_share_within_one() {
    let Some(store) = common::store().await else {
        return;
    };
    let pool = store.pool();
    let saga = saga_name();
    const N: u32 = 8;

    let (mut a, mut b, mut c) = (Vec::new(), Vec::new(), Vec::new());
    for _ in 0..8 {
        a = rebalance(pool, &saga, "wa", N).await;
        b = rebalance(pool, &saga, "wb", N).await;
        c = rebalance(pool, &saga, "wc", N).await;
    }
    let sets = [&a, &b, &c].map(|v| v.iter().copied().collect::<BTreeSet<_>>());

    // Complete + disjoint.
    let union: BTreeSet<i32> = sets.iter().flatten().copied().collect();
    assert_eq!(
        union,
        (0..N as i32).collect(),
        "every shard owned: {union:?}"
    );
    let total: usize = sets.iter().map(|s| s.len()).sum();
    assert_eq!(total, N as usize, "no shard double-owned (sum {total})");

    // ceil(8/3) = 3 → each worker owns at most 3, none starves.
    for s in &sets {
        assert!((1..=3).contains(&s.len()), "share within bounds: {s:?}");
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn dead_worker_shards_reabsorbed_gracefully() {
    let Some(store) = common::store().await else {
        return;
    };
    let pool = store.pool();
    let saga = saga_name();
    const N: u32 = 8;

    // Converge two workers.
    for _ in 0..6 {
        rebalance(pool, &saga, "wa", N).await;
        rebalance(pool, &saga, "wb", N).await;
    }

    // wb shuts down gracefully: release its shards + deregister membership.
    shard_lease::release_all(pool, &saga, "wb").await.unwrap();
    shard_lease::deregister_worker(pool, &saga, "wb")
        .await
        .unwrap();

    // wa reabsorbs everything: live=1 → target=N, and wb's shards are free.
    let mut a = Vec::new();
    for _ in 0..3 {
        a = rebalance(pool, &saga, "wa", N).await;
    }
    assert_eq!(
        a,
        (0..N as i32).collect::<Vec<_>>(),
        "survivor reabsorbs all"
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn expired_worker_is_taken_over() {
    let Some(store) = common::store().await else {
        return;
    };
    let pool = store.pool();
    let saga = saga_name();
    const N: u32 = 4;

    // The ttl + liveness window are fleet-wide constants — a worker is "alive"
    // per whatever window the *querier* uses, so all workers must agree. Use a
    // short shared value here so a crashed worker (one that just stops ticking,
    // no graceful release/deregister) ages out fast. The convergence loop runs
    // without sleeps so no lease lapses mid-loop.
    let short = Duration::from_millis(400);
    for _ in 0..4 {
        shard_lease::rebalance(pool, &saga, "wa", N, short, short, true)
            .await
            .unwrap();
        shard_lease::rebalance(pool, &saga, "wb", N, short, short, true)
            .await
            .unwrap();
    }

    // wb "crashes": we simply stop ticking it. Sleep past the shared window so
    // both its membership heartbeat and its shard leases lapse.
    tokio::time::sleep(Duration::from_millis(700)).await;

    // wa rebalances: wb no longer counts (membership expired) and its shard
    // leases are expired → wa reclaims its own (renew-by-worker_id) and takes
    // over wb's expired shards.
    let mut a = Vec::new();
    for _ in 0..3 {
        a = shard_lease::rebalance(pool, &saga, "wa", N, short, short, true)
            .await
            .unwrap();
    }
    assert_eq!(
        a,
        (0..N as i32).collect::<Vec<_>>(),
        "survivor takes over the crashed worker's shards"
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn hysteresis_holds_excess_when_release_disabled() {
    let Some(store) = common::store().await else {
        return;
    };
    let pool = store.pool();
    let saga = saga_name();
    const N: u32 = 8;

    // wa alone grabs all 8.
    shard_lease::rebalance(pool, &saga, "wa", N, TTL, WINDOW, true)
        .await
        .unwrap();

    // wb appears; wa ticks with release_excess=false → keeps its 8 despite
    // target dropping to 4 (hysteresis window). wb gets nothing yet.
    let a = shard_lease::rebalance(pool, &saga, "wa", N, TTL, WINDOW, false)
        .await
        .unwrap();
    let b = shard_lease::rebalance(pool, &saga, "wb", N, TTL, WINDOW, false)
        .await
        .unwrap();
    assert_eq!(a.len(), 8, "release disabled → wa keeps excess");
    assert_eq!(b.len(), 0, "wb can't take live shards while wa holds them");

    // Once wa stops holding back (release_excess=true), it sheds to target.
    let a = shard_lease::rebalance(pool, &saga, "wa", N, TTL, WINDOW, true)
        .await
        .unwrap();
    assert_eq!(a.len(), 4, "release enabled → wa sheds to target");
}
