use crate::protocol::{DocumentFull, DocumentFullVec};
use crate::schema::INDEX_SCHEMA;
use crate::storage::deletes::Deletes;
use crate::storage::incremental::IncrementalIndexMeta;
use crate::storage::meta::StorageMeta;
use crate::tokenizer::CharabiaTokenizer;
use anyhow::{anyhow, Result};
use lazy_static::lazy_static;
use log::*;
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::cmp::{max, Ordering};
use std::sync::Arc;
use std::{env, path::*};
use sysinfo::{MemoryRefreshKind, RefreshKind, System};
use tantivy::directory::MmapDirectory;
use tantivy::indexer::{NoMergePolicy, UserOperation};
use tantivy::{Directory, Index, IndexWriter, ReloadPolicy, Searcher, TantivyDocument, Term};
use tokio::fs::*;
use tokio::task::spawn_blocking;

lazy_static! {
    pub static ref INDEX_VERSION_META_FILE_PATTERN: Regex =
        Regex::new(r#"^meta-(?<start>\d+)-(?<end>\d+).json$"#).unwrap();
    pub static ref DEFAULT_INDEXING_MEMORY_BUDGET: usize = {
        let kind = RefreshKind::new().with_memory(MemoryRefreshKind::new().with_ram());
        let mut sys = System::new_with_specifics(kind);
        sys.refresh_specifics(kind);
        max(128 * 1024 * 1024, sys.total_memory() * 7 / 10) as usize
    };
    pub static ref FTS_MAX_SEGMENT_DOCS: u32 = {
        env::var("FTS_MAX_SEGMENT_DOCS")
            .ok()
            .and_then(|s| s.parse::<u32>().ok())
            .unwrap_or(50_000_000)
    };
}

const END_VERSION_CAPTURE_GROUP_NAME: &str = "end";
const START_VERSION_CAPTURE_GROUP_NAME: &str = "start";
// const TARGET_INDEX_INSTANCES: usize = 3;

#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub struct IndexMeta {
    storage_path: String,
    index_name: String,
    incremental_meta: IncrementalIndexMeta,
}

pub async fn remove_file_ignore_error(path: &Path) {
    if let Err(e) = tokio::fs::remove_file(path).await {
        if e.kind() != std::io::ErrorKind::NotFound {
            warn!("Failed to remove file '{}': {}", path.display(), e)
        }
    }
}

impl IndexMeta {
    pub fn current(&self) -> &IncrementalIndexMeta {
        &self.incremental_meta
    }

    pub fn end_version(&self) -> u64 {
        self.incremental_meta.end_version()
    }

    pub fn start_version(&self) -> u64 {
        self.incremental_meta.start_version()
    }

    pub fn incremental_meta(&self) -> IncrementalIndexMeta {
        self.incremental_meta.clone()
    }

    pub async fn new(
        storage_path: &str,
        index_name: &str,
        start_version: u64,
        end_version: u64,
    ) -> Result<Self> {
        debug!(
            "Creating IndexMeta with storage_path: {}, index_name: {}, start_version: {}, end_version: {}",
            storage_path, index_name, start_version, end_version,
        );
        let mut latest_instance = IncrementalIndexMeta::default();
        latest_instance.set_meta(
            index_name.to_string(),
            start_version,
            end_version,
            latest_instance.schema_version().clone(),
        );
        let meta = IndexMeta {
            storage_path: storage_path.into(),
            incremental_meta: latest_instance.clone(),
            index_name: index_name.into(),
        };
        debug!(
            "Creating index meta dir in IndexMeta: {:?}, {:?}",
            meta,
            meta.index_path()
        );
        match create_dir_all(meta.index_path()).await {
            Ok(_) => {}
            Err(e) => {
                if e.kind() != std::io::ErrorKind::AlreadyExists {
                    return Err(anyhow::anyhow!("Failed to create directory: {}", e));
                }
            }
        };
        Ok(meta)
    }

    pub async fn load_local_by_explicit_versions(
        storage_path: &str,
        index_name: &str,
        start_version: u64,
        end_version: u64,
    ) -> Result<Option<Self>> {
        let mut has_index_dir = false;
        let mut has_existing_meta = false;
        let mut current = IncrementalIndexMeta::default();
        let mut files = read_dir(format!("{}/{}", storage_path, index_name)).await?;
        while let Some(entry) = files.next_entry().await? {
            if entry.file_type().await?.is_dir() {
                if entry.file_name().to_str().unwrap()
                    == format!("{}-{}", start_version, end_version)
                {
                    has_index_dir = true;
                }
                continue;
            }

            debug!("Loading local index meta: {:?}", entry.path());
            if let Some(file_name) = entry.file_name().to_str() {
                if let Some(caps) = INDEX_VERSION_META_FILE_PATTERN.captures(file_name) {
                    if let (Some(start_str), Some(end_str)) = (
                        caps.name(START_VERSION_CAPTURE_GROUP_NAME),
                        caps.name(END_VERSION_CAPTURE_GROUP_NAME),
                    ) {
                        if start_str.as_str().parse::<u64>()? == start_version
                            && end_str.as_str().parse::<u64>()? == end_version
                        {
                            current = serde_json::from_str::<IncrementalIndexMeta>(
                                &read_to_string(entry.path()).await?,
                            )?;
                            current.set_meta(
                                index_name.to_string(),
                                start_version,
                                end_version,
                                current.schema_version().clone(),
                            );
                            has_existing_meta = true;
                            continue;
                        }
                    }
                }
            }
        }

        if !has_existing_meta {
            return Ok(None);
        }

        if !has_index_dir {
            return Err(anyhow!(
                "Explicit index directory {}-{} not found in {}",
                start_version,
                end_version,
                index_name
            ));
        }

        Ok(Some(Self {
            storage_path: storage_path.into(),
            incremental_meta: current,
            index_name: index_name.into(),
        }))
    }

    pub async fn load_local(storage_path: &str, index_name: &str) -> Result<Option<Self>> {
        let mut current = IncrementalIndexMeta::default();
        let mut count = 0;
        let mut files = read_dir(format!("{}/{}", storage_path, index_name)).await?;
        let mut dirs = Vec::new();
        while let Some(entry) = files.next_entry().await? {
            if entry.file_type().await?.is_dir() {
                dirs.push(entry.file_name().to_str().unwrap().to_string());
                continue;
            }

            debug!("Loading local index meta: {:?}", entry.path());
            if let Some(file_name) = entry.file_name().to_str() {
                if let Some(caps) = INDEX_VERSION_META_FILE_PATTERN.captures(file_name) {
                    if let (Some(start_version), Some(end_version)) = (
                        caps.name(START_VERSION_CAPTURE_GROUP_NAME),
                        caps.name(END_VERSION_CAPTURE_GROUP_NAME),
                    ) {
                        let start_version = start_version.as_str().parse::<u64>()?;
                        let end_version = end_version.as_str().parse::<u64>()?;
                        match end_version.cmp(&current.end_version()) {
                            Ordering::Greater => {
                                current = serde_json::from_str::<IncrementalIndexMeta>(
                                    &read_to_string(entry.path()).await?,
                                )?;
                            }
                            Ordering::Equal => {
                                if start_version < current.start_version() {
                                    current = serde_json::from_str::<IncrementalIndexMeta>(
                                        &read_to_string(entry.path()).await?,
                                    )?;
                                } else {
                                    continue;
                                }
                            }
                            _ => continue,
                        }

                        count += 1;
                        current.set_meta(
                            index_name.to_string(),
                            start_version,
                            end_version,
                            current.schema_version().clone(),
                        );
                    }
                }
            }
        }

        if count == 0 {
            return Ok(None);
        }

        // check if the current index is in the list of available indexes
        if !dirs.contains(&format!(
            "{}-{}",
            current.start_version(),
            current.end_version()
        )) {
            warn!("Current index is not available: {:?}", current);
            return Ok(None);
        }

        Ok(Some(Self {
            storage_path: storage_path.into(),
            incremental_meta: current,
            index_name: index_name.into(),
        }))
    }

    async fn open_or_create(&self, index_path: &str) -> Result<Index> {
        let directory = self.open_directory(index_path).await?;
        let index = spawn_blocking(move || Index::open_or_create(directory, INDEX_SCHEMA.clone()))
            .await??;
        Ok(Self::post_open_index(index))
    }

    fn post_open_index(index: Index) -> Index {
        index
            .tokenizers()
            .register(CharabiaTokenizer::name(), CharabiaTokenizer::new(2, 20));
        index
    }

    async fn open_directory(&self, index_path: &str) -> Result<Box<dyn Directory>> {
        let index_path = Path::new(index_path).to_path_buf();
        Ok(Box::new(
            spawn_blocking(move || MmapDirectory::open(index_path)).await??,
        ))
    }

    pub async fn writer(
        &self,
        memory_budget_in_bytes: impl Into<Option<usize>>,
    ) -> Result<IndexWriter> {
        self.writer_for_index(
            self.open_or_create(&self.index_path()).await?,
            memory_budget_in_bytes,
        )
        .await
    }

    async fn writer_for_index(
        &self,
        index: Index,
        memory_budget_in_bytes: impl Into<Option<usize>>,
    ) -> Result<IndexWriter> {
        let memory_budget_in_bytes = memory_budget_in_bytes
            .into()
            .unwrap_or(*DEFAULT_INDEXING_MEMORY_BUDGET);
        Ok(spawn_blocking(move || {
            index
                .writer(memory_budget_in_bytes)
                .inspect(|w| w.set_merge_policy(Box::new(NoMergePolicy)))
        })
        .await??)
    }

    pub async fn searcher(&self) -> Result<Searcher> {
        let index = self.open_or_create(&self.index_path()).await?;
        spawn_blocking(move || {
            Ok(index
                .reader_builder()
                .reload_policy(ReloadPolicy::Manual)
                .try_into()?
                .searcher())
        })
        .await?
    }

    fn index_path(&self) -> String {
        format!(
            "{}/{}/{}-{}",
            self.storage_path,
            self.index_name,
            self.start_version(),
            self.end_version()
        )
    }

    pub async fn compact(&self) -> Result<()> {
        let mut writer = self.writer(None).await?;

        let mut current_segment_count = 0;
        let mut preparing_segment_ids = Vec::new();
        for segment in writer
            .index()
            .searchable_segments()
            .as_deref()
            .unwrap_or(&[])
        {
            info!(
                "Compacting segment: {:?}, size: {:?}",
                segment.id().uuid_string(),
                segment.meta().num_docs()
            );
            if segment.meta().num_docs() >= *FTS_MAX_SEGMENT_DOCS {
                info!(
                    "Skipping segment {:?} as it has {} documents, which is >= {}",
                    segment.id(),
                    segment.meta().num_docs(),
                    *FTS_MAX_SEGMENT_DOCS
                );
                continue; // Skip maxed out segments
            }

            preparing_segment_ids.push(segment.id());
            current_segment_count += segment.meta().num_docs();
            if current_segment_count >= *FTS_MAX_SEGMENT_DOCS {
                let segment_meta = writer.merge(preparing_segment_ids.as_slice()).await?;
                info!(
                    "Compacted index {}-{:?} into {:?}",
                    self.index_name, &preparing_segment_ids, segment_meta
                );
                current_segment_count = 0;
                preparing_segment_ids.clear();
            }
        }

        let segment_meta = writer.merge(preparing_segment_ids.as_slice()).await?;
        info!(
            "Compacted index {}-{:?} into {:?}",
            self.index_name, &preparing_segment_ids, segment_meta
        );

        Ok(())
    }
}

pub async fn compact_index(storage_path: &str, index_name: &str) -> Result<Option<IndexMeta>> {
    match IndexMeta::load_local(storage_path, index_name).await? {
        Some(index_meta) => index_meta.compact().await.map(|_| Some(index_meta)),
        None => Ok(None),
    }
}

pub async fn apply_deleteds(
    storage_path: &str,
    storage: Arc<dyn cloud::Storage>,
    storage_meta: &StorageMeta,
) -> Result<()> {
    let deleteds = Deletes::load_deleteds(storage.clone()).await?;
    info!("Applying deleteds: {}", deleteds.len());
    for index in storage_meta.indexes.clone() {
        let updated_writer = Writer::new(
            storage_path,
            &index.0,
            storage_meta.first_start_version(&index.0),
            storage_meta.latest_end_version(&index.0),
            None,
        );
        updated_writer
            .update(DocumentFullVec::from(&deleteds).0)
            .await?;
    }
    Ok(())
}

#[derive(Debug, Default)]
pub struct Writer {
    storage_path: String,
    index_name: String,
    start_version: u64,
    end_version: u64,
    memory_budget: Option<usize>,
}

impl Writer {
    pub fn new(
        storage_path: &str,
        index_name: &str,
        start_version: u64,
        end_version: u64,
        memory_budget: impl Into<Option<usize>>,
    ) -> Self {
        debug!(
            "Creating Writer with storage_path: {}, index_name: {}, start_version: {}, end_version: {}",
            storage_path, index_name, start_version, end_version,
        );
        Self {
            storage_path: storage_path.into(),
            index_name: index_name.into(),
            start_version,
            end_version,
            memory_budget: memory_budget.into(),
        }
    }

    fn write_documents(
        writer: &mut IndexWriter,
        docs: Vec<DocumentFull>,
    ) -> anyhow::Result<Vec<String>> {
        let mut deleted = Vec::new();
        writer.run(docs.into_iter().map(|document| {
            if let Some(content) = document.content {
                let mut doc = TantivyDocument::default();
                doc.add_text(*crate::schema::DOC_ID_FIELD, &document.doc_id);
                doc.add_text(*crate::schema::TENANT_TOKEN_FIELD, document.tenant_token);
                doc.add_text(*crate::schema::CONTENT_FIELD, content);
                doc.add_u64(*crate::schema::VERSION_FIELD, document.version);
                doc.add_text(*crate::schema::FIELD_NAME_FIELD, document.field_name);
                doc.add_u64(*crate::schema::COLLECTION_ID_FIELD, document.collection_id);
                UserOperation::Add(doc)
            } else {
                deleted.push(document.doc_id.clone());
                // TODO confirm doc_id is unique across tenants
                UserOperation::Delete(Term::from_field_text(
                    *crate::schema::DOC_ID_FIELD,
                    &document.doc_id,
                ))
            }
        }))?;
        Ok(deleted)
    }

    pub async fn update(&self, documents: Vec<DocumentFull>) -> Result<(IndexMeta, Vec<String>)> {
        let index_meta = IndexMeta::new(
            &self.storage_path,
            &self.index_name,
            self.start_version,
            self.end_version,
        )
        .await?;
        if documents.is_empty() {
            return Ok((index_meta, Vec::new()));
        }

        let mut writer = index_meta.writer(self.memory_budget).await?;
        let deleteds = spawn_blocking(move || {
            let deleted = Self::write_documents(&mut writer, documents)?;
            writer.commit()?;
            Ok::<Vec<String>, anyhow::Error>(deleted)
        })
        .await??;
        Ok((index_meta, deleteds))
    }
}