use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use miette::{IntoDiagnostic, Result};
use rand::RngExt;
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::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 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,
    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>>,
    /// ids assigned to spawned tasks
    next_task_id: Arc<AtomicUsize>,
}

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()),
            next_task_id: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub(super) async fn spawn_firehose_ingestor(
        &self,
        source: &FirehoseSource,
        shared: &FirehoseShared,
        delay_startup: bool,
    ) -> Result<()> {
        use std::sync::atomic::AtomicI64;
        let state = &self.state;

        let start = db::get_firehose_cursor(&state.db, &source.url).await?;
        // insert into relay_cursors if not already present; existing in-memory cursor takes precedence
        let _ = 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 enabled = state.firehose_enabled.subscribe();
        let ingestor = FirehoseIngestor::new(
            state.clone(),
            shared.buffer_tx.clone(),
            source.url.clone(),
            source.is_pds,
            state.filter.clone(),
            enabled,
            shared.verify_signatures,
            shared.max_failures,
        )
        .await;

        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 we dont cause thundering herd problems
                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 could 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 };
        self.tasks.upsert_async(source.url.clone(), handle).await;

        Ok(())
    }

    /// 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,
                    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 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 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());
        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.commit().into_diagnostic()?;
                db.persist()
            })
            .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(())
    }

    /// add PDS sources discovered from a seed relay.
    ///
    /// seed pages can contain thousands of hosts. persist the whole page in one
    /// batch and stagger startup so a fresh relay does not stampede DNS, TCP,
    /// TLS, and remote PDS websocket endpoints at once.
    pub(super) async fn add_seeded_sources(&self, urls: Vec<Url>) -> Result<usize> {
        let shared = self
            .shared
            .get()
            .ok_or_else(|| miette::miette!("firehose worker not started"))?;

        let mut sources = Vec::with_capacity(urls.len());
        for url in urls {
            if self.is_source_known(&url) {
                continue;
            }
            sources.push(FirehoseSource { url, is_pds: true });
        }

        if sources.is_empty() {
            return Ok(0);
        }

        self.state
            .db
            .run({
                let sources = sources.clone();
                move |db| {
                    let mut batch = db.inner.batch();
                    for source in &sources {
                        let value = rmp_serde::to_vec(&db::FirehoseSourceMeta {
                            is_pds: source.is_pds,
                        })
                        .map_err(|e| {
                            miette::miette!("failed to serialize firehose source meta: {e}")
                        })?;
                        batch.insert(
                            &db.crawler,
                            keys::firehose_source_key(source.url.as_str()),
                            &value,
                        );
                    }
                    batch.commit().into_diagnostic()?;
                    db.persist()
                }
            })
            .await?;

        let mut added = 0usize;
        for source in sources {
            if self
                .known_sources
                .insert_async(source.url.clone(), true)
                .await
                .is_ok()
            {
                self.spawn_firehose_ingestor(&source, shared, true).await?;
                added += 1;
            }
        }

        Ok(added)
    }

    /// 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> {
        if self.known_sources.contains_async(url).await {
            let url_str = url.to_string();
            self.state
                .db
                .run(move |db| {
                    db.crawler
                        .remove(keys::firehose_source_key(&url_str))
                        .into_diagnostic()?;
                    db.persist()
                })
                .await?;
            self.known_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, is_pds: bool) -> Result<()> {
        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(())
    }
}