//! stream positions, assigned while committing.
//!
//! every event stream position (`/stream` ids, `/subscribe` keys, relay
//! seqs) is handed out inside [`Sequencer::commit`], which holds one lock across assigning the
//! positions, committing the batch that writes their rows, and broadcasting
//! the events. so by construction:
//!
//! - position order is commit order. every position at or below a sequence's
//!   head is committed, so a missing row below the head is a permanent hole,
//!   never a write still in flight.
//! - a broadcast is only sent once its row is readable.
//! - broadcasts arrive in position order, across every writer.
//! - positions never repeat, even across restarts: a commit that goes past
//!   a counter's reservation reserves the next block in the same batch, and
//!   reopening resumes after the reservation, so positions that leave no row
//!   or whose rows get pruned are never handed out again. a crash skips the
//!   rest of the block, which is just more holes.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, PoisonError};

use fjall::OwnedWriteBatch;
use miette::{IntoDiagnostic, Result, WrapErr};

use super::schema::{Cursors, Ks};

/// the last position a counter hands out. every stream position fits in an
/// i64, the width relay seqs and jetstream `time_us` have on the wire.
const LAST_POSITION: u64 = i64::MAX as u64;

/// the one place stream positions are handed out.
pub(crate) struct Sequencer {
    lock: Mutex<()>,
    /// where each counter's next position is written down
    marks: Ks<Cursors>,
}

impl Sequencer {
    pub(crate) fn new(marks: Ks<Cursors>) -> Self {
        Self {
            lock: Mutex::new(()),
            marks,
        }
    }

    /// assign positions and stage their rows, commit, then publish, all under
    /// one lock. positions only advance if the batch commits.
    pub(crate) fn commit<'a, T, R>(
        &self,
        mut batch: OwnedWriteBatch,
        stage: impl FnOnce(&mut Assign<'a>, &mut OwnedWriteBatch) -> Result<T>,
        publish: impl FnOnce(T) -> R,
    ) -> Result<R> {
        let _guard = self.lock.lock().unwrap_or_else(PoisonError::into_inner);
        let mut assign = Assign {
            pending: Vec::new(),
            reserving: Vec::new(),
        };
        let staged = stage(&mut assign, &mut batch)?;
        assign.mark(&self.marks, &mut batch);
        batch.commit().into_diagnostic()?;
        assign.advance();
        Ok(publish(staged))
    }
}

/// positions taken by one sequenced commit. only [`Sequencer::commit`] can
/// make one, so positions can't be taken anywhere else.
pub(crate) struct Assign<'a> {
    /// the value each touched counter moves to once the batch commits
    pending: Vec<(&'a Counter, u64)>,
    /// the reservation each counter moves to once the batch commits, for the
    /// ones this commit took past theirs
    reserving: Vec<(&'a Counter, u64)>,
}

impl<'a> Assign<'a> {
    /// take `counter`'s next value. the first take in a commit begins at
    /// `start(committed next)` and later ones count up from there.
    fn take(&mut self, counter: &'a Counter, start: impl FnOnce(u64) -> u64) -> Result<u64> {
        let slot = match self
            .pending
            .iter()
            .position(|(pending, _)| std::ptr::eq(*pending, counter))
        {
            Some(slot) => slot,
            None => {
                let next = start(counter.next.load(Ordering::Acquire));
                self.pending.push((counter, next));
                self.pending.len() - 1
            }
        };
        let taken = self.pending[slot].1;
        miette::ensure!(taken <= LAST_POSITION, "stream positions are exhausted");
        self.pending[slot].1 = taken + 1;
        Ok(taken)
    }

    /// reserve the next block for each counter this commit took past its
    /// reservation, in the same batch. the rest write nothing, because a mark
    /// per commit fills the cursors memtable, and every flush of it holds up
    /// all writers on fsyncs
    fn mark(&mut self, marks: &Ks<Cursors>, batch: &mut OwnedWriteBatch) {
        for &(counter, next) in &self.pending {
            if next > counter.reserved.load(Ordering::Acquire) {
                let reserved = next.max(next.saturating_add(counter.block).min(LAST_POSITION));
                batch.insert(marks, counter.mark, reserved.to_be_bytes());
                self.reserving.push((counter, reserved));
            }
        }
    }

    fn advance(self) {
        for (counter, next) in self.pending {
            counter.next.store(next, Ordering::Release);
        }
        for (counter, reserved) in self.reserving {
            counter.reserved.store(reserved, Ordering::Release);
        }
    }
}

/// a position counter. `mark` holds a reservation: no commit has taken a
/// position at or past it, so reopening the database resumes there, past
/// positions that left no row or whose rows were pruned.
struct Counter {
    /// the first position not committed yet
    next: AtomicU64,
    /// the committed reservation, only moved inside a sequenced commit
    reserved: AtomicU64,
    /// how far past `next` a commit reserves when it runs out
    block: u64,
    mark: &'static [u8],
}

impl Counter {
    const fn new(mark: &'static [u8], block: u64) -> Self {
        Self {
            next: AtomicU64::new(0),
            reserved: AtomicU64::new(0),
            block,
            mark,
        }
    }

    /// continue at `next`, or at the reservation if a commit got further.
    fn resume_at(&self, marks: &Ks<Cursors>, next: u64) -> Result<()> {
        let marked = marks
            .get(self.mark)
            .into_diagnostic()?
            .map(|mark| {
                <[u8; 8]>::try_from(mark.as_ref())
                    .map(u64::from_be_bytes)
                    .into_diagnostic()
                    .wrap_err("a stream position mark must be 8 bytes")
            })
            .transpose()?;
        let reserved = marked.unwrap_or(0);
        let next = next.max(reserved);
        // only a corrupted or hand-edited database gets here
        miette::ensure!(
            next <= LAST_POSITION,
            "the database's next stream position {next} leaves no room for more"
        );
        self.next.store(next, Ordering::Release);
        self.reserved.store(reserved, Ordering::Release);
        Ok(())
    }
}

/// a counter whose values are only taken inside a sequenced commit.
pub(crate) struct Sequence(Counter);

/// positions a sequence reserves at a time. big enough that its mark is
/// rewritten once every thousand or so commits, so the cursors memtable
/// practically never fills up and flushes, and still nothing next to the
/// position space a crash can skip.
pub(in crate::db) const SEQUENCE_BLOCK: u64 = 1024;

impl Sequence {
    pub(crate) const fn new(mark: &'static [u8]) -> Self {
        Self(Counter::new(mark, SEQUENCE_BLOCK))
    }

    /// the last committed position, if any.
    pub(crate) fn head(&self) -> Option<u64> {
        self.next().checked_sub(1)
    }

    /// the position the next commit will take first.
    pub(crate) fn next(&self) -> u64 {
        self.0.next.load(Ordering::Acquire)
    }

    /// continue numbering after `last`, the last stored position, or at 1
    /// for none, when opening the database.
    pub(in crate::db) fn resume_after(&self, marks: &Ks<Cursors>, last: Option<u64>) -> Result<()> {
        self.0.resume_at(marks, next_after(last).max(1))
    }

    #[cfg(test)]
    pub(in crate::db) fn resume_at(&self, next: u64) {
        self.0.next.store(next, Ordering::Release);
    }

    pub(crate) fn take<'a>(&'a self, assign: &mut Assign<'a>) -> Result<u64> {
        assign.take(&self.0, |next| next)
    }
}

/// microsecond timestamps that strictly increase in commit order, even when
/// the wall clock stalls or steps back. a commit's timestamps are consecutive,
/// starting at the wall time or just after the previous commit's last,
/// whichever is later, so the wall clock is read once per commit.
#[cfg(feature = "jetstream")]
pub(crate) struct Clock(Counter);

/// how far ahead a clock reserves, in microseconds. a commit only reserves
/// when the wall clock passes the reservation, so this is about one mark a
/// second, and after a crash timestamps run at most this far ahead of the
/// wall clock until it catches up.
#[cfg(feature = "jetstream")]
const CLOCK_BLOCK_US: u64 = 1_000_000;

#[cfg(feature = "jetstream")]
impl Clock {
    pub(crate) const fn new(mark: &'static [u8]) -> Self {
        Self(Counter::new(mark, CLOCK_BLOCK_US))
    }

    /// continue after `last`, the last stored timestamp, when opening the
    /// database.
    pub(in crate::db) fn resume_after(&self, marks: &Ks<Cursors>, last: Option<u64>) -> Result<()> {
        self.0.resume_at(marks, next_after(last))
    }

    /// start the next commit `by` ahead of the wall clock, as if the wall
    /// clock had stepped back after it.
    #[cfg(all(test, feature = "indexer_stream"))]
    pub(crate) fn run_ahead(&self, by: std::time::Duration) {
        let now = u64::try_from(chrono::Utc::now().timestamp_micros()).unwrap();
        let ahead = now + u64::try_from(by.as_micros()).unwrap();
        self.0.next.store(ahead, Ordering::Release);
    }

    pub(crate) fn tick<'a>(&'a self, assign: &mut Assign<'a>) -> Result<u64> {
        assign.take(&self.0, |next| {
            let now = u64::try_from(chrono::Utc::now().timestamp_micros()).unwrap_or(0);
            next.max(now)
        })
    }
}

/// the position after `last`, saturating so a stored `u64::MAX` fails
/// resuming instead of wrapping.
fn next_after(last: Option<u64>) -> u64 {
    last.map_or(0, |last| last.saturating_add(1))
}

/// send `events` in position order. a run of rows that aren't broadcast is
/// announced by one `marker` for its last position, sent before the next
/// broadcast, so a subscriber reads them from the db before moving past them.
pub(crate) fn publish_in_order<B, P>(
    tx: &tokio::sync::broadcast::Sender<B>,
    events: impl IntoIterator<Item = (P, Option<B>)>,
    marker: impl Fn(P) -> B,
) {
    let mut unannounced = None;
    for (position, event) in events {
        let Some(event) = event else {
            unannounced = Some(position);
            continue;
        };
        if let Some(position) = unannounced.take() {
            let _ = tx.send(marker(position));
        }
        let _ = tx.send(event);
    }
    if let Some(position) = unannounced {
        let _ = tx.send(marker(position));
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;

    fn db() -> (tempfile::TempDir, crate::db::Db) {
        let tmp = tempfile::tempdir().unwrap();
        let db = crate::db::Db::open(&Config {
            database_path: tmp.path().to_path_buf(),
            ..Default::default()
        })
        .unwrap();
        (tmp, db)
    }

    #[test]
    fn positions_advance_only_when_the_batch_commits() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let seq = Sequence::new(b"test_seq");
        seq.resume_at(5);

        let failed = sequencer.commit(
            db.inner.batch(),
            |assign, _| {
                assert_eq!((seq.take(assign)?, seq.take(assign)?), (5, 6));
                Err::<(), _>(miette::miette!("stage failed"))
            },
            |_| {},
        );
        assert!(failed.is_err());
        assert_eq!(seq.next(), 5);

        let taken = sequencer
            .commit(
                db.inner.batch(),
                |assign, _| {
                    let taken = [seq.take(assign)?, seq.take(assign)?];
                    // readers never see positions that haven't committed
                    assert_eq!(seq.head(), Some(4));
                    Ok(taken)
                },
                |taken| taken,
            )
            .unwrap();
        assert_eq!(taken, [5, 6]);
        assert_eq!(seq.head(), Some(6));
    }

    #[test]
    fn positions_stop_at_the_last_one() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let seq = Sequence::new(b"test_seq");
        seq.resume_at(LAST_POSITION);

        let failed = sequencer.commit(
            db.inner.batch(),
            |assign, _| {
                assert_eq!(seq.take(assign)?, LAST_POSITION);
                seq.take(assign)
            },
            |_| {},
        );
        assert!(failed.is_err());
        assert_eq!(seq.next(), LAST_POSITION);
    }

    #[test]
    fn a_stored_position_needs_room_after_it() {
        let (_tmp, db) = db();
        let seq = Sequence::new(b"test_seq");
        seq.resume_after(&db.cursors, None).unwrap();
        assert_eq!(seq.next(), 1);
        seq.resume_after(&db.cursors, Some(LAST_POSITION - 1))
            .unwrap();
        assert_eq!(seq.next(), LAST_POSITION);
        assert!(seq.resume_after(&db.cursors, Some(LAST_POSITION)).is_err());
        assert!(seq.resume_after(&db.cursors, Some(u64::MAX)).is_err());
    }

    #[test]
    fn a_counter_resumes_past_what_its_commits_took() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let seq = Sequence::new(b"test_seq");
        seq.resume_after(&db.cursors, None).unwrap();
        sequencer
            .commit(
                db.inner.batch(),
                |assign, _| (0..3).try_for_each(|_| seq.take(assign).map(drop)),
                |_| {},
            )
            .unwrap();

        // none of those positions has a row to resume after, so it resumes at the reservation
        let reopened = Sequence::new(b"test_seq");
        reopened.resume_after(&db.cursors, None).unwrap();
        assert_eq!(reopened.next(), 4 + SEQUENCE_BLOCK);
        // rows further along still win
        reopened
            .resume_after(&db.cursors, Some(SEQUENCE_BLOCK + 9))
            .unwrap();
        assert_eq!(reopened.next(), SEQUENCE_BLOCK + 10);
    }

    #[test]
    fn positions_taken_past_a_reservation_are_not_repeated_after_a_crash() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let seq = Sequence::new(b"test_seq");
        seq.resume_after(&db.cursors, None).unwrap();
        let mut last = 0;
        for _ in 0..SEQUENCE_BLOCK + 500 {
            last = sequencer
                .commit(db.inner.batch(), |assign, _| seq.take(assign), |t| t)
                .unwrap();
        }

        // a crash keeps only what committed, so a fresh counter reading the marks is one
        let reopened = Sequence::new(b"test_seq");
        reopened.resume_after(&db.cursors, None).unwrap();
        let next = reopened.next();
        assert!(next > last, "{next} was already taken");
        assert!(
            next <= last + 1 + SEQUENCE_BLOCK,
            "{next} skipped more than a block"
        );
    }

    #[test]
    fn an_exact_mark_from_before_reservations_resumes_where_it_stopped() {
        let (_tmp, db) = db();
        // older databases wrote the next position itself on every commit
        db.cursors.insert(b"test_seq", 7_u64.to_be_bytes()).unwrap();
        let sequencer = Sequencer::new(db.cursors.clone());
        let seq = Sequence::new(b"test_seq");
        seq.resume_after(&db.cursors, None).unwrap();
        assert_eq!(seq.next(), 7);
        let taken = sequencer
            .commit(db.inner.batch(), |assign, _| seq.take(assign), |t| t)
            .unwrap();
        assert_eq!(taken, 7);

        let reopened = Sequence::new(b"test_seq");
        reopened.resume_after(&db.cursors, None).unwrap();
        assert_eq!(reopened.next(), 8 + SEQUENCE_BLOCK);
    }

    #[test]
    fn a_hundred_thousand_commits_dont_fill_the_cursors_memtable() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let seq = Sequence::new(b"stream_next|test_ids");
        seq.resume_after(&db.cursors, None).unwrap();
        for _ in 0..100_000 {
            sequencer
                .commit(db.inner.batch(), |assign, _| seq.take(assign), |_| {})
                .unwrap();
        }
        // a full memtable asks a background worker to rotate it, so give that a moment
        let deadline = std::time::Instant::now() + std::time::Duration::from_millis(500);
        while std::time::Instant::now() < deadline {
            assert_eq!(
                db.cursors.sealed_memtable_count() + db.cursors.table_count(),
                0
            );
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
    }

    #[test]
    fn a_position_mark_is_read_big_endian() {
        let (_tmp, db) = db();
        let seq = Sequence::new(b"test_seq");
        db.cursors
            .insert(b"test_seq", [0, 0, 0, 0, 0, 0, 1, 2])
            .unwrap();
        seq.resume_after(&db.cursors, None).unwrap();
        assert_eq!(seq.next(), 258);

        db.cursors.insert(b"test_seq", [0; 7]).unwrap();
        assert!(seq.resume_after(&db.cursors, None).is_err());
    }

    #[cfg(feature = "jetstream")]
    #[test]
    fn clock_increases_even_when_ahead_of_the_wall_clock() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let clock = Clock::new(b"test_clock");
        let ahead = u64::try_from(chrono::Utc::now().timestamp_micros()).unwrap() + 60_000_000;
        clock.resume_after(&db.cursors, Some(ahead)).unwrap();

        // a failed commit takes nothing
        let failed = sequencer.commit(
            db.inner.batch(),
            |assign, _| {
                clock.tick(assign)?;
                Err::<(), _>(miette::miette!("stage failed"))
            },
            |_| {},
        );
        assert!(failed.is_err());

        let ticks = sequencer
            .commit(
                db.inner.batch(),
                |assign, _| Ok([clock.tick(assign)?, clock.tick(assign)?]),
                |ticks| ticks,
            )
            .unwrap();
        assert_eq!(ticks, [ahead + 1, ahead + 2]);
    }

    #[cfg(feature = "jetstream")]
    #[test]
    fn a_commit_reads_the_wall_clock_once() {
        let (_tmp, db) = db();
        let sequencer = Sequencer::new(db.cursors.clone());
        let clock = Clock::new(b"test_clock");
        let wall = || u64::try_from(chrono::Utc::now().timestamp_micros()).unwrap();

        let before = wall();
        let ticks = sequencer
            .commit(
                db.inner.batch(),
                |assign, _| {
                    let first = clock.tick(assign)?;
                    // slow staging doesn't spread the commit's timestamps out
                    std::thread::sleep(std::time::Duration::from_millis(5));
                    Ok([first, clock.tick(assign)?, clock.tick(assign)?])
                },
                |ticks| ticks,
            )
            .unwrap();
        assert!(ticks[0] >= before);
        assert_eq!(ticks, [ticks[0], ticks[0] + 1, ticks[0] + 2]);

        // the next commit starts at the wall time, which has moved on
        let next = sequencer
            .commit(db.inner.batch(), |assign, _| clock.tick(assign), |t| t)
            .unwrap();
        assert!(next >= ticks[0] + 5_000);
    }

    #[test]
    fn unbroadcast_runs_are_announced_before_the_next_broadcast() {
        let (tx, mut rx) = tokio::sync::broadcast::channel(16);
        publish_in_order(
            &tx,
            [
                (1, None),
                (2, None),
                (3, Some("live 3")),
                (4, Some("live 4")),
                (5, None),
            ],
            |position| match position {
                2 => "marker 2",
                5 => "marker 5",
                _ => panic!("unexpected marker for {position}"),
            },
        );
        let received: Vec<_> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
        assert_eq!(received, ["marker 2", "live 3", "live 4", "marker 5"]);
    }
}