use anyhow::{anyhow, Result};
use fts_index::index::{apply_deleteds, IndexMeta};
use fts_index::protocol::{Query, SearchHit};
use fts_index::schema::{DOC_ID_FIELD, VERSION_FIELD};
use fts_index::storage::incremental::gc_orphan_index_dirs;
use fts_index::storage::meta::StorageMeta;
use log::*;
use std::collections::HashMap;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::Arc;
use std::thread::sleep;
use tantivy::collector::TopDocs;
use tantivy::schema::Value;
use tantivy::TantivyDocument;
use tokio::sync::{Mutex, RwLock};

pub struct IndexSearcher {
    meta: IndexMeta,
    searcher: tantivy::Searcher,
}

impl IndexSearcher {
    async fn new(meta: IndexMeta) -> Result<Self> {
        let searcher = meta.searcher().await?;
        Ok(Self { meta, searcher })
    }

    fn search(&self, query: &Query, limit: u32) -> Result<Vec<SearchHit>> {
        let tantivy_query: Box<dyn tantivy::query::Query> = query.into();
        info!(
            "Search index with query: {:?}",
            serde_json::to_string(query)
        );
        let result: Vec<SearchHit> = self
            .searcher
            .search(&tantivy_query, &TopDocs::with_limit(limit as usize))?
            .into_iter()
            .flat_map(|(score, doc_address)| {
                self.searcher
                    .doc(doc_address)
                    .map_err(|e| anyhow!("Could not retrieve document: {}", e))
                    .and_then(|doc: TantivyDocument| {
                        let doc_id = doc
                            .get_first(*DOC_ID_FIELD)
                            .ok_or(anyhow!("Document ID not found in search result"))?;
                        let doc_id = doc_id
                            .as_str()
                            .ok_or(anyhow!("Document ID is not a string"))?
                            .into();
                        let version = doc
                            .get_first(*VERSION_FIELD)
                            .ok_or(anyhow!("Document version not found in search result"))?;
                        let version = version
                            .as_u64()
                            .ok_or(anyhow!("Document version is not a number"))?;

                        Ok(SearchHit {
                            doc_id,
                            version,
                            score,
                        })
                    })
            })
            .collect();

        // unique by SearchHit.doc_id
        let mut unique_result = HashMap::new();
        result.into_iter().for_each(|hit| {
            unique_result.insert(hit.doc_id.clone(), hit);
        });

        Ok(unique_result.into_values().collect())
    }
}

#[derive(Default)]
pub struct Searcher {
    indexes: RwLock<HashMap<String, Arc<IndexSearcher>>>,
    current_version: Mutex<u64>,
    storage_path: String,
    updating_index: AtomicBool,
}

impl Searcher {
    pub(crate) async fn new(
        storage_path: &str,
        cloud: Option<Arc<dyn cloud::Storage>>,
    ) -> Result<Self> {
        let searcher = Searcher {
            storage_path: storage_path.into(),
            ..Default::default()
        };

        match cloud {
            Some(c) => {
                if let Err(e) = searcher.init_cloud(c).await {
                    error!("Failed to initialize cloud storage: {}", e);
                }
            }
            None => searcher.init_local().await?,
        }

        Ok(searcher)
    }

    pub(crate) async fn version(&self) -> u64 {
        *self.current_version.lock().await
    }

    async fn init_local(&self) -> Result<()> {
        let mut version_lock = self.current_version.lock().await;
        let storage_meta = StorageMeta::load_latest_local(&self.storage_path).await?;
        debug!("Loaded local storage meta: {:?}", storage_meta);
        let new_version = storage_meta.version;
        *self.indexes.write().await = self.load_indexes(&storage_meta, false).await?;
        *version_lock = new_version;
        Ok(())
    }

    async fn load_local_by_storage_meta_explicit(&self, storage_meta: &StorageMeta) -> Result<()> {
        let mut version_lock = self.current_version.lock().await;
        let new_version = storage_meta.version;

        // explicit load indexes from storage_meta: fist-start & last-end
        *self.indexes.write().await = self.load_indexes(storage_meta, true).await?;
        debug!("Loaded indexes: {:?}", self.indexes.read().await.keys());
        *version_lock = new_version;
        Ok(())
    }

    async fn init_cloud(&self, cloud: Arc<dyn cloud::Storage>) -> Result<()> {
        let storage_meta;
        let mut retry = 5;
        let mut active_try = true;
        loop {
            // try load cloud version first
            if active_try {
                active_try = false;
                if let Ok(v) = StorageMeta::get_cloud_version(cloud.clone()).await {
                    storage_meta = if let Ok(m) =
                        StorageMeta::load_from_cloud_by_version(cloud.clone(), v).await
                    {
                        m
                    } else {
                        warn!(
                            "Failed to load storage meta by cloud version, {} retrying...",
                            retry
                        );
                        sleep(std::time::Duration::from_secs(10));
                        continue;
                    };
                    break;
                }
            }

            if retry == 0 {
                return Err(anyhow!("Failed to load storage meta from cloud"));
            }

            // load latest cloud version
            let cloud_versions = StorageMeta::get_cloud_versions(cloud.clone()).await?;

            let mut storage_meta_opt = None;
            for version in cloud_versions {
                match StorageMeta::load_from_cloud_by_version(cloud.clone(), version).await {
                    Ok(meta) => {
                        if meta.sum_incrementals() == 0 {
                            continue;
                        }
                        storage_meta_opt = Some(meta);
                        break;
                    }
                    Err(e) => {
                        warn!("Failed to load storage meta from cloud: {}", e);
                    }
                }
            }

            if let Some(meta) = storage_meta_opt {
                storage_meta = meta;
                break;
            }
            warn!(
                "Failed to load storage meta by latest version, {} retrying...",
                retry
            );
            sleep(std::time::Duration::from_secs(10));
            retry -= 1;
        }

        storage_meta.lock_cloud_version(cloud.clone(), true).await?;

        // stoage-meta is exists in local, so assume that the version is downloaded
        if StorageMeta::get_local_versions(&self.storage_path)
            .await?
            .contains(&storage_meta.version)
        {
            info!(
                "Storage-meta version {} is already downloaded [init_cloud].",
                storage_meta.version
            );
        } else {
            storage_meta
                .download_all_incrementals_and_merge(cloud.clone(), &self.storage_path)
                .await?;
            info!("Inited cloud storage-meta {}.", storage_meta.version);
        }

        apply_deleteds(&self.storage_path, cloud.clone(), &storage_meta).await?;
        self.load_local_by_storage_meta_explicit(&storage_meta)
            .await?;

        storage_meta.lock_cloud_version(cloud, false).await?;

        // a pod that restarts and never shifts version would otherwise never collect
        if let Err(e) = self.local_gc().await {
            warn!("Failed to run local GC on init: {}", e);
        }

        Ok(())
    }

    async fn load_indexes(
        &self,
        meta: &StorageMeta,
        explicit: bool,
        // compact: bool,
    ) -> Result<HashMap<String, Arc<IndexSearcher>>> {
        let mut indexes: HashMap<String, Arc<IndexSearcher>> = HashMap::new();
        let mut instance_count = 0;
        for (index_name, _) in meta.indexes.clone().into_iter() {
            let index_meta = if explicit {
                IndexMeta::load_local_by_explicit_versions(
                    &self.storage_path,
                    &index_name,
                    meta.first_start_version(&index_name),
                    meta.latest_end_version(&index_name),
                )
                .await?
            } else {
                IndexMeta::load_local(&self.storage_path, &index_name).await?
            };

            if let Some(index_meta) = index_meta {
                indexes.insert(
                    index_name.to_string(),
                    Arc::new(IndexSearcher::new(index_meta).await?),
                );
                instance_count += 1;
            }
        }

        if instance_count == 0 {
            return Err(anyhow!("No index instance loaded"));
        }

        Ok(indexes)
    }

    pub(crate) async fn get_versions(&self, index_name: &str) -> Result<(u64, u64)> {
        let indexes = self.indexes.read().await;
        if let Some(index_searcher) = indexes.get(index_name) {
            return Ok((
                index_searcher.meta.start_version(),
                index_searcher.meta.end_version(),
            ));
        }
        Ok((0, 0))
    }

    pub(crate) async fn search(
        &self,
        index_name: &str,
        query: &Query,
        limit: u32,
    ) -> Result<(Vec<SearchHit>, u64)> {
        if let Some(index_searcher) = self.indexes.read().await.get(index_name) {
            return Ok((
                index_searcher.search(query, limit)?,
                index_searcher.meta.end_version(),
            ));
        }
        Err(anyhow!("Index '{}' does not exist", index_name))
    }

    pub async fn shift_index_version(
        &self,
        storage: Arc<dyn cloud::Storage>,
        version: u64,
    ) -> Result<u64> {
        match self
            .updating_index
            .compare_exchange(false, true, SeqCst, SeqCst)
        {
            Ok(_) => {}
            Err(_) => return Err(anyhow!("Conflict update is running")),
        }

        let result = self.shift_index_version_inner(storage, version).await;
        self.updating_index.store(false, SeqCst);
        result
    }

    async fn shift_index_version_inner(
        &self,
        storage: Arc<dyn cloud::Storage>,
        version: u64,
    ) -> Result<u64> {
        if version == self.version().await {
            return Ok(version);
        }

        let storage_meta =
            StorageMeta::load_from_cloud_by_version(storage.clone(), version).await?;

        storage_meta
            .lock_cloud_version(storage.clone(), true)
            .await?;

        // stoage-meta is exists in local, so assume that the version is downloaded
        if StorageMeta::get_local_versions(&self.storage_path)
            .await?
            .contains(&storage_meta.version)
        {
            info!(
                "Storage-meta version {} is already downloaded [shift_index_version_inner].",
                storage_meta.version
            );
        } else {
            storage_meta
                .download_all_incrementals_and_merge(storage.clone(), &self.storage_path)
                .await?;
            info!(
                "Shifted storage-meta {} -> {}.",
                self.version().await,
                storage_meta.version
            );
        }

        apply_deleteds(&self.storage_path, storage.clone(), &storage_meta).await?;
        self.load_local_by_storage_meta_explicit(&storage_meta)
            .await?;

        storage_meta
            .lock_cloud_version(storage.clone(), false)
            .await?;

        match self.local_gc().await {
            Ok(_) => {}
            Err(e) => warn!("Failed to run local GC: {}", e),
        }

        Ok(version)
    }

    async fn local_gc(&self) -> Result<()> {
        let versions = StorageMeta::get_local_versions(&self.storage_path).await?;
        let current_version = self.version().await;
        debug!("Starting gc around version {}", current_version);
        for version in versions {
            // keep the live version and its neighbours; `current_version + 1` because
            // `version - 1` underflows at 0
            if version == current_version
                || version.abs_diff(current_version) <= 1
            {
                continue;
            }

            debug!("GCing version {}", version);
            // one unreadable or half-collected version must not stop the rest of the sweep
            match StorageMeta::load_local_by_version(&self.storage_path, version).await {
                Ok(storage_meta) => match storage_meta.local_gc(&self.storage_path).await {
                    Ok(_) => debug!("GC version {} done", version),
                    Err(e) => warn!("GC version {} incomplete: {}", version, e),
                },
                Err(e) => warn!("Failed to load local storage meta {}: {}", version, e),
            }
        }

        // dirs whose meta file was already lost are invisible to the loop above
        match StorageMeta::load_local_by_version(&self.storage_path, current_version).await {
            Ok(current_meta) => {
                for index_name in current_meta.indexes.keys() {
                    let keep = (
                        current_meta.first_start_version(index_name),
                        current_meta.latest_end_version(index_name),
                    );
                    match gc_orphan_index_dirs(&self.storage_path, index_name, keep).await {
                        Ok(n) if n > 0 => {
                            info!("Removed {} orphan index dir(s) for index {}", n, index_name)
                        }
                        Ok(_) => {}
                        Err(e) => warn!("Failed to sweep orphan dirs for {}: {}", index_name, e),
                    }
                }
            }
            Err(e) => warn!(
                "Skipped orphan sweep, cannot load current storage meta {}: {}",
                current_version, e
            ),
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use fts_index::storage::meta::{Incremental, Index};

    fn meta_json(version: u64, index_name: &str, start: u64, end: u64) -> String {
        let meta = StorageMeta {
            version,
            indexes: HashMap::from([(
                index_name.to_string(),
                Index::new(1.0f32.into(), vec![Incremental::new(start, end)]),
            )]),
        };
        serde_json::to_string(&meta).unwrap()
    }

    #[tokio::test]
    async fn test_local_gc_continues_past_unreadable_meta() {
        let storage_path = "/tmp/fts-searcher-gc-continue-test";
        let _ = tokio::fs::remove_dir_all(storage_path).await;
        tokio::fs::create_dir_all(format!("{}/idx/0-500", storage_path))
            .await
            .unwrap();

        // get_local_versions returns descending, so 9 is visited before 5
        tokio::fs::write(format!("{}/meta-9.json", storage_path), "{ not json")
            .await
            .unwrap();
        tokio::fs::write(
            format!("{}/meta-5.json", storage_path),
            meta_json(5, "idx", 0, 500),
        )
        .await
        .unwrap();
        tokio::fs::write(format!("{}/idx/meta-0-500.json", storage_path), "{}")
            .await
            .unwrap();

        let searcher = Searcher {
            storage_path: storage_path.to_string(),
            ..Default::default()
        };
        *searcher.current_version.lock().await = 20;

        assert!(searcher.local_gc().await.is_ok());
        assert!(
            !std::path::Path::new(&format!("{}/meta-5.json", storage_path)).exists(),
            "version 5 must still be collected after version 9 failed to load"
        );
        assert!(!std::path::Path::new(&format!("{}/idx/0-500", storage_path)).exists());
        assert!(std::path::Path::new(&format!("{}/meta-9.json", storage_path)).exists());

        let _ = tokio::fs::remove_dir_all(storage_path).await;
    }

    #[tokio::test]
    async fn test_local_gc_handles_version_zero() {
        let storage_path = "/tmp/fts-searcher-gc-zero-test";
        let _ = tokio::fs::remove_dir_all(storage_path).await;
        tokio::fs::create_dir_all(storage_path).await.unwrap();
        tokio::fs::write(
            format!("{}/meta-0.json", storage_path),
            meta_json(0, "idx", 0, 0),
        )
        .await
        .unwrap();

        let searcher = Searcher {
            storage_path: storage_path.to_string(),
            ..Default::default()
        };
        *searcher.current_version.lock().await = 5;

        // `version - 1` overflows here and panics under debug overflow checks
        assert!(searcher.local_gc().await.is_ok());
        assert!(!std::path::Path::new(&format!("{}/meta-0.json", storage_path)).exists());

        let _ = tokio::fs::remove_dir_all(storage_path).await;
    }

    #[tokio::test]
    async fn test_local_gc_sweeps_orphan_dirs() {
        let storage_path = "/tmp/fts-searcher-gc-sweep-test";
        let _ = tokio::fs::remove_dir_all(storage_path).await;

        // live version 20 serves range 0-300; 0-100 is an orphan with no meta file
        tokio::fs::create_dir_all(format!("{}/idx/0-300", storage_path))
            .await
            .unwrap();
        tokio::fs::create_dir_all(format!("{}/idx/0-100", storage_path))
            .await
            .unwrap();
        tokio::fs::write(
            format!("{}/meta-20.json", storage_path),
            meta_json(20, "idx", 0, 300),
        )
        .await
        .unwrap();
        tokio::fs::write(format!("{}/idx/meta-0-300.json", storage_path), "{}")
            .await
            .unwrap();
        // the immediate neighbours of the live version are kept, not collected
        tokio::fs::write(
            format!("{}/meta-19.json", storage_path),
            meta_json(19, "idx", 0, 200),
        )
        .await
        .unwrap();
        tokio::fs::write(
            format!("{}/meta-21.json", storage_path),
            meta_json(21, "idx", 0, 400),
        )
        .await
        .unwrap();

        let searcher = Searcher {
            storage_path: storage_path.to_string(),
            ..Default::default()
        };
        *searcher.current_version.lock().await = 20;

        assert!(searcher.local_gc().await.is_ok());
        assert!(
            !std::path::Path::new(&format!("{}/idx/0-100", storage_path)).exists(),
            "orphan dir must be reclaimed by the sweep"
        );
        assert!(std::path::Path::new(&format!("{}/idx/0-300", storage_path)).exists());
        assert!(std::path::Path::new(&format!("{}/meta-20.json", storage_path)).exists());
        assert!(
            std::path::Path::new(&format!("{}/meta-19.json", storage_path)).exists(),
            "version below the live one must be kept"
        );
        assert!(
            std::path::Path::new(&format!("{}/meta-21.json", storage_path)).exists(),
            "version above the live one must be kept"
        );

        let _ = tokio::fs::remove_dir_all(storage_path).await;
    }
}