use crate::types::{RepoMetadata, RepoState};

use fjall::{Database, Keyspace, PersistMode};
use miette::{Context, IntoDiagnostic, Result};
use scc::HashMap;
use smol_str::SmolStr;

use std::collections::BTreeSet;

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

pub mod compaction;
pub mod counts;
#[cfg(feature = "indexer")]
use counts::CountDeltaReservation;
pub use counts::{CountDeltas, load_count_delta_watermark, set_ks_count};
pub mod ephemeral;
pub mod filter;
pub mod keys;
#[cfg(feature = "indexer")]
mod lifecycle_counts;
pub mod migration;
pub mod pds_meta;
#[cfg(feature = "indexer")]
pub(crate) mod storage_mode;
pub mod types;

pub mod keyspaces;
mod open;
pub mod registry;
pub mod schema;
mod train;
mod txn;

#[cfg(feature = "indexer")]
pub use keyspaces::IndexerDb;
#[cfg(feature = "jetstream")]
pub(crate) use keyspaces::JetstreamDb;
#[cfg(feature = "relay")]
pub(crate) use keyspaces::RelayDb;
#[cfg(feature = "indexer_stream")]
pub use keyspaces::StreamDb;

pub use schema::Ks;

#[cfg(feature = "indexer")]
use lifecycle_counts::LifecycleCountBatch;
pub(crate) use txn::{Txn, TxnCommitTimings};

use tracing::error;

#[derive(Clone)]
pub struct Db {
    pub inner: Arc<Database>,
    pub path: std::path::PathBuf,
    pub repos: Ks<schema::Repos>,
    pub repo_metadata: Ks<schema::RepoMetadata>,
    pub cursors: Ks<schema::Cursors>,
    pub counts: Ks<schema::Counts>,
    pub filter: Ks<schema::Filter>,
    pub crawler: Ks<schema::Crawler>,
    #[cfg(feature = "indexer")]
    pub indexer: IndexerDb,
    #[cfg(feature = "indexer_stream")]
    pub stream: StreamDb,
    #[cfg(feature = "jetstream")]
    pub(crate) jetstream: JetstreamDb,
    #[cfg(feature = "relay")]
    pub(crate) relay: RelayDb,
    #[cfg(feature = "backlinks")]
    pub(crate) backlinks: Ks<schema::Backlinks>,
    pub counts_map: Arc<HashMap<SmolStr, u64>>,
    next_count_delta_id: Arc<AtomicU64>,
    count_delta_checkpoint_watermark: Arc<AtomicU64>,
    count_delta_gc_watermark: Arc<AtomicU64>,
    count_delta_in_flight: Arc<Mutex<BTreeSet<u64>>>,
    pub(crate) compaction_running: Arc<std::sync::atomic::AtomicBool>,
    #[cfg(test)]
    pub(crate) persist_failures: Arc<AtomicUsize>,
    /// 256 lock-sharded mutexes keyed by one byte of the trimmed DID. this is
    /// synchronization only, not storage partitioning: unrelated DIDs may
    /// conservatively share a lock. record mutations hold the selected lock
    /// from read through commit so operator redaction cannot race them.
    #[cfg(feature = "indexer")]
    pub(crate) record_write_locks: Arc<[std::sync::Mutex<()>; 256]>,
}

/// lock index for a DID. skip the method tag so all PLC DIDs do not collapse
/// onto one mutex.
#[cfg(feature = "indexer")]
pub(crate) fn record_lock_index_for_did(did: &jacquard_common::types::did::Did) -> u8 {
    let prefix = keys::record_prefix_did(did);
    match prefix.first() {
        Some(&types::TAG_PLC) | Some(&types::TAG_WEB) => prefix.get(1).copied().unwrap_or(0),
        Some(&b) => b,
        None => 0,
    }
}

impl Db {
    /// runs synchronous database work without blocking the async runtime.
    pub async fn run<T, F>(&self, f: F) -> Result<T>
    where
        T: Send + 'static,
        F: FnOnce(&Db) -> Result<T> + Send + 'static,
    {
        let db = self.clone();
        tokio::task::spawn_blocking(move || f(&db))
            .await
            .into_diagnostic()?
    }

    pub fn persist(&self) -> Result<()> {
        #[cfg(test)]
        if self
            .persist_failures
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
                if remaining > 0 {
                    Some(remaining - 1)
                } else {
                    None
                }
            })
            .is_ok()
        {
            return Err(miette::miette!("injected persist failure"));
        }

        #[cfg(not(feature = "__persist_sync_all"))]
        const MODE: PersistMode = PersistMode::Buffer;
        #[cfg(feature = "__persist_sync_all")]
        const MODE: PersistMode = PersistMode::SyncAll;
        self.inner.persist(MODE).into_diagnostic()?;
        Ok(())
    }

    #[cfg(test)]
    pub(crate) fn fail_next_persist_for_test(&self) {
        self.persist_failures.fetch_add(1, Ordering::Release);
    }

    pub async fn compact(&self) -> Result<()> {
        use std::sync::atomic::Ordering;
        if self
            .compaction_running
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            return Err(miette::miette!("compaction already in progress"));
        }
        struct Guard(Arc<std::sync::atomic::AtomicBool>);
        impl Drop for Guard {
            fn drop(&mut self) {
                self.0.store(false, Ordering::Release);
            }
        }
        let _guard = Guard(self.compaction_running.clone());

        let compact = |name: &'static str, ks: Keyspace| async move {
            tracing::info!("compacting keyspace {name}...");
            tokio::task::spawn_blocking(move || ks.major_compact().into_diagnostic())
                .await
                .into_diagnostic()??;
            tracing::info!("compacted keyspace {name}");
            Ok::<(), miette::Report>(())
        };

        let tasks: Vec<_> = self
            .all_keyspaces()
            .into_iter()
            .map(|(name, ks)| compact(name, ks))
            .collect();

        futures::future::try_join_all(tasks).await?;

        Ok(())
    }
}

#[cfg(feature = "indexer")]
mod indexer;

#[cfg(feature = "indexer")]
pub use indexer::*;

#[derive(serde::Serialize, serde::Deserialize, Default)]
pub(crate) struct FirehoseSourceMeta {
    #[serde(default)]
    pub(crate) is_pds: bool,
}

pub fn set_firehose_cursor(db: &Db, relay: &Url, cursor: i64) -> Result<()> {
    db.cursors
        .insert(
            keys::firehose_cursor_key_from_url(relay),
            cursor.to_be_bytes(),
        )
        .into_diagnostic()
}

pub async fn get_firehose_cursor(db: &Db, relay: &Url) -> Result<Option<i64>> {
    let key = keys::firehose_cursor_key_from_url(relay);
    db.run(move |db| db.cursors.get(key).into_diagnostic())
        .await?
        .map(|v| {
            Ok(i64::from_be_bytes(
                v.as_ref()
                    .try_into()
                    .into_diagnostic()
                    .wrap_err("cursor is not 8 bytes")?,
            ))
        })
        .transpose()
}

#[cfg(feature = "indexer")]
pub fn ser_repo_meta(state: &RepoMetadata) -> Result<Vec<u8>> {
    rmp_serde::to_vec(&state).into_diagnostic()
}

pub fn deser_repo_meta(bytes: &[u8]) -> Result<RepoMetadata> {
    rmp_serde::from_slice(bytes).into_diagnostic()
}

pub fn ser_repo_state(state: &RepoState) -> Result<Vec<u8>> {
    rmp_serde::to_vec(&state).into_diagnostic()
}

pub fn deser_repo_state<'b>(bytes: &'b [u8]) -> Result<RepoState<'b>> {
    rmp_serde::from_slice(bytes).into_diagnostic()
}

pub fn check_poisoned(e: &fjall::Error) {
    if matches!(e, fjall::Error::Poisoned) {
        error!("!!! DATABASE POISONED !!! exiting");
        std::process::exit(10);
    }
}

pub fn check_poisoned_report(e: &miette::Report) {
    let Some(err) = e.downcast_ref::<fjall::Error>() else {
        return;
    };
    self::check_poisoned(err);
}

/// Load the persisted `(day, count)` pair for the daily PDS add counter.
/// Returns `None` if no entry exists or the stored data is malformed.
pub fn load_pds_daily_adds(db: &Db) -> Result<Option<(u64, u64)>> {
    let Some(val) = db.cursors.get(keys::PDS_DAILY_ADDS_KEY).into_diagnostic()? else {
        return Ok(None);
    };
    if val.len() < 16 {
        miette::bail!("malformed pds daily limit value");
    }
    let day = u64::from_be_bytes(val[..8].try_into().into_diagnostic()?);
    let count = u64::from_be_bytes(val[8..].try_into().into_diagnostic()?);
    Ok(Some((day, count)))
}

/// Stage the daily PDS add counter alongside the source it admits.
pub(crate) fn stage_pds_daily_adds(
    batch: &mut fjall::OwnedWriteBatch,
    db: &Db,
    day: u64,
    count: u64,
) {
    let mut value = [0u8; 16];
    value[..8].copy_from_slice(&day.to_be_bytes());
    value[8..].copy_from_slice(&count.to_be_bytes());
    batch.insert(&db.cursors, keys::PDS_DAILY_ADDS_KEY, value);
}

/// Loads persisted sources and their explicit trusted-admin provenance bit.
/// A missing marker is deliberately public, fail-closed on legacy rows.
pub fn load_persisted_firehose_sources(
    db: &crate::db::Db,
) -> Result<Vec<(crate::config::FirehoseSource, bool)>> {
    use crate::db::keys::{FIREHOSE_SOURCE_PREFIX, trusted_firehose_source_key};

    let mut sources = Vec::new();
    for entry in db.crawler.prefix(FIREHOSE_SOURCE_PREFIX) {
        let (key, val) = entry.into_inner().into_diagnostic()?;
        let url_bytes = &key[FIREHOSE_SOURCE_PREFIX.len()..];
        let url_str = std::str::from_utf8(url_bytes).into_diagnostic()?;
        let url = Url::parse(url_str).into_diagnostic()?;
        let meta: FirehoseSourceMeta = rmp_serde::from_slice(&val)
            .map_err(|e| miette::miette!("failed to deserialize firehose source meta: {e}"))?;
        let trusted = db
            .crawler
            .contains_key(trusted_firehose_source_key(url_str))
            .into_diagnostic()?;
        sources.push((
            crate::config::FirehoseSource {
                url,
                is_pds: meta.is_pds,
            },
            trusted,
        ));
    }
    Ok(sources)
}

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

    #[tokio::test]
    async fn test_db_compact_concurrency_guard() -> Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let cfg = Config {
            database_path: tmp.path().to_path_buf(),
            ..Config::default()
        };

        let db = Db::open(&cfg)?;

        // Manually mark compaction as running
        db.compaction_running.store(true, Ordering::SeqCst);

        // Attempting to compact should now fail
        let res = db.compact().await;
        assert!(res.is_err());
        assert!(res.unwrap_err().to_string().contains("already in progress"));

        // Release the lock
        db.compaction_running.store(false, Ordering::SeqCst);

        // Compacting should now succeed
        let res = db.compact().await;
        assert!(res.is_ok());

        Ok(())
    }
    #[tokio::test]
    async fn legacy_firehose_row_without_trusted_marker_is_public() -> Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let config = Config {
            database_path: tmp.path().to_path_buf(),
            ..Config::default()
        };
        let db = Db::open(&config)?;
        let url = "wss://legacy-pds.example/";
        let value = rmp_serde::to_vec(&FirehoseSourceMeta { is_pds: true }).into_diagnostic()?;
        db.crawler
            .insert(keys::firehose_source_key(url), value)
            .into_diagnostic()?;
        let sources = load_persisted_firehose_sources(&db)?;
        assert_eq!(sources.len(), 1);
        assert_eq!(sources[0].0.url, Url::parse(url).unwrap());
        assert!(!sources[0].1, "missing marker must fail closed to public");
        Ok(())
    }

    #[test]
    fn firehose_source_meta_golden_fixture() {
        let encoded = rmp_serde::to_vec(&FirehoseSourceMeta { is_pds: true }).unwrap();
        assert_eq!(hex::encode(&encoded), "91c3");
        let decoded: FirehoseSourceMeta = rmp_serde::from_slice(&encoded).unwrap();
        assert!(decoded.is_pds);

        let legacy_default: FirehoseSourceMeta = rmp_serde::from_slice(&[0x90]).unwrap();
        assert!(!legacy_default.is_pds);
    }

    #[tokio::test]
    async fn firehose_cursor_big_endian_roundtrip() -> Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let cfg = Config {
            database_path: tmp.path().to_path_buf(),
            ..Config::default()
        };
        let db = Db::open(&cfg)?;
        let relay = Url::parse("https://relay.example").into_diagnostic()?;
        let cursor = i64::from_be_bytes(*b"stinkpot");

        set_firehose_cursor(&db, &relay, cursor)?;
        let raw = db
            .cursors
            .get(keys::firehose_cursor_key_from_url(&relay))
            .into_diagnostic()?
            .unwrap();
        assert_eq!(raw.as_ref(), b"stinkpot");
        assert_eq!(get_firehose_cursor(&db, &relay).await?, Some(cursor));
        Ok(())
    }
}