//! P4.6 — sharded worker-fleet scale validation (the saga-effector throughput
//! metric: instances/sec to terminal, 1 vs N workers).
//!
//! Unlike `stress_pg_chaos::scaling_node_fanout` (host-driven `deliver` with a
//! *manually* disjoint partition), this exercises the thing P4.1–P4.3 built:
//! autonomous [`SagaWorker`]s that self-partition `num_shards` via shard
//! leases and each scan only the timeouts they own. M instances are triggered
//! suspended on a short deadline; once it lapses, a fleet of K workers drains
//! them to terminal with no host involvement. Strong scaling — same M, more
//! workers should finish faster — answers "how many concurrent saga effectors,
//! 1 vs N workers/shards" and proves the scan distributes (H4).
//!
//! Run on a FRESH DB (principle 10):
//!   docker compose down -v && docker compose up -d --wait postgres
//!   MIRE_SHARDSCALE=1 cargo test --release -p mire-sagas --test scale_sharded \
//!     -- --nocapture --test-threads=1
//! Gated on MIRE_SHARDSCALE + DATABASE_URL.

#![allow(dead_code)]

use std::sync::Arc;
use std::time::{Duration, Instant};

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

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

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

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

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

#[derive(Debug, Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "scs.work-req")]
struct WorkReq {
    id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, EventData)]
#[mire(event_type = "scs.work-done")]
struct WorkDone;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct StepDone;

async fn work(State(s): State<St>) -> Emit {
    // Short deadline: nobody answers, so the worker fleet times it out.
    Emit::event(WorkReq { id: s.id.clone() }).timeout(Duration::from_millis(50))
}

struct ScaleSaga;
impl Saga for ScaleSaga {
    type State = St;
    fn name() -> &'static str {
        "scale-sharded-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(
            "work",
            task(work)
                .on_response(|_: WorkDone| StepDone)
                .terminal_success(),
        )
    }
}

const N_SHARDS: u32 = 16;
const M: usize = 1500; // instances to drain per fleet size

async fn runner(url: &str, pool_size: u32) -> PgSagaRunner<ScaleSaga> {
    let pool = PgPoolOptions::new()
        .max_connections(pool_size)
        .connect(url)
        .await
        .expect("connect");
    PgSagaRunner::new(
        ScaleSaga::build(SagaBuilder::<ScaleSaga>::for_saga()),
        mire::EventStore::new(pool),
    )
    .with_num_shards(N_SHARDS)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn scaling_sharded_worker_drain() {
    let Some(url) = std::env::var("DATABASE_URL").ok() else {
        eprintln!("skipping: DATABASE_URL not set");
        return;
    };
    if std::env::var("MIRE_SHARDSCALE").is_err() {
        eprintln!("skipping scaling_sharded_worker_drain; set MIRE_SHARDSCALE=1 to run");
        return;
    }

    // Migrate once up front on a dedicated connection (principle 7/10).
    {
        let p = PgPoolOptions::new()
            .max_connections(1)
            .connect(&url)
            .await
            .unwrap();
        let s = mire::EventStore::new(p);
        s.migrate().await.unwrap();
        mire_sagas::migrate(&s).await.unwrap();
    }

    // Small dedicated pool just for the drain-progress count query, so it
    // never competes with the workers' pools.
    let count_pool = PgPoolOptions::new()
        .max_connections(2)
        .connect(&url)
        .await
        .unwrap();

    eprintln!(
        "\nsharded worker-fleet drain (num_shards={N_SHARDS}, {M} instances, \
         strong scaling):"
    );
    eprintln!("  workers   instances/s   drain   efficiency-vs-1");
    let mut base_rate: Option<f64> = None;

    for &k in &[1usize, 2, 4] {
        // Fresh id space per fleet size so awaits never overlap across runs.
        let run_tag = uuid::Uuid::new_v4().simple().to_string();
        let ids: Vec<String> = (0..M).map(|i| format!("scs-{k}-{run_tag}-{i}")).collect();

        // Trigger all M suspended (untimed setup). Use a dedicated pool.
        let trig = Arc::new(runner(&url, 16).await);
        {
            // Concurrent triggers to keep setup quick.
            use tokio::sync::Semaphore;
            let sem = Arc::new(Semaphore::new(24));
            let mut hs = Vec::new();
            for id in &ids {
                let trig = trig.clone();
                let sem = sem.clone();
                let id = id.clone();
                hs.push(tokio::spawn(async move {
                    let _p = sem.acquire().await.unwrap();
                    trig.trigger::<Trigger>(&Requested { id: id.clone() })
                        .await
                        .expect("trigger")
                        .expect("started");
                }));
            }
            for h in hs {
                h.await.unwrap();
            }
        }

        // Let every 50ms deadline lapse before the timed drain begins.
        tokio::time::sleep(Duration::from_millis(150)).await;

        // Spin up K autonomous workers (own pool each), poll fast. TIMED until
        // all M reach terminal.
        let cancel = CancellationToken::new();
        let mut handles = Vec::with_capacity(k);
        for _ in 0..k {
            let w = SagaWorker::new(runner(&url, 16).await)
                .poll_interval(Duration::from_millis(20))
                .shard_lease_ttl(Duration::from_secs(2));
            let c = cancel.clone();
            handles.push(tokio::spawn(async move { w.run(c).await }));
        }

        let t0 = Instant::now();
        // Poll the count of non-terminal instances via the awaits table: once
        // no awaits remain for this run's instances, the drain is complete.
        // (Cheaper than loading M envelopes each poll.)
        loop {
            let remaining: i64 = sqlx::query_scalar(
                "SELECT count(*) FROM es_saga_awaits WHERE saga_name = $1 AND instance_id LIKE $2",
            )
            .bind(ScaleSaga::name())
            .bind(format!("scs-{k}-{run_tag}-%"))
            .fetch_one(&count_pool)
            .await
            .unwrap();
            if remaining == 0 {
                break;
            }
            if t0.elapsed() > Duration::from_secs(120) {
                panic!("drain did not complete within 120s ({remaining} awaits left, k={k})");
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
        let drain = t0.elapsed().as_secs_f64();
        cancel.cancel();
        for h in handles {
            h.await.unwrap().unwrap();
        }

        let rate = M as f64 / drain;
        let efficiency = match base_rate {
            None => {
                base_rate = Some(rate);
                1.0
            }
            Some(b) => rate / b,
        };
        eprintln!("  {k:>7}   {rate:>11.0}   {drain:>5.2}s   {efficiency:>5.2}x");

        // Correctness spot-check: a sample of instances are FailedTerminal.
        for id in ids.iter().step_by(M / 20) {
            assert_eq!(
                trig.status(id).await.unwrap(),
                Some(SagaInstanceStatus::FailedTerminal),
                "{id} should have timed out terminal"
            );
        }
    }
}
