use std::time::Duration;

use miette::{IntoDiagnostic, Result};
use smol_str::SmolStr;
use tokio::task::JoinSet;
use tracing::{error, info, warn};
use url::Url;

use super::pds_repos::{PdsListing, PdsRepoLister};
use super::worker::CursorUpdate;
use crate::control::new_pds_crawl;
use crate::db::keys;
use crate::net::{PublicHost, PublicHttpClient};
use crate::util::WatchEnabledExt;

const MAX_CONCURRENT_CRAWLS: usize = 4;
const RETRY_DELAY: Duration = Duration::from_secs(5 * 60);

/// runs the one-off `listRepos` walk queued for each PDS we start following,
/// see [`new_pds_crawl`].
#[derive(Clone)]
pub(crate) struct NewPdsCrawler {
    pub(crate) lister: PdsRepoLister,
    pub(crate) http: PublicHttpClient,
}

impl NewPdsCrawler {
    pub(crate) async fn run(mut self) {
        loop {
            self.lister
                .enabled
                .wait_enabled(self.lister.component)
                .await;
            let all_done = self.crawl_queued().await.unwrap_or_else(|err| {
                error!(err = ?err, "failed to read queued new PDS crawls");
                false
            });

            // a wake that lands mid round is kept by the notify, so the next round sees its host
            let state = self.lister.checker.state.clone();
            let woken = state.new_pds_crawl_notify.notified();
            if all_done {
                woken.await;
            } else {
                tokio::select! {
                    _ = woken => {}
                    _ = tokio::time::sleep(RETRY_DELAY) => {}
                }
            }
        }
    }

    /// returns whether every queued host got through.
    async fn crawl_queued(&self) -> Result<bool> {
        let queued = self.lister.checker.state.db.run(load_queued).await?;
        let mut all_done = true;
        let mut crawls = JoinSet::new();
        for (host, cursor) in queued {
            if crawls.len() >= MAX_CONCURRENT_CRAWLS
                && let Some(result) = crawls.join_next().await
            {
                all_done &= crawl_succeeded(result);
            }
            let crawler = self.clone();
            crawls.spawn(async move {
                let result = crawler.crawl(&host, cursor).await;
                (host, result)
            });
        }
        while let Some(result) = crawls.join_next().await {
            all_done &= crawl_succeeded(result);
        }
        Ok(all_done)
    }

    async fn crawl(self, host: &str, cursor: Option<SmolStr>) -> Result<()> {
        let key = keys::new_pds_crawl_key(host);
        let public = match PublicHost::parse(host) {
            Ok(public) => public,
            Err(error) => {
                warn!(%host, %error, "queued PDS is not a public host, dropping its crawl");
                return self.forget(key).await;
            }
        };
        // an operator removing the source means they don't want this host either
        let source_key = keys::firehose_source_key(public.firehose_url().as_str());
        let followed = self
            .lister
            .checker
            .state
            .db
            .run(move |db| db.crawler.contains_key(source_key).into_diagnostic())
            .await?;
        if !followed {
            info!(%host, "PDS is no longer followed, dropping its crawl");
            return self.forget(key).await;
        }

        info!(%host, resuming = cursor.is_some(), "listing repos of a newly followed PDS");
        let pds = public.https_endpoint().url().clone();
        self.list_host(host, key, &pds, cursor).await
    }

    async fn list_host(
        mut self,
        host: &str,
        key: Vec<u8>,
        pds: &Url,
        cursor: Option<SmolStr>,
    ) -> Result<()> {
        let http = self.http.clone();
        // we don't know its size, and one account paces it at the tier's base rate
        let listing = self
            .lister
            .list(pds, 1, &http, cursor, |next| {
                Ok(vec![CursorUpdate {
                    key: key.clone(),
                    // the last page deletes the key in the same commit as its repos
                    value: next
                        .map(|next| new_pds_crawl::encode_cursor(Some(next)))
                        .transpose()?,
                }])
            })
            .await?;
        match listing {
            PdsListing::Complete => info!(%host, "listed repos of a newly followed PDS"),
            PdsListing::Banned => {
                info!(%host, "PDS was banned, dropping its crawl");
                self.forget(key).await?;
            }
        }
        Ok(())
    }

    async fn forget(&self, key: Vec<u8>) -> Result<()> {
        self.lister
            .send_cursor_updates(vec![CursorUpdate { key, value: None }])
            .await
    }
}

fn load_queued(db: &crate::db::Db) -> Result<Vec<(String, Option<SmolStr>)>> {
    db.cursors
        .prefix(keys::NEW_PDS_CRAWL_PREFIX)
        .map(|entry| {
            let (key, value) = entry.into_inner().into_diagnostic()?;
            let host = std::str::from_utf8(&key[keys::NEW_PDS_CRAWL_PREFIX.len()..])
                .into_diagnostic()?
                .to_owned();
            Ok((host, new_pds_crawl::decode_cursor(&value)?))
        })
        .collect()
}

fn crawl_succeeded(
    result: std::result::Result<(String, Result<()>), tokio::task::JoinError>,
) -> bool {
    match result {
        Ok((_, Ok(()))) => true,
        Ok((host, Err(err))) => {
            warn!(%host, err = ?err, "new PDS crawl failed, retrying later");
            false
        }
        Err(err) => {
            error!(err = ?err, "new PDS crawl task failed");
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;
    use crate::crawler::{CrawlerStats, CrawlerWorker, InFlight, SignalChecker};
    use crate::state::AppState;
    use axum::extract::{Query, State};
    use axum::response::{IntoResponse, Response};
    use axum::{Json, Router, http::StatusCode, routing::get};
    use jacquard_common::types::string::Did;
    use serde_json::json;
    use std::collections::HashMap;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    use tempfile::tempdir;

    const HEAD: &str = "bafyreidf747c4x3lps3k4n357l3a3r57k3k465743k573k465743k5";

    struct MockPds {
        requests: AtomicUsize,
        fail_second_page: AtomicBool,
    }

    /// serves two `listRepos` pages, one unknown repo on each.
    async fn mock_pds() -> Result<(std::net::SocketAddr, Arc<MockPds>)> {
        async fn list_repos(
            State(pds): State<Arc<MockPds>>,
            Query(query): Query<HashMap<String, String>>,
        ) -> Response {
            pds.requests.fetch_add(1, Ordering::SeqCst);
            if !query.contains_key("cursor") {
                return Json(json!({
                    "cursor": "next",
                    "repos": [{"did": "did:web:first.example", "head": HEAD, "rev": "3j6s746574657"}]
                }))
                .into_response();
            }
            if pds.fail_second_page.load(Ordering::SeqCst) {
                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
            }
            Json(json!({
                "repos": [{"did": "did:web:second.example", "head": HEAD, "rev": "3j6s746574657"}]
            }))
            .into_response()
        }

        let pds = Arc::new(MockPds {
            requests: AtomicUsize::new(0),
            fail_second_page: AtomicBool::new(false),
        });
        let app = Router::new()
            .route("/xrpc/com.atproto.sync.listRepos", get(list_repos))
            .with_state(pds.clone());
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .into_diagnostic()?;
        let addr = listener.local_addr().into_diagnostic()?;
        tokio::spawn(async move { axum::serve(listener, app).await });
        Ok((addr, pds))
    }

    /// a crawler whose `pds.example` points at `addr`, with a running worker.
    fn crawler(state: &Arc<AppState>, addr: std::net::SocketAddr) -> Result<NewPdsCrawler> {
        let stats = CrawlerStats::new(state.clone(), Vec::new(), state.throttler.clone());
        let (worker, tx) = CrawlerWorker::new(state.clone(), 100, 50, stats.clone());
        tokio::spawn(worker.run());
        Ok(NewPdsCrawler {
            lister: PdsRepoLister {
                checker: SignalChecker {
                    http: PublicHttpClient::test(reqwest::Client::new()),
                    state: state.clone(),
                    throttler: state.throttler.clone(),
                },
                in_flight: InFlight::new(),
                tx,
                enabled: state.crawler_enabled.subscribe(),
                stats,
                component: "new PDS crawler",
            },
            http: PublicHttpClient::test(
                reqwest::Client::builder()
                    .resolve("pds.example", addr)
                    .build()
                    .into_diagnostic()?,
            ),
        })
    }

    fn test_state(dir: &tempfile::TempDir) -> Result<Arc<AppState>> {
        Ok(Arc::new(AppState::new(&Config {
            database_path: dir.path().to_path_buf(),
            full_network: true,
            enable_crawler: Some(true),
            ..Default::default()
        })?))
    }

    fn queue(state: &AppState, host: &str, cursor: Option<&str>) -> Result<()> {
        let cursor = cursor.map(SmolStr::new);
        state
            .db
            .cursors
            .insert(
                keys::new_pds_crawl_key(host),
                new_pds_crawl::encode_cursor(cursor.as_ref())?,
            )
            .into_diagnostic()
    }

    fn tracked(state: &AppState, did: &'static str) -> Result<bool> {
        state
            .db
            .repos
            .contains_key(keys::repo_key(&Did::new_static(did)?))
            .into_diagnostic()
    }

    /// the worker commits batches on its own task, so wait for it to catch up.
    async fn eventually(check: impl Fn() -> Result<bool>) -> Result<()> {
        tokio::time::timeout(Duration::from_secs(5), async {
            while !check()? {
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
            Ok(())
        })
        .await
        .into_diagnostic()?
    }

    #[tokio::test]
    async fn a_broken_walk_keeps_its_cursor_and_the_last_page_clears_it() -> Result<()> {
        let dir = tempdir().into_diagnostic()?;
        let state = test_state(&dir)?;
        let (addr, pds) = mock_pds().await?;
        let crawler = crawler(&state, addr)?;
        let url = Url::parse("http://pds.example/").into_diagnostic()?;
        let key = keys::new_pds_crawl_key("pds.example");
        queue(&state, "pds.example", None)?;

        pds.fail_second_page.store(true, Ordering::SeqCst);
        assert!(
            crawler
                .clone()
                .list_host("pds.example", key.clone(), &url, None)
                .await
                .is_err()
        );
        eventually(|| {
            Ok(tracked(&state, "did:web:first.example")?
                && load_queued(&state.db)? == vec![("pds.example".into(), Some("next".into()))])
        })
        .await?;

        // the 500 put the PDS in backoff, which the producer waits out between rounds
        pds.fail_second_page.store(false, Ordering::SeqCst);
        state.throttler.get_handle(&url).await.record_success();
        let [(host, cursor)] = <[_; 1]>::try_from(load_queued(&state.db)?).unwrap();
        crawler.list_host(&host, key, &url, cursor).await?;
        eventually(|| {
            Ok(tracked(&state, "did:web:second.example")? && load_queued(&state.db)?.is_empty())
        })
        .await?;
        // first page, the failed second page, then the second page again
        assert_eq!(pds.requests.load(Ordering::SeqCst), 3);
        Ok(())
    }

    #[tokio::test]
    async fn banned_unfollowed_and_private_hosts_drop_their_crawl() -> Result<()> {
        let dir = tempdir().into_diagnostic()?;
        let state = test_state(&dir)?;
        let (addr, pds) = mock_pds().await?;
        let crawler = crawler(&state, addr)?;

        queue(&state, "pds.example", None)?;
        state.pds_meta.rcu(|meta| {
            let mut next = (**meta).clone();
            next.update_host_entry("pds.example", |entry| {
                entry.status = crate::pds_meta::HostStatus::Banned;
            });
            next
        });
        let url = Url::parse("http://pds.example/").into_diagnostic()?;
        crawler
            .clone()
            .list_host(
                "pds.example",
                keys::new_pds_crawl_key("pds.example"),
                &url,
                None,
            )
            .await?;

        queue(&state, "gone.example", None)?;
        crawler.clone().crawl("gone.example", None).await?;
        queue(&state, "127.0.0.1", None)?;
        crawler.crawl("127.0.0.1", None).await?;

        eventually(|| Ok(load_queued(&state.db)?.is_empty())).await?;
        assert_eq!(pds.requests.load(Ordering::SeqCst), 0);
        Ok(())
    }
}