//! Shard-ownership leases for the worker fleet (plan §P4.2).
//!
//! A [`SagaWorker`](crate::SagaWorker) owns a *subset* of the saga's shards
//! and only scans those — timeouts, signal drain, recovery (P4.3). Ownership
//! is recorded in `es_saga_shard_leases` and maintained by [`rebalance`],
//! called once per worker tick.
//!
//! **This is work distribution, not a correctness fence.** Three distinct
//! mechanisms, do not conflate them (principle 6):
//! - the **instance lease** (`es_saga_leases`) is the mutual-exclusion unit
//!   for effects — at-most-one active processor per instance;
//! - **OCC** on the envelope stream fences the log — a stale worker's commit
//!   is discarded;
//! - **shard leases** (here) only decide *which worker scans which
//!   partition*. A shard transiently owned by two workers is safe (the
//!   instance lease still serialises any instance they both reach); a shard
//!   owned by nobody for a tick just delays its scans. So this layer is
//!   deliberately boring: best-effort, no fence token, eventual convergence.
//!
//! ## Convergence policy
//!
//! Each worker targets an even share without any central coordinator:
//! `target = ceil(num_shards / live_workers)`. The denominator is **fleet
//! membership**, counted from `es_saga_worker_heartbeats` (workers that
//! heartbeated within the liveness window) — *not* the number of shards
//! currently held. That distinction is load-bearing: a brand-new worker
//! holds zero shards, so if incumbents counted holders they'd see themselves
//! as the sole owner, keep `target = num_shards`, and never shed — the
//! newcomer (which can't steal a live shard lease) could never get in. By
//! heartbeating presence regardless of holdings, the newcomer lifts every
//! incumbent's denominator immediately, so they shed within a tick.
//!
//! Per tick a worker (1) heartbeats its presence and renews the shards it
//! holds, (2) if it holds fewer than target, claims free/expired shards up to
//! target, (3) if it holds more than target, releases the excess for peers to
//! pick up. A dying worker ages out of the count and its leases lapse;
//! survivors' denominator drops and they reabsorb the shards. Convergence
//! need not be instant or exact — overlap and brief gaps are both safe.

use std::time::Duration;

use mire::EventStoreError;
use sqlx::PgPool;
use tracing::{debug, trace};

/// Renew every shard this worker currently holds, extending each lease by
/// `ttl`. Returns the shard ids still owned afterward. A shard that was
/// stolen while we were away (its `worker_id` is now someone else's) is not
/// matched and silently drops out of our owned set — exactly the takeover we
/// want. `WHERE worker_id = $me` ignores `leased_until`: if a shard is still
/// ours but briefly lapsed and nobody took it, we reclaim it here.
pub async fn renew_owned(
    pool: &PgPool,
    saga_name: &str,
    worker_id: &str,
    ttl: Duration,
) -> Result<Vec<i32>, EventStoreError> {
    let ttl_secs = ttl.as_secs_f64().to_string();
    let owned: Vec<i32> = sqlx::query_scalar(
        "UPDATE es_saga_shard_leases
            SET leased_until = now() + ($1 || ' seconds')::interval,
                heartbeat_at = now()
          WHERE saga_name = $2 AND worker_id = $3
        RETURNING shard_id",
    )
    .bind(ttl_secs)
    .bind(saga_name)
    .bind(worker_id)
    .fetch_all(pool)
    .await?;
    Ok(owned)
}

/// Record (or refresh) this worker's presence in the fleet membership
/// registry. Called once per tick before assignment so the worker is counted
/// in the target-share denominator even when it holds no shards yet.
pub async fn heartbeat_worker(
    pool: &PgPool,
    saga_name: &str,
    worker_id: &str,
) -> Result<(), EventStoreError> {
    sqlx::query(
        "INSERT INTO es_saga_worker_heartbeats (saga_name, worker_id, heartbeat_at)
         VALUES ($1, $2, now())
         ON CONFLICT (saga_name, worker_id) DO UPDATE SET heartbeat_at = now()",
    )
    .bind(saga_name)
    .bind(worker_id)
    .execute(pool)
    .await?;
    Ok(())
}

/// Number of workers that heartbeated within `window` — the fleet-membership
/// denominator for the target share. See the module-level convergence note
/// for why this counts presence, not shard holdings.
pub async fn live_worker_count(
    pool: &PgPool,
    saga_name: &str,
    window: Duration,
) -> Result<i64, EventStoreError> {
    let window_secs = window.as_secs_f64().to_string();
    let n: i64 = sqlx::query_scalar(
        "SELECT count(*) FROM es_saga_worker_heartbeats
          WHERE saga_name = $1
            AND heartbeat_at > now() - ($2 || ' seconds')::interval",
    )
    .bind(saga_name)
    .bind(window_secs)
    .fetch_one(pool)
    .await?;
    Ok(n)
}

/// Remove this worker's membership row on graceful shutdown so the fleet
/// count drops immediately rather than waiting out the liveness window.
pub async fn deregister_worker(
    pool: &PgPool,
    saga_name: &str,
    worker_id: &str,
) -> Result<(), EventStoreError> {
    sqlx::query("DELETE FROM es_saga_worker_heartbeats WHERE saga_name = $1 AND worker_id = $2")
        .bind(saga_name)
        .bind(worker_id)
        .execute(pool)
        .await?;
    Ok(())
}

/// Shards that are unclaimed or whose lease has expired — acquisition
/// candidates — up to `limit`, lowest shard id first (deterministic so two
/// workers racing for the same free shards contend on the same rows and the
/// PK fence resolves it). `generate_series` covers shard ids that have no row
/// yet (never claimed).
async fn acquirable_shards(
    pool: &PgPool,
    saga_name: &str,
    num_shards: u32,
    limit: i64,
) -> Result<Vec<i32>, EventStoreError> {
    if limit <= 0 {
        return Ok(Vec::new());
    }
    let shards: Vec<i32> = sqlx::query_scalar(
        "SELECT s::int AS shard_id
           FROM generate_series(0, $1::int - 1) s
           LEFT JOIN es_saga_shard_leases l
             ON l.saga_name = $2 AND l.shard_id = s
          WHERE l.shard_id IS NULL OR l.leased_until < now()
          ORDER BY s
          LIMIT $3",
    )
    .bind(num_shards as i32)
    .bind(saga_name)
    .bind(limit)
    .fetch_all(pool)
    .await?;
    Ok(shards)
}

/// Try to claim one shard, valid for `ttl`. Succeeds (`Ok(true)`) when the
/// shard is unclaimed, its lease has expired, or it is already ours;
/// `Ok(false)` when a different live worker holds it (lost the race). Same
/// single-statement steal-if-expired fence as the instance lease.
pub async fn try_acquire(
    pool: &PgPool,
    saga_name: &str,
    shard_id: i32,
    worker_id: &str,
    ttl: Duration,
) -> Result<bool, EventStoreError> {
    let ttl_secs = ttl.as_secs_f64().to_string();
    let row: Option<i32> = sqlx::query_scalar(
        "INSERT INTO es_saga_shard_leases
            (saga_name, shard_id, worker_id, leased_until, heartbeat_at)
         VALUES ($1, $2, $3, now() + ($4 || ' seconds')::interval, now())
         ON CONFLICT (saga_name, shard_id) DO UPDATE
            SET worker_id    = EXCLUDED.worker_id,
                leased_until = EXCLUDED.leased_until,
                heartbeat_at = now()
          WHERE es_saga_shard_leases.leased_until < now()
             OR es_saga_shard_leases.worker_id = EXCLUDED.worker_id
         RETURNING shard_id",
    )
    .bind(saga_name)
    .bind(shard_id)
    .bind(worker_id)
    .bind(ttl_secs)
    .fetch_optional(pool)
    .await?;
    let acquired = row.is_some();
    if acquired {
        trace!(saga_name, shard_id, worker_id, "shard lease acquired");
    } else {
        trace!(
            saga_name,
            shard_id, worker_id, "shard lease contended (live holder)"
        );
    }
    Ok(acquired)
}

/// Release one shard we hold so a peer can take it immediately (voluntary
/// rebalance / graceful shutdown). Best-effort: only frees the row if it is
/// still ours.
pub async fn release(
    pool: &PgPool,
    saga_name: &str,
    shard_id: i32,
    worker_id: &str,
) -> Result<(), EventStoreError> {
    sqlx::query(
        "UPDATE es_saga_shard_leases
            SET leased_until = '-infinity'
          WHERE saga_name = $1 AND shard_id = $2 AND worker_id = $3",
    )
    .bind(saga_name)
    .bind(shard_id)
    .bind(worker_id)
    .execute(pool)
    .await?;
    trace!(saga_name, shard_id, worker_id, "shard lease released");
    Ok(())
}

/// Release every shard this worker holds (graceful shutdown). Survivors pick
/// them up on their next tick rather than waiting out the TTL.
pub async fn release_all(
    pool: &PgPool,
    saga_name: &str,
    worker_id: &str,
) -> Result<(), EventStoreError> {
    sqlx::query(
        "UPDATE es_saga_shard_leases
            SET leased_until = '-infinity'
          WHERE saga_name = $1 AND worker_id = $2",
    )
    .bind(saga_name)
    .bind(worker_id)
    .execute(pool)
    .await?;
    Ok(())
}

/// One convergence tick. Heartbeats presence, renews held shards, then moves
/// toward an even share (`ceil(num_shards / live_workers)`): claims
/// free/expired shards if under target, releases the excess if over. Returns
/// the shards owned **after** this tick — the set the worker filters its
/// scans by until the next tick.
///
/// `ttl` bounds how long an owned shard lease stays valid without renewal;
/// `liveness_window` bounds how long a worker counts toward the fleet size
/// without heartbeating (use a small multiple of the tick interval, e.g.
/// `ttl`). **Both are fleet-wide constants** — a worker is judged "alive" per
/// whatever window the *querying* worker uses, so every worker must pass the
/// same `liveness_window` (and `ttl`), exactly as they must agree on
/// `num_shards`. A mismatch makes the fleet disagree on who is alive and the
/// share lopsided. `release_excess` lets the caller apply hysteresis: pass
/// `false`
/// for the first few ticks after a membership change so a flapping peer
/// doesn't cause shards to ping-pong; pass `true` once stable. Acquiring up
/// to target is always safe to do every tick (it only ever fills gaps).
pub async fn rebalance(
    pool: &PgPool,
    saga_name: &str,
    worker_id: &str,
    num_shards: u32,
    ttl: Duration,
    liveness_window: Duration,
    release_excess: bool,
) -> Result<Vec<i32>, EventStoreError> {
    // 1. Register presence (counted in the denominator even at 0 shards) and
    //    renew what we hold.
    heartbeat_worker(pool, saga_name, worker_id).await?;
    let mut owned = renew_owned(pool, saga_name, worker_id, ttl).await?;

    // 2. Target share = ceil(num_shards / live workers). Membership, not
    //    holdings — see the module note.
    let live = live_worker_count(pool, saga_name, liveness_window)
        .await?
        .max(1) as u32;
    let target = num_shards.div_ceil(live);

    let owned_before = owned.len();
    debug!(
        saga_name,
        worker_id,
        num_shards,
        live_workers = live,
        target_shards = target,
        owned_before,
        "shard rebalance tick",
    );

    // 3. Under target → claim free/expired shards up to target.
    if (owned.len() as u32) < target {
        let need = target - owned.len() as u32;
        let candidates = acquirable_shards(pool, saga_name, num_shards, need as i64).await?;
        for shard in candidates {
            if (owned.len() as u32) >= target {
                break;
            }
            if try_acquire(pool, saga_name, shard, worker_id, ttl).await? {
                owned.push(shard);
            }
        }
    } else if release_excess && (owned.len() as u32) > target {
        // 4. Over target → shed the excess (highest shard ids first) so peers
        //    can pick them up. Drop the released shards from our owned set.
        owned.sort_unstable();
        while (owned.len() as u32) > target {
            if let Some(shard) = owned.pop() {
                release(pool, saga_name, shard, worker_id).await?;
            }
        }
    }

    owned.sort_unstable();
    if owned.len() != owned_before {
        debug!(
            saga_name,
            worker_id,
            owned_after = owned.len(),
            "shard ownership changed after rebalance",
        );
    }
    Ok(owned)
}
