use std::sync::Arc;
#[cfg(test)]
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use miette::{IntoDiagnostic, Result};
use rand::RngExt;
#[cfg(test)]
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info};
use url::Url;

use crate::config::FirehoseSource;
use crate::db::{self, keys};
#[cfg(feature = "firehose-diagnostics")]
use crate::ingest::firehose_stats::{FirehoseStatsSnapshot, RelayWorkerStatsSnapshot};
use crate::ingest::{BufferTx, firehose::FirehoseIngestor};
use crate::pds_daily_limit::PdsDailyReservation;
use crate::pds_discovery::{PdsHostError, canonical_pds_url};
use crate::state::AppState;

pub(super) struct FirehoseIngestorHandle {
    id: usize,
    cancel: CancellationToken,
}

impl Drop for FirehoseIngestorHandle {
    fn drop(&mut self) {
        self.cancel.cancel();
    }
}

pub(super) struct FirehoseShared {
    pub(super) buffer_tx: BufferTx,
    pub(super) verify_signatures: bool,
    pub(super) max_failures: usize,
}

/// a successful validated PDS admission.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PdsAdmission {
    StartedNew,
    RestartedKnown,
    AlreadyRunning,
}

/// a public policy rejection from validated PDS admission.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum PdsAdmissionRejection {
    InvalidHost(PdsHostError),
    HostBanned,
    DailyLimitReached,
}

pub(crate) type PdsAdmissionResult = std::result::Result<PdsAdmission, PdsAdmissionRejection>;

pub(super) struct SeededSourceAdds {
    pub(super) added: usize,
    pub(super) known: usize,
    pub(super) rejected: usize,
}

/// a snapshot of a single firehose relay's runtime state.
#[derive(Debug, Clone, serde::Serialize)]
pub struct FirehosePdsInfo {
    pub host: String,
    pub seq: i64,
    pub account_count: u64,
    pub status: &'static str,
}

/// details for the most recent recorded firehose source failure.
#[derive(Debug, Clone, serde::Serialize)]
pub struct FirehoseFailureInfo {
    pub at: i64,
    pub kind: String,
    pub detail: String,
}

/// a snapshot of a single firehose relay's runtime state.
#[derive(Debug, Clone, serde::Serialize)]
pub struct FirehoseSourceInfo {
    pub url: Url,
    /// true when this is a direct PDS connection; enables host authority enforcement.
    pub is_pds: bool,
    /// true when outbound sockets enforce the public-address boundary.
    pub public: bool,
    pub running: bool,
    pub failing: bool,
    pub throttled: bool,
    pub consecutive_failures: usize,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub throttled_until: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub retry_in_secs: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub last_failure: Option<FirehoseFailureInfo>,
    #[cfg(feature = "firehose-diagnostics")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stats: Option<FirehoseStatsSnapshot>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub host_status: Option<&'static str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pds: Option<FirehosePdsInfo>,
}

/// feature-gated runtime diagnostics for firehose ingestion internals.
#[cfg(feature = "firehose-diagnostics")]
#[derive(Debug, Clone, serde::Serialize)]
pub struct FirehoseDiagnosticsInfo {
    pub relay_worker: RelayWorkerStatsSnapshot,
}

/// runtime control over the firehose ingestor component.
#[derive(Clone)]
pub struct FirehoseHandle {
    pub(super) state: Arc<AppState>,
    /// set once by [`Hydrant::run`]; `None` means run() has not been called yet.
    pub(super) shared: Arc<std::sync::OnceLock<FirehoseShared>>,
    /// per-relay running tasks, keyed by url.
    pub(super) tasks: Arc<scc::HashMap<Url, FirehoseIngestorHandle>>,
    /// known source urls → is_pds flag; includes API-added (db-persisted) and static config sources.
    pub(super) known_sources: Arc<scc::HashMap<Url, bool>>,
    /// Sources admitted through the public boundary. Persisted legacy rows
    /// without a trusted marker are inserted here; absence is reserved for
    /// current operator/admin sources.
    pub(super) public_sources: Arc<scc::HashMap<Url, bool>>,
    /// ids assigned to spawned tasks
    next_task_id: Arc<AtomicUsize>,
    #[cfg(test)]
    admission_pause: Arc<Mutex<Option<(Arc<Notify>, Arc<Notify>)>>>,
    #[cfg(test)]
    spawn_enabled: Arc<std::sync::atomic::AtomicBool>,
}

impl FirehoseHandle {
    pub(super) fn new(state: Arc<AppState>) -> Self {
        Self {
            state,
            shared: Arc::new(std::sync::OnceLock::new()),
            tasks: Arc::new(scc::HashMap::new()),
            known_sources: Arc::new(scc::HashMap::new()),
            public_sources: Arc::new(scc::HashMap::new()),
            next_task_id: Arc::new(AtomicUsize::new(0)),
            #[cfg(test)]
            admission_pause: Arc::new(Mutex::new(None)),
            #[cfg(test)]
            spawn_enabled: Arc::new(std::sync::atomic::AtomicBool::new(true)),
        }
    }

    pub(super) async fn spawn_firehose_ingestor(
        &self,
        source: &FirehoseSource,
        shared: &FirehoseShared,
        delay_startup: bool,
    ) -> Result<()> {
        let public = self.public_sources.contains_sync(&source.url);
        let prepared = self
            .prepare_firehose_ingestor(source, shared, public)
            .await?;
        self.spawn_prepared_firehose_ingestor(source, prepared, delay_startup)
            .await;
        Ok(())
    }

    async fn prepare_firehose_ingestor(
        &self,
        source: &FirehoseSource,
        shared: &FirehoseShared,
        public: bool,
    ) -> Result<(Option<i64>, FirehoseIngestor)> {
        let start = db::get_firehose_cursor(&self.state.db, &source.url).await?;
        let ingestor = FirehoseIngestor::new(
            self.state.clone(),
            shared.buffer_tx.clone(),
            source.url.clone(),
            source.is_pds,
            public,
            self.state.filter.clone(),
            self.state.firehose_enabled.subscribe(),
            shared.verify_signatures,
            shared.max_failures,
        )
        .await;
        Ok((start, ingestor))
    }

    async fn spawn_prepared_firehose_ingestor(
        &self,
        source: &FirehoseSource,
        (start, ingestor): (Option<i64>, FirehoseIngestor),
        delay_startup: bool,
    ) {
        #[cfg(test)]
        if !self.spawn_enabled.load(Ordering::Relaxed) {
            drop(ingestor);
            return;
        }
        use std::sync::atomic::AtomicI64;

        // existing in-memory cursor state takes precedence.
        let _ = self
            .state
            .firehose_cursors
            .insert_async(source.url.clone(), AtomicI64::new(start.unwrap_or(0)))
            .await;
        info!(relay = %source.url, source.is_pds, cursor = ?start, "starting firehose ingestor");

        let id = self.next_task_id.fetch_add(1, Ordering::Relaxed);
        let cancel = CancellationToken::new();
        tokio::spawn({
            let relay_url = source.url.clone();
            let is_pds = source.is_pds;
            let tasks = self.tasks.clone();
            let token = cancel.clone();
            async move {
                // jitter connection start so a seeded page does not create a thundering herd.
                if delay_startup {
                    let max_jitter_ms = if is_pds && !cfg!(debug_assertions) {
                        60_000
                    } else {
                        2_000
                    };
                    let jitter_ms = rand::rng().random_range(0u64..max_jitter_ms);
                    debug!(
                        relay = %relay_url,
                        is_pds,
                        jitter_ms,
                        "delaying firehose ingestor startup"
                    );
                    tokio::select! {
                        _ = tokio::time::sleep(Duration::from_millis(jitter_ms)) => {}
                        _ = token.cancelled() => {
                            info!(relay = %relay_url, "firehose ingestor cancelled");
                            return;
                        }
                    }
                }
                tokio::select! {
                    res = ingestor.run() => {
                        // only remove our own entry because an upsert can replace us.
                        tasks.remove_if_async(&relay_url, |h| h.id == id).await;
                        match res {
                            Ok(()) => info!(relay = %relay_url, "firehose shut down!"),
                            Err(e) => error!(relay = %relay_url, err = %e, "firehose ingestor exited with error"),
                        }
                    },
                    _ = token.cancelled() => {
                        info!(relay = %relay_url, "firehose ingestor cancelled");
                    }
                }
            }
        });

        let handle = FirehoseIngestorHandle { id, cancel };
        // `upsert_async` returns the old handle. ignoring it drops and cancels the old task.
        self.tasks.upsert_async(source.url.clone(), handle).await;
    }

    #[cfg(test)]
    pub(crate) fn disable_spawns_for_test(&self) {
        self.spawn_enabled.store(false, Ordering::Relaxed);
    }

    #[cfg(test)]
    pub(crate) fn startable_for_test(&self) {
        let (buffer_tx, _) = BufferTx::channel(1);
        assert!(
            self.shared
                .set(FirehoseShared {
                    buffer_tx,
                    verify_signatures: false,
                    max_failures: 1,
                })
                .is_ok(),
            "test firehose shared must be set once"
        );
    }

    #[cfg(test)]
    pub(crate) fn pause_validated_admission_after_ban_check(&self) -> (Arc<Notify>, Arc<Notify>) {
        let pause = (Arc::new(Notify::new()), Arc::new(Notify::new()));
        *self
            .admission_pause
            .lock()
            .expect("admission test pause lock must not be poisoned") = Some(pause.clone());
        pause
    }

    #[cfg(test)]
    async fn pause_validated_admission_after_ban_check_if_requested(&self) {
        let pause = self
            .admission_pause
            .lock()
            .expect("admission test pause lock must not be poisoned")
            .clone();
        if let Some((checked, resume)) = pause {
            checked.notify_one();
            resume.notified().await;
        }
    }

    /// enable firehose ingestion, no-op if already enabled.
    pub fn enable(&self) {
        self.state.firehose_enabled.send_replace(true);
    }
    /// disable firehose ingestion, in-flight messages complete before pausing.
    pub fn disable(&self) {
        self.state.firehose_enabled.send_replace(false);
    }
    /// returns the current enabled state of firehose ingestion.
    pub fn is_enabled(&self) -> bool {
        *self.state.firehose_enabled.borrow()
    }

    /// returns `true` if this URL is already a known firehose source.
    /// either currently running or persisted (e.g. the host is offline but was previously added).
    pub fn is_source_known(&self, url: &Url) -> bool {
        self.known_sources.contains_sync(url)
    }

    /// return `true` if this source has a running firehose task (eg. its not offline).
    pub fn is_source_running(&self, url: &Url) -> bool {
        self.tasks.contains_sync(url)
    }

    /// list all known firehose sources, including offline ones pending retry.
    pub async fn list_sources(&self) -> Vec<FirehoseSourceInfo> {
        let now = chrono::Utc::now().timestamp();
        let meta = self.state.pds_meta.load();
        let mut out = Vec::with_capacity(self.known_sources.capacity());
        self.known_sources
            .iter_async(|url, &is_pds| {
                let running = self.tasks.contains_sync(url);
                let throttle = self.state.throttler.snapshot(url);
                let pds = is_pds.then(|| self.pds_info(url, &meta)).flatten();
                let host_status = pds.as_ref().map(|pds| pds.status);
                #[cfg(feature = "firehose-diagnostics")]
                let stats = self.state.firehose_stats.snapshot(url);
                let last_failure =
                    throttle
                        .last_failure
                        .clone()
                        .map(|failure| FirehoseFailureInfo {
                            at: failure.at,
                            kind: failure.kind,
                            detail: failure.detail,
                        });
                out.push(FirehoseSourceInfo {
                    url: url.clone(),
                    is_pds,
                    public: self.public_sources.contains_sync(url),
                    running,
                    failing: throttle.is_failing(),
                    throttled: throttle.is_throttled(now),
                    consecutive_failures: throttle.consecutive_failures,
                    throttled_until: (throttle.throttled_until != 0)
                        .then_some(throttle.throttled_until),
                    retry_in_secs: throttle.retry_in_secs(now),
                    last_failure,
                    #[cfg(feature = "firehose-diagnostics")]
                    stats,
                    host_status,
                    pds,
                });
                true
            })
            .await;
        out.sort_unstable_by(|a, b| a.url.as_str().cmp(b.url.as_str()));
        out
    }

    #[cfg(feature = "firehose-diagnostics")]
    pub fn diagnostics(&self) -> FirehoseDiagnosticsInfo {
        FirehoseDiagnosticsInfo {
            relay_worker: self.state.firehose_stats.relay_worker_snapshot(),
        }
    }

    fn pds_info(&self, url: &Url, meta: &crate::pds_meta::PdsMeta) -> Option<FirehosePdsInfo> {
        let host = url.host_str()?;
        let seq = self
            .state
            .firehose_cursors
            .peek_with(url, |_, cursor| cursor.load(Ordering::SeqCst))
            .or_else(|| {
                self.state
                    .db
                    .cursors
                    .get(keys::firehose_cursor_key_from_url(url))
                    .ok()
                    .flatten()
                    .and_then(|bytes| bytes.as_ref().try_into().ok().map(i64::from_be_bytes))
            })
            .unwrap_or(0);
        let account_count = self
            .state
            .db
            .get_count_sync(&keys::pds_account_count_key(host));
        Some(FirehosePdsInfo {
            host: host.to_string(),
            seq,
            account_count,
            status: meta.status(host).as_str(),
        })
    }

    /// add a direct PDS source through the public discovery policy.
    #[cfg_attr(not(feature = "relay"), allow(dead_code))]
    pub(crate) async fn add_source_validated(&self, url: Url) -> Result<PdsAdmissionResult> {
        self.admit_validated_source(url, false).await
    }

    async fn admit_validated_source(
        &self,
        url: Url,
        delay_startup: bool,
    ) -> Result<PdsAdmissionResult> {
        let (host, url) = match canonical_pds_url(&url) {
            Ok(source) => source,
            Err(error) => return Ok(Err(PdsAdmissionRejection::InvalidHost(error))),
        };
        let _admission = self.state.pds_admission.lock().await;

        if self.state.pds_meta.load().is_banned(host.as_str()) {
            return Ok(Err(PdsAdmissionRejection::HostBanned));
        }
        #[cfg(test)]
        self.pause_validated_admission_after_ban_check_if_requested()
            .await;
        if self.is_source_running(&url) {
            return Ok(Ok(PdsAdmission::AlreadyRunning));
        }
        if self.is_source_known(&url) {
            // Do not route this through admin add_source: that trusted override
            // intentionally clears public provenance. A validated request stays
            // public across restart and persistence.
            let shared = self
                .shared
                .get()
                .ok_or_else(|| miette::miette!("firehose worker not started"))?;
            let is_pds = self
                .known_sources
                .read_async(&url, |_, &is_pds| is_pds)
                .await
                .unwrap_or(true);
            let source = FirehoseSource {
                url: url.clone(),
                is_pds,
            };
            let throttle = self.state.throttler.get_handle(&url).await;
            throttle.record_success();
            let prepared = self
                .prepare_firehose_ingestor(&source, shared, self.public_sources.contains_sync(&url))
                .await?;
            self.spawn_prepared_firehose_ingestor(&source, prepared, delay_startup)
                .await;
            return Ok(Ok(PdsAdmission::RestartedKnown));
        }

        let Some(reservation) = self.state.pds_daily_limit.reserve() else {
            return Ok(Err(PdsAdmissionRejection::DailyLimitReached));
        };
        let shared = match self.shared.get() {
            Some(shared) => shared,
            None => {
                self.state.pds_daily_limit.refund(reservation);
                return Err(miette::miette!("firehose worker not started"));
            }
        };
        let source = FirehoseSource { url, is_pds: true };
        // all fallible startup preparation happens before taking a quota slot's write.
        let prepared = match self.prepare_firehose_ingestor(&source, shared, true).await {
            Ok(prepared) => prepared,
            Err(error) => {
                self.state.pds_daily_limit.refund(reservation);
                return Err(error);
            }
        };

        match self
            .persist_validated_source(&source, host.as_str(), reservation)
            .await
        {
            Ok(true) => {}
            Ok(false) => {
                self.state.pds_daily_limit.refund(reservation);
                return Ok(Err(PdsAdmissionRejection::HostBanned));
            }
            Err(error) => {
                self.state.pds_daily_limit.refund(reservation);
                return Err(error);
            }
        }

        let _ = self
            .known_sources
            .insert_async(source.url.clone(), true)
            .await;
        let _ = self
            .public_sources
            .insert_async(source.url.clone(), true)
            .await;
        let throttle = self.state.throttler.get_handle(&source.url).await;
        throttle.record_success();
        self.spawn_prepared_firehose_ingestor(&source, prepared, delay_startup)
            .await;
        Ok(Ok(PdsAdmission::StartedNew))
    }

    async fn persist_validated_source(
        &self,
        source: &FirehoseSource,
        host: &str,
        reservation: PdsDailyReservation,
    ) -> Result<bool> {
        let key = keys::firehose_source_key(source.url.as_str());
        let marker_key = keys::trusted_firehose_source_key(source.url.as_str());
        let status_key = db::pds_meta::pds_status_key(host);
        let ban_key = db::pds_meta::pds_ban_key(host);
        let value =
            rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds: true }).map_err(|error| {
                miette::miette!("failed to serialize firehose source meta: {error}")
            })?;
        let source_url = source.url.clone();
        self.state
            .db
            .run(move |db| {
                let operator_banned = db.filter.contains_key(ban_key).into_diagnostic()?;
                let observed_banned = db
                    .filter
                    .get(status_key)
                    .into_diagnostic()?
                    .map(|status| {
                        rmp_serde::from_slice::<crate::pds_meta::HostStatus>(status.as_ref())
                            .into_diagnostic()
                    })
                    .transpose()?
                    .is_some_and(|status| status == crate::pds_meta::HostStatus::Banned);
                if operator_banned || observed_banned {
                    return Ok(false);
                }

                let mut batch = db.inner.batch();
                if let Some((day, count)) = reservation.persisted() {
                    db::stage_pds_daily_adds(&mut batch, db, day, count);
                }
                batch.insert(&db.crawler, key, value);
                batch.remove(&db.crawler, marker_key);
                batch.commit().into_diagnostic()?;
                if let Err(error) = db.persist() {
                    tracing::error!(
                        source = %source_url,
                        err = %error,
                        "PDS source committed but durability persist failed"
                    );
                }
                Ok(true)
            })
            .await
    }

    /// add PDS sources discovered from a seed relay.
    ///
    /// each accepted source persists its own quota reservation with its source
    /// record. startup stays staggered so a page cannot stampede remote hosts.
    pub(super) async fn add_seeded_sources(&self, urls: Vec<Url>) -> Result<SeededSourceAdds> {
        let mut adds = SeededSourceAdds {
            added: 0,
            known: 0,
            rejected: 0,
        };
        for url in urls {
            if self.is_source_known(&url) {
                adds.known += 1;
                continue;
            }
            match self.admit_validated_source(url.clone(), true).await? {
                Ok(PdsAdmission::StartedNew) => adds.added += 1,
                Ok(PdsAdmission::RestartedKnown | PdsAdmission::AlreadyRunning) => {
                    adds.known += 1;
                }
                Err(rejection) => {
                    adds.rejected += 1;
                    tracing::warn!(source = %url, ?rejection, "rejected seeded PDS source");
                }
            }
        }
        Ok(adds)
    }

    /// add a new firehose source at runtime, persisting it to the database.
    ///
    /// if a source with the same URL already exists, it is replaced: the
    /// running task is stopped and a new one is started with the new `is_pds`
    /// setting. existing cursor state for the URL is preserved.
    pub async fn add_source(&self, url: Url, is_pds: bool) -> Result<()> {
        let _admission = self.state.pds_admission.lock().await;
        let shared = self
            .shared
            .get()
            .ok_or_else(|| miette::miette!("firehose worker not started"))?;

        // persist to db first
        let key = keys::firehose_source_key(url.as_str());
        let marker_key = keys::trusted_firehose_source_key(url.as_str());
        self.state
            .db
            .run(move |db| {
                let mut batch = db.inner.batch();
                let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds }).map_err(|e| {
                    miette::miette!("failed to serialize firehose source meta: {e}")
                })?;
                batch.insert(&db.crawler, key, &value);
                batch.insert(&db.crawler, marker_key, []);
                batch.commit().into_diagnostic()?;
                db.persist()
            })
            .await?;
        self.public_sources.remove_async(&url).await;

        let _ = self.known_sources.upsert_async(url.clone(), is_pds).await;

        // reset failure state so the fresh task gets a clean slate.
        // if the previous task exited after max failures, the failure counter
        // would otherwise cause the new task to exit immediately.
        let throttle = self.state.throttler.get_handle(&url).await;
        throttle.record_success();

        self.spawn_firehose_ingestor(&FirehoseSource { url, is_pds }, shared, false)
            .await?;

        Ok(())
    }

    /// remove a firehose source at runtime.
    ///
    /// returns `true` if the source was found and removed, `false` otherwise.
    /// if the source was added via the API, it is removed from the database;
    /// if it came from the static config, only the running task is stopped.
    pub async fn remove_source(&self, url: &Url) -> Result<bool> {
        let _admission = self.state.pds_admission.lock().await;
        if self.known_sources.contains_async(url).await {
            let url_str = url.to_string();
            let marker_key = keys::trusted_firehose_source_key(&url_str);
            self.state
                .db
                .run(move |db| {
                    db.crawler
                        .remove(keys::firehose_source_key(&url_str))
                        .into_diagnostic()?;
                    db.crawler.remove(marker_key).into_diagnostic()?;
                    db.persist()
                })
                .await?;
            self.known_sources.remove_async(url).await;
            self.public_sources.remove_async(url).await;
        }

        Ok(self.tasks.remove_async(url).await.is_some())
    }

    /// restart an offline firehose source without touching the database or daily limits.
    pub(super) async fn restart_source(&self, url: Url) -> Result<()> {
        let _admission = self.state.pds_admission.lock().await;
        let Some(is_pds) = self
            .known_sources
            .read_async(&url, |_, &is_pds| is_pds)
            .await
        else {
            return Ok(());
        };
        if self.tasks.contains_sync(&url) {
            return Ok(());
        }
        let shared = self
            .shared
            .get()
            .ok_or_else(|| miette::miette!("firehose worker not started"))?;

        // clear the failure counter so the new task isn't immediately terminated
        let throttle = self.state.throttler.get_handle(&url).await;
        throttle.record_success();

        self.spawn_firehose_ingestor(&FirehoseSource { url, is_pds }, shared, true)
            .await
    }

    /// reset the stored firehose cursor for a given URL.
    pub async fn reset_cursor(&self, url: &str) -> Result<()> {
        let url = Url::parse(url).into_diagnostic()?;
        let key = keys::firehose_cursor_key_from_url(&url);
        self.state
            .db
            .run(move |db| {
                db.cursors.remove(key).into_diagnostic()?;
                db.persist()
            })
            .await?;

        self.state.firehose_cursors.remove_async(&url).await;

        Ok(())
    }
}

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

    fn test_config(dir: &tempfile::TempDir) -> Config {
        Config {
            database_path: dir.path().to_path_buf(),
            ..Default::default()
        }
    }

    fn test_config_with_limit(dir: &tempfile::TempDir, new_host_limit: u64) -> Config {
        Config {
            database_path: dir.path().to_path_buf(),
            new_host_limit: Some(new_host_limit),
            ..Default::default()
        }
    }

    fn startable(hydrant: &Hydrant) {
        hydrant.firehose.startable_for_test();
    }

    async fn persisted_sources(hydrant: &Hydrant) -> Result<Vec<(FirehoseSource, bool)>> {
        hydrant
            .state
            .db
            .run(crate::db::load_persisted_firehose_sources)
            .await
    }

    /// `upsert_async` returns the previous handle, whose drop cancels its task.
    #[tokio::test]
    async fn replacing_a_source_cancels_the_previous_ingestor() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(&dir)).await?;
        startable(&hydrant);
        let url = Url::parse("wss://replace.example/").into_diagnostic()?;

        hydrant.firehose.add_source(url.clone(), true).await?;
        let (first_id, first_token) = hydrant
            .firehose
            .tasks
            .read_async(&url, |_, handle| (handle.id, handle.cancel.clone()))
            .await
            .expect("first source must be running");
        assert!(!first_token.is_cancelled(), "a fresh handle is live");

        hydrant.firehose.add_source(url.clone(), true).await?;

        let (second_id, second_token) = hydrant
            .firehose
            .tasks
            .read_async(&url, |_, handle| (handle.id, handle.cancel.clone()))
            .await
            .expect("a source must still be running");
        assert_ne!(first_id, second_id, "the replacement mints a new task id");
        assert!(
            first_token.is_cancelled(),
            "replacing the handle must cancel the previous ingestor"
        );
        assert!(!second_token.is_cancelled(), "the new ingestor stays live");
        assert_eq!(hydrant.firehose.tasks.len(), 1, "one handle per url");
        assert!(hydrant.firehose.is_source_running(&url));

        hydrant.firehose.tasks.remove_async(&url).await;
        assert!(second_token.is_cancelled(), "dropping the handle stops it");
        assert!(!hydrant.firehose.is_source_running(&url));
        Ok(())
    }

    #[tokio::test]
    async fn admin_override_clears_public_provenance_in_memory_and_db() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(&dir)).await?;
        startable(&hydrant);
        hydrant.firehose.disable_spawns_for_test();
        let url = Url::parse("wss://public.example/").into_diagnostic()?;

        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Ok(PdsAdmission::StartedNew)
        );
        assert!(hydrant.firehose.public_sources.contains_sync(&url));
        assert!(
            persisted_sources(&hydrant)
                .await?
                .iter()
                .any(|(source, trusted)| { source.url == url && !trusted })
        );

        hydrant.firehose.add_source(url.clone(), false).await?;
        assert!(
            !hydrant.firehose.public_sources.contains_sync(&url),
            "trusted admin replacement must not inherit public provenance"
        );
        assert!(
            persisted_sources(&hydrant)
                .await?
                .iter()
                .any(|(source, trusted)| { source.url == url && *trusted })
        );
        Ok(())
    }

    #[tokio::test]
    async fn validated_request_preserves_a_known_trusted_source() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(&dir)).await?;
        startable(&hydrant);
        hydrant.firehose.disable_spawns_for_test();
        let url = Url::parse("wss://operator.example/").into_diagnostic()?;

        hydrant.firehose.add_source(url.clone(), false).await?;
        let throttle = hydrant.state.throttler.get_handle(&url).await;
        throttle.record_failure();
        assert!(hydrant.state.throttler.snapshot(&url).is_failing());

        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Ok(PdsAdmission::RestartedKnown)
        );
        assert!(!hydrant.state.throttler.snapshot(&url).is_failing());
        assert_eq!(
            hydrant
                .firehose
                .known_sources
                .read_async(&url, |_, is_pds| *is_pds)
                .await,
            Some(false)
        );
        assert!(!hydrant.firehose.public_sources.contains_sync(&url));
        assert!(
            persisted_sources(&hydrant)
                .await?
                .iter()
                .any(|(source, trusted)| source.url == url && *trusted && !source.is_pds)
        );
        Ok(())
    }

    #[tokio::test]
    async fn validated_admission_rejects_invalid_and_private_hosts() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(&dir)).await?;

        for url in [
            "wss://127.0.0.1/",
            "ws://public.example/",
            "wss://public.example/path",
        ] {
            assert!(matches!(
                hydrant
                    .firehose
                    .add_source_validated(Url::parse(url).into_diagnostic()?)
                    .await?,
                Err(PdsAdmissionRejection::InvalidHost(_))
            ));
        }
        assert!(persisted_sources(&hydrant).await?.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn seeded_banned_host_is_not_persisted_or_started() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(&dir)).await?;
        startable(&hydrant);
        hydrant.pds.ban("blocked.example").await?;
        let url = Url::parse("wss://blocked.example/").into_diagnostic()?;

        let adds = hydrant
            .firehose
            .add_seeded_sources(vec![url.clone()])
            .await?;

        assert_eq!(adds.added, 0);
        assert_eq!(adds.rejected, 1);
        assert!(!hydrant.firehose.is_source_known(&url));
        assert!(!hydrant.firehose.is_source_running(&url));
        assert!(persisted_sources(&hydrant).await?.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn seeded_admission_stops_at_the_daily_limit() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?;
        startable(&hydrant);
        let first = Url::parse("wss://first-seed.example/").into_diagnostic()?;
        let second = Url::parse("wss://second-seed.example/").into_diagnostic()?;

        let adds = hydrant
            .firehose
            .add_seeded_sources(vec![first.clone(), second.clone()])
            .await?;

        assert_eq!(adds.added, 1);
        assert_eq!(adds.rejected, 1);
        assert_eq!(persisted_sources(&hydrant).await?.len(), 1);
        assert!(matches!(
            crate::db::load_pds_daily_adds(&hydrant.state.db)?,
            Some((_, 1))
        ));
        hydrant.firehose.tasks.remove_async(&first).await;
        Ok(())
    }

    #[tokio::test]
    async fn committed_admission_is_not_refunded_when_persist_fails() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?;
        startable(&hydrant);
        let url = Url::parse("wss://post-commit-persist.example/").into_diagnostic()?;
        hydrant.state.db.fail_next_persist_for_test();

        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Ok(PdsAdmission::StartedNew)
        );
        assert!(hydrant.firehose.is_source_running(&url));
        assert!(matches!(
            crate::db::load_pds_daily_adds(&hydrant.state.db)?,
            Some((_, 1))
        ));
        assert_eq!(
            hydrant
                .firehose
                .add_source_validated(
                    Url::parse("wss://second-post-commit.example/").into_diagnostic()?
                )
                .await?,
            Err(PdsAdmissionRejection::DailyLimitReached)
        );

        hydrant.firehose.tasks.remove_async(&url).await;
        Ok(())
    }

    #[tokio::test]
    async fn running_validated_source_does_not_spend_another_slot() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?;
        startable(&hydrant);
        let url = Url::parse("wss://running.example/").into_diagnostic()?;

        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Ok(PdsAdmission::StartedNew)
        );
        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Ok(PdsAdmission::AlreadyRunning)
        );
        assert!(matches!(
            crate::db::load_pds_daily_adds(&hydrant.state.db)?,
            Some((_, 1))
        ));
        hydrant.firehose.tasks.remove_async(&url).await;
        Ok(())
    }

    #[tokio::test]
    async fn legacy_raw_ban_is_enforced_and_removed_by_unban() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let cfg = test_config(&dir);
        let raw_host = "BANNED.EXAMPLE.";
        let url = Url::parse("wss://banned.example/").into_diagnostic()?;

        {
            let hydrant = Hydrant::new(cfg.clone()).await?;
            let mut batch = hydrant.state.db.inner.batch();
            crate::db::pds_meta::set_status(
                &mut batch,
                &hydrant.state.db.filter,
                raw_host,
                crate::pds_meta::HostStatus::Banned,
            )?;
            batch.commit().into_diagnostic()?;
            crate::db::migration::rewind_version_for_test(&hydrant.state.db, 9)?;
            hydrant.state.db.persist()?;
        }

        let hydrant = Hydrant::new(cfg.clone()).await?;
        hydrant.firehose.startable_for_test();
        assert_eq!(
            hydrant.state.pds_meta.load().status("banned.example"),
            crate::pds_meta::HostStatus::Active
        );
        assert!(
            hydrant
                .state
                .db
                .filter
                .contains_key(crate::db::pds_meta::pds_ban_key("banned.example"))
                .into_diagnostic()?
        );
        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Err(PdsAdmissionRejection::HostBanned)
        );

        hydrant.pds.unban(raw_host).await?;
        assert!(!hydrant.pds.is_banned("banned.example"));
        assert!(
            !hydrant
                .state
                .db
                .filter
                .contains_key(crate::db::pds_meta::pds_ban_key("banned.example"))
                .into_diagnostic()?
        );
        drop(hydrant);

        let reopened = Hydrant::new(cfg).await?;
        assert!(!reopened.pds.is_banned("banned.example"));
        Ok(())
    }

    #[tokio::test]
    async fn durable_ban_prevents_validated_admission() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config_with_limit(&dir, 1)).await?;
        startable(&hydrant);
        let host = "durably-banned.example";
        let url = Url::parse("wss://durably-banned.example/").into_diagnostic()?;

        hydrant
            .state
            .db
            .run(move |db| {
                let mut batch = db.inner.batch();
                crate::db::pds_meta::set_status(
                    &mut batch,
                    &db.filter,
                    host,
                    crate::pds_meta::HostStatus::Banned,
                )?;
                batch.commit().into_diagnostic()?;
                db.persist()
            })
            .await?;
        assert!(
            !hydrant.pds.is_banned(host),
            "the durable check is independent of the RCU snapshot"
        );

        assert_eq!(
            hydrant.firehose.add_source_validated(url.clone()).await?,
            Err(PdsAdmissionRejection::HostBanned)
        );
        assert!(!hydrant.firehose.is_source_known(&url));
        assert!(!hydrant.firehose.is_source_running(&url));
        assert!(persisted_sources(&hydrant).await?.is_empty());
        assert_eq!(crate::db::load_pds_daily_adds(&hydrant.state.db)?, None);
        Ok(())
    }

    #[tokio::test]
    async fn inflight_admission_starts_before_its_ban_is_published() -> Result<()> {
        let dir = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(&dir)).await?;
        startable(&hydrant);
        let host = "inflight-ban.example";
        let url = Url::parse("wss://inflight-ban.example/").into_diagnostic()?;
        let (checked, resume) = hydrant.firehose.pause_validated_admission_after_ban_check();

        let admission = tokio::spawn({
            let firehose = hydrant.firehose.clone();
            let url = url.clone();
            async move { firehose.add_source_validated(url).await }
        });
        checked.notified().await;

        let mut ban = tokio::spawn({
            let pds = hydrant.pds.clone();
            async move { pds.ban(host).await }
        });
        assert!(
            tokio::time::timeout(Duration::from_millis(100), &mut ban)
                .await
                .is_err(),
            "the ban cannot publish while admission owns its lock"
        );

        resume.notify_one();
        assert_eq!(
            admission
                .await
                .map_err(|error| miette::miette!("admission task failed: {error}"))??,
            Ok(PdsAdmission::StartedNew)
        );
        assert!(hydrant.firehose.is_source_running(&url));

        ban.await
            .map_err(|error| miette::miette!("ban task failed: {error}"))??;
        assert!(hydrant.pds.is_banned(host));
        hydrant.firehose.tasks.remove_async(&url).await;
        Ok(())
    }
}