use serde_aux::prelude::*;
use std::{
    cmp::Reverse,
    collections::HashMap,
    fs::{File, create_dir_all, remove_dir_all},
    io::{BufWriter, Write},
    path::Path,
    sync::Arc,
};

use crate::{
    data::{Data, DataSource},
    index::{Match, Search, SearchParamsExt},
    utils::{load_u32_vec, progress_bar},
};
use anyhow::{Result, anyhow};
use ordered_float::OrderedFloat;
use serde::Deserialize;
use tantivy::Index;
use tantivy::query::QueryParser;
use tantivy::schema::*;
use tantivy::{IndexReader, collector::TopDocs};

struct Inner {
    data: Data,
    index: Index,
    reader: IndexReader,
    parser: QueryParser,
    field_to_data: Vec<u32>,
}

impl std::fmt::Debug for Inner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Inner")
            .field("data", &self.data)
            .field("index", &self.index)
            .field("reader", &"IndexReader { ... }")
            .field("parser", &"QueryParser { ... }")
            .finish()
    }
}

#[derive(Debug, Clone)]
pub struct FullTextIndex {
    inner: Arc<Inner>,
}

impl FullTextIndex {
    fn field_to_column(&self, field_id: u32) -> usize {
        let field_to_data = &self.inner.field_to_data;

        let mut field_id = field_id as usize;
        let data_id = field_to_data[field_id] as usize;

        let mut offset = 0;
        while field_id > 0 && field_to_data[field_id - 1] == data_id as u32 {
            field_id -= 1;
            offset += 1;
        }
        offset
    }
}

#[derive(Debug, Clone, Deserialize)]
pub struct FullTextSearchParams {
    #[serde(
        default = "crate::index::default_k",
        deserialize_with = "deserialize_number_from_string"
    )]
    pub k: usize,
    #[serde(
        default,
        rename = "min-score",
        deserialize_with = "deserialize_option_number_from_string"
    )]
    pub min_score: Option<f32>,
    #[serde(default, deserialize_with = "deserialize_bool_from_anything")]
    pub exact: bool,
}

impl SearchParamsExt for FullTextSearchParams {
    fn k(&self) -> usize {
        self.k
    }

    fn exact(&self) -> bool {
        self.exact
    }
}

impl Search for FullTextIndex {
    type Data = Data;
    type Query<'e> = &'e str;
    type BuildParams = ();
    type SearchParams = FullTextSearchParams;

    fn build(data: &Self::Data, index_dir: &Path, _params: &Self::BuildParams) -> Result<()>
    where
        Self: Sized,
    {
        let mut schema_builder = Schema::builder();
        schema_builder.add_text_field("text", TEXT);
        schema_builder.add_u64_field("field_id", STORED | INDEXED);
        let schema = schema_builder.build();

        if index_dir.exists() {
            // need to remove here because tantivy does not support overwriting existing index
            remove_dir_all(index_dir)?;
        };

        create_dir_all(index_dir)?;
        let index = Index::create_in_dir(index_dir, schema.clone())?;
        let mut index_writer = index.writer(100_000_000)?;

        let text_field = schema.get_field("text")?;
        let field_id_field = schema.get_field("field_id")?;

        let total_fields = data.total_fields();
        let pb = progress_bar("Building full-text index", Some(total_fields as u64))?;

        let mut field_to_data_file =
            BufWriter::new(File::create(index_dir.join("index.field-to-data"))?);
        let mut field_id: u32 = 0;
        for (id, fields) in data.items() {
            for field in fields {
                if field_id == u32::MAX {
                    return Err(anyhow!("too many fields, max {} supported", u32::MAX));
                }

                if !field.is_text() {
                    return Err(anyhow!("Full-text index can only be built on text fields"));
                }

                let mut doc = TantivyDocument::default();
                doc.add_text(text_field, field.as_str());
                doc.add_u64(field_id_field, field_id as u64);
                index_writer.add_document(doc)?;

                field_id += 1;
                field_to_data_file.write_all(&id.to_le_bytes())?;

                pb.inc(1);
            }
        }

        pb.finish_with_message("Full-text index built");

        index_writer.commit()?;

        Ok(())
    }

    fn load(data: Self::Data, index_dir: &Path) -> Result<Self>
    where
        Self: Sized,
    {
        let mut index = Index::open_in_dir(index_dir)?;
        index.set_default_multithread_executor()?;

        let reader = index.reader()?;

        let text_field = index.schema().get_field("text")?;
        let parser = QueryParser::for_index(&index, vec![text_field]);

        let field_to_data = load_u32_vec(&index_dir.join("index.field-to-data"))?;

        Ok(FullTextIndex {
            inner: Arc::new(Inner {
                data,
                index,
                reader,
                parser,
                field_to_data,
            }),
        })
    }

    fn search(&self, query: Self::Query<'_>, params: &Self::SearchParams) -> Result<Vec<Match>> {
        let index = &self.inner.index;
        let schema = index.schema();
        let field_id_field = schema.get_field("field_id")?;

        let query = self.inner.parser.parse_query(query)?;

        let search_k = params.search_k(&self.inner.data);
        let top_docs = TopDocs::with_limit(search_k);

        let searcher = self.inner.reader.searcher();

        let mut matches = Vec::with_capacity(params.k);
        for (score, doc) in searcher.search(&query, &top_docs)? {
            if let Some(min_score) = params.min_score
                && score < min_score
            {
                continue;
            }

            let doc: HashMap<_, _> = searcher.doc(doc)?;
            let field_id = doc
                .get(&field_id_field)
                .and_then(|v| v.as_u64())
                .ok_or_else(|| anyhow!("Missing 'field_id' field in document"))?
                as u32;

            let data_id = self.inner.field_to_data[field_id as usize];

            matches.push((data_id, field_id, score));
        }

        // Sort by data_id, then by score descending, then by field_id
        matches.sort_by_key(|&(data_id, field_id, score)| {
            (data_id, Reverse(OrderedFloat(score)), field_id)
        });

        // Deduplicate by data_id (keeping the best score)
        matches.dedup_by(|a, b| a.0 == b.0);

        // Sort by score descending, then by data_id
        matches.sort_by_key(|&(data_id, _field_id, score)| (Reverse(OrderedFloat(score)), data_id));

        // Take top k
        let matches: Vec<Match> = matches
            .into_iter()
            .map(|(data_id, field_id, score)| {
                Match::WithField(data_id, self.field_to_column(field_id), score)
            })
            .take(params.k)
            .collect();

        Ok(matches)
    }

    fn search_with_filter<F>(
        &self,
        _query: Self::Query<'_>,
        _params: &Self::SearchParams,
        _filter: F,
    ) -> Result<Vec<Match>>
    where
        F: Fn(u32) -> bool,
    {
        Err(anyhow!(
            "Filtered search is not supported for full-text index"
        ))
    }

    fn data(&self) -> &Self::Data {
        &self.inner.data
    }

    fn index_type(&self) -> &'static str {
        "full-text"
    }
}
