//! stream subscribers against the real write path, with a live commit landing
//! while a backfill batch is still open: the interleaving that once made
//! both streams skip the backfill (hydrant-abq).

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

use jacquard_common::types::did::Did;
use jacquard_common::types::string::Tid;
use tokio::sync::mpsc;

use crate::config::Config;
use crate::db::Txn;
use crate::db::types::{DbAction, DbRkey, DbTid};
use crate::state::AppState;

use super::StreamOptions;

const BACKFILLED: &str = "did:plc:aaaaaaaaaaaaaaaaaaaaaaaa";
const LIVE: &str = "did:plc:zzzzzzzzzzzzzzzzzzzzzzzz";
const COLLECTION: &str = "app.bsky.feed.post";
const REPO_SIZE: usize = 100;

fn state() -> (tempfile::TempDir, Arc<AppState>, StreamOptions) {
    let tmp = tempfile::tempdir().unwrap();
    let config = Config {
        database_path: tmp.path().to_path_buf(),
        ..Default::default()
    };
    let state = Arc::new(AppState::new(&config).unwrap());
    (tmp, state, StreamOptions::from_config(&config))
}

fn put(records: &mut crate::db::RecordTxn<'_, '_, '_>, n: usize) {
    let block = crate::car::CarBlock::from_body(
        serde_ipld_dagcbor::to_vec(&serde_json::json!({ "$type": COLLECTION, "n": n })).unwrap(),
    );
    let rkey = DbRkey::Str(smol_str::format_smolstr!("r{n}"));
    records
        .put_record(COLLECTION, &rkey, &block, DbAction::Create)
        .unwrap();
}

fn rev() -> DbTid {
    DbTid::from(&Tid::now_0())
}

/// the backfill worker's write: a whole repo in one uncommitted batch.
fn stage_backfill(state: &AppState) -> Txn<'_> {
    let did = Did::new_static(BACKFILLED).unwrap();
    let mut txn = Txn::new(&state.db);
    let mut records = txn.backfill_records(state, &rev(), &did).unwrap();
    (0..REPO_SIZE).for_each(|n| put(&mut records, n));
    records.finish().unwrap();
    txn
}

/// one live commit, the way the indexer shard writes it.
fn commit_live(state: &AppState) {
    commit_live_records(state, REPO_SIZE..REPO_SIZE + 1);
}

/// one live commit writing a record for each of `records`.
fn commit_live_records(state: &AppState, records: std::ops::Range<usize>) {
    let did = Did::new_static(LIVE).unwrap();
    let mut txn = Txn::new(&state.db);
    let mut staged = txn.records(state, &rev(), &did).unwrap();
    records.for_each(|n| put(&mut staged, n));
    staged.finish().unwrap();
    txn.commit().unwrap();
}

/// dids of what a subscriber sent, collected until `until` holds or a
/// deadline passes.
struct Seen<T> {
    rx: mpsc::Receiver<T>,
    did_of: fn(T) -> String,
    dids: Vec<String>,
}

impl<T> Seen<T> {
    fn until(&mut self, until: impl Fn(&[String]) -> bool) -> &mut Self {
        let deadline = Instant::now() + Duration::from_secs(5);
        while Instant::now() < deadline && !until(&self.dids) {
            match self.rx.try_recv() {
                Ok(item) => self.dids.push((self.did_of)(item)),
                Err(_) => std::thread::sleep(Duration::from_millis(5)),
            }
        }
        self
    }

    fn live(&mut self) -> &mut Self {
        self.until(|dids| dids.iter().any(|did| did == LIVE))
    }

    /// wait a moment longer, for anything sent that shouldn't have been
    fn settle(&mut self) -> &mut Self {
        std::thread::sleep(Duration::from_millis(100));
        while let Ok(item) = self.rx.try_recv() {
            self.dids.push((self.did_of)(item));
        }
        self
    }

    fn counts(&self) -> (usize, usize) {
        let count = |did| self.dids.iter().filter(|seen| *seen == did).count();
        (count(BACKFILLED), count(LIVE))
    }
}

fn stream(state: &Arc<AppState>, opts: StreamOptions, cursor: Option<u64>) -> Seen<StreamItem> {
    let (tx, rx) = mpsc::channel(4096);
    let state = state.clone();
    std::thread::spawn(move || {
        super::event_stream_thread::<super::Decoded>(state, tx, cursor, opts)
    });
    Seen {
        rx,
        did_of: |item| item.unwrap().record.unwrap().did.as_str().to_owned(),
        dids: Vec::new(),
    }
}

type StreamItem = Result<crate::control::Event, crate::control::StreamError>;

/// wait until `n` live-only subscribers have subscribed, so they see what
/// follows.
fn wait_subscribed(n: usize, receivers: impl Fn() -> usize) {
    let deadline = Instant::now() + Duration::from_secs(5);
    while receivers() < n {
        assert!(Instant::now() < deadline, "subscriber never subscribed");
        std::thread::sleep(Duration::from_millis(5));
    }
}

#[test]
fn stream_tail_gets_a_backfill_committed_after_a_live_commit() {
    let (_tmp, state, opts) = state();
    let mut seen = stream(&state, opts, None);
    wait_subscribed(1, || state.db.stream.event_tx.receiver_count());

    let backfill = stage_backfill(&state);
    commit_live(&state);
    seen.live();
    backfill.commit().unwrap();

    seen.until(|dids| dids.len() > REPO_SIZE).settle();
    assert_eq!(seen.counts(), (REPO_SIZE, 1));
}

#[test]
fn stream_replay_gets_a_backfill_committed_after_a_live_commit() {
    let (_tmp, state, opts) = state();
    let backfill = stage_backfill(&state);
    commit_live(&state);
    let mut seen = stream(&state, opts, Some(0));
    seen.live();
    backfill.commit().unwrap();

    seen.until(|dids| dids.len() > REPO_SIZE).settle();
    assert_eq!(seen.counts(), (REPO_SIZE, 1));
}

#[cfg(feature = "jetstream")]
mod jetstream {
    use super::*;
    use crate::control::{JetstreamFilter, JetstreamSubscriberOptions};

    type Item = Result<bytes::Bytes, crate::control::JetstreamStreamError>;

    fn subscribe(
        state: &Arc<AppState>,
        opts: StreamOptions,
        cursor: Option<i64>,
        wanted_event_types: &[&str],
    ) -> Seen<Item> {
        subscribe_with_capacity(state, opts, 4096, cursor, wanted_event_types)
    }

    fn subscribe_with_capacity(
        state: &Arc<AppState>,
        opts: StreamOptions,
        capacity: usize,
        cursor: Option<i64>,
        wanted_event_types: &[&str],
    ) -> Seen<Item> {
        let (tx, rx) = mpsc::channel(capacity);
        let state = state.clone();
        let wanted: Vec<_> = wanted_event_types.iter().map(|t| t.to_string()).collect();
        let filter =
            JetstreamFilter::new(JetstreamSubscriberOptions::parse(&[], &[], 0, &wanted).unwrap());
        std::thread::spawn(move || {
            super::super::jetstream_stream_thread(state, tx, cursor, filter, opts)
        });
        Seen {
            rx,
            did_of: |item| {
                let json: serde_json::Value = serde_json::from_slice(&item.unwrap()).unwrap();
                json["did"].as_str().unwrap().to_owned()
            },
            dids: Vec::new(),
        }
    }

    fn a_second_ago() -> i64 {
        chrono::Utc::now().timestamp_micros() - 1_000_000
    }

    #[test]
    fn replay_gets_a_backfill_committed_after_a_live_commit() {
        let (_tmp, state, opts) = state();
        let cursor = a_second_ago();
        let backfill = stage_backfill(&state);
        commit_live(&state);
        let mut seen = subscribe(&state, opts, Some(cursor), &["live", "historical"]);
        seen.live();
        backfill.commit().unwrap();

        seen.until(|dids| dids.len() > REPO_SIZE).settle();
        assert_eq!(seen.counts(), (REPO_SIZE, 1));
    }

    #[test]
    fn tail_gets_backfills_only_when_asked() {
        let (_tmp, state, opts) = state();
        let mut historical = subscribe(&state, opts, None, &["live", "historical"]);
        let mut default = subscribe(&state, opts, None, &[]);
        wait_subscribed(2, || state.db.jetstream.tx.receiver_count());

        let backfill = stage_backfill(&state);
        commit_live(&state);
        backfill.commit().unwrap();

        historical.until(|dids| dids.len() > REPO_SIZE).settle();
        assert_eq!(historical.counts(), (REPO_SIZE, 1));
        default.live().settle();
        assert_eq!(default.counts(), (0, 1));
    }

    #[test]
    fn a_lagging_tail_gets_backfills_only_when_asked() {
        // more live events than the broadcast channel holds, per commit
        const LAG: usize = 600;
        let (_tmp, state, opts) = state();
        let opts = StreamOptions {
            pending_event_limit: 1,
            ..opts
        };
        let mut historical =
            subscribe_with_capacity(&state, opts, 64, None, &["live", "historical"]);
        let mut default = subscribe_with_capacity(&state, opts, 64, None, &[]);
        wait_subscribed(2, || state.db.jetstream.tx.receiver_count());

        // nobody reads, so both subscribers lag and catch up from the db
        commit_live_records(&state, REPO_SIZE..REPO_SIZE + LAG);
        stage_backfill(&state).commit().unwrap();
        commit_live_records(&state, REPO_SIZE + LAG..REPO_SIZE + 2 * LAG);

        historical
            .until(|dids| dids.len() >= REPO_SIZE + 2 * LAG)
            .settle();
        assert_eq!(historical.counts(), (REPO_SIZE, 2 * LAG));
        default.until(|dids| dids.len() >= 2 * LAG).settle();
        assert_eq!(default.counts(), (0, 2 * LAG));
    }

    #[test]
    fn a_cursor_ahead_of_the_wall_clock_replays() {
        let (_tmp, state, opts) = state();
        // committed time_us stays ahead of the wall clock after it steps back
        state.db.jetstream.clock.run_ahead(Duration::from_secs(60));
        commit_live(&state);
        let seen_through = state.db.jetstream.head().unwrap().unwrap().time_us;
        commit_live_records(&state, REPO_SIZE + 1..REPO_SIZE + 2);

        let cursor = i64::try_from(seen_through + 1).unwrap();
        let mut seen = subscribe(&state, opts, Some(cursor), &[]);
        seen.live().settle();
        assert_eq!(seen.counts(), (0, 1));
    }

    #[test]
    fn a_redaction_doesnt_reach_events_already_queued() {
        let (_tmp, state, opts) = state();
        let (tx, mut rx) = mpsc::channel(2);
        let filter = JetstreamFilter::new(JetstreamSubscriberOptions::default());
        let thread_state = state.clone();
        std::thread::spawn(move || {
            super::super::jetstream_stream_thread(thread_state, tx, None, filter, opts)
        });
        // live bodies are only kept inline while /stream has a subscriber too
        let _stream = stream(&state, opts, None);
        wait_subscribed(1, || state.db.jetstream.tx.receiver_count());
        wait_subscribed(1, || state.db.stream.event_tx.receiver_count());
        let taken = || wait_subscribed(1, || 1 - state.db.jetstream.tx.len().min(1));

        // fill the output and leave the subscriber waiting to send the next,
        // so the one after sits unrendered in its backlog
        commit_live_records(&state, REPO_SIZE..REPO_SIZE + 2);
        taken();
        commit_live_records(&state, REPO_SIZE + 2..REPO_SIZE + 3);
        taken();
        let did = Did::new_static(LIVE).unwrap();
        let rkey = DbRkey::Str(smol_str::format_smolstr!("r{}", REPO_SIZE + 2));
        crate::db::redact_record_bodies(
            &state.db,
            &did,
            COLLECTION,
            &rkey,
            crate::types::DeleteBodyTarget::Head,
        )
        .unwrap();

        let commits: Vec<serde_json::Value> = (0..3)
            .map(|_| {
                let item = rx.blocking_recv().unwrap().unwrap();
                serde_json::from_slice::<serde_json::Value>(&item).unwrap()["commit"].take()
            })
            .collect();
        // the third was queued before the redaction completed, so it goes out as committed
        assert!(
            commits.iter().all(|commit| commit.get("record").is_some()),
            "{commits:?}"
        );
    }

    #[test]
    fn live_json_is_what_a_replay_renders() {
        let (_tmp, state, opts) = state();
        // live bodies are only kept inline while /stream has a subscriber too
        let _stream = stream(&state, opts, None);
        wait_subscribed(1, || state.db.stream.event_tx.receiver_count());
        let mut broadcasts = state.db.jetstream.tx.subscribe();
        let cursor = a_second_ago();
        commit_live(&state);

        let crate::types::JetstreamBroadcast::Live(live) = broadcasts.blocking_recv().unwrap()
        else {
            panic!("a live commit is broadcast live");
        };
        let json = live.json.expect("built while a subscriber listened");
        let live_json: serde_json::Value = serde_json::from_slice(&json).unwrap();

        let (tx, mut rx) = mpsc::channel(4);
        let filter = JetstreamFilter::new(JetstreamSubscriberOptions::default());
        let thread_state = state.clone();
        std::thread::spawn(move || {
            super::super::jetstream_stream_thread(thread_state, tx, Some(cursor), filter, opts)
        });
        let replayed = rx.blocking_recv().unwrap().unwrap();
        let replayed: serde_json::Value = serde_json::from_slice(&replayed).unwrap();
        assert_eq!(live_json, replayed);

        let body = serde_json::json!({ "$type": COLLECTION, "n": REPO_SIZE });
        let cid = crate::car::CarBlock::from_body(serde_ipld_dagcbor::to_vec(&body).unwrap())
            .cid()
            .to_string();
        let rev = replayed["commit"]["rev"].as_str().unwrap();
        assert!(Tid::new(rev).is_ok(), "{rev}");
        assert_eq!(
            replayed,
            serde_json::json!({
                "did": LIVE,
                "time_us": live.position.time_us,
                "kind": "commit",
                "commit": {
                    "rev": rev,
                    "operation": "create",
                    "collection": COLLECTION,
                    "rkey": format!("r{REPO_SIZE}"),
                    "record": body,
                    "cid": cid,
                    "live": true,
                },
            })
        );
    }

    #[test]
    fn replay_sends_backfills_by_default() {
        let (_tmp, state, opts) = state();
        let cursor = a_second_ago();
        stage_backfill(&state).commit().unwrap();
        commit_live(&state);

        let mut seen = subscribe(&state, opts, Some(cursor), &[]);
        seen.until(|dids| dids.len() > REPO_SIZE).settle();
        assert_eq!(seen.counts(), (REPO_SIZE, 1));

        let mut live_only = subscribe(&state, opts, Some(cursor), &["live"]);
        live_only.live().settle();
        assert_eq!(live_only.counts(), (0, 1));
    }
}