// Copyright (C) 2021 Bosutech XXI S.L.
//
// nucliadb is offered under the AGPL v3.0 and as commercial software.
// For commercial licensing, contact us at info@nuclia.com.
//
// AGPL:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
//
use crate::query_io::translate_expression;
use crate::request_types::{ParagraphSearchRequest, ParagraphSuggestRequest};
use crate::set_query::SetQuery;
use itertools::Itertools;
use nidx_protos::StreamRequest;
use nidx_types::prefilter::PrefilterResult;
use nidx_types::query_language::BooleanExpression;
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use tantivy::query::*;
use tantivy::schema::{Facet, IndexRecordOption};
use tantivy::{DocId, InvertedIndexReader, Term};

use crate::fuzzy_query::FuzzyTermQuery;
use crate::schema::ParagraphSchema;
use crate::stop_words::is_stop_word;
type QueryP = (Occur, Box<dyn Query>);

// Used to identify the terms matched by tantivy
#[derive(Clone)]
pub struct TermCollector {
    pub eterms: HashSet<String>,
    pub fterms: HashMap<DocId, Vec<(Arc<InvertedIndexReader>, u64)>>,
}
impl Default for TermCollector {
    fn default() -> Self {
        Self::new()
    }
}
impl TermCollector {
    pub fn new() -> TermCollector {
        TermCollector {
            fterms: HashMap::new(),
            eterms: HashSet::new(),
        }
    }
    pub fn log_eterm(&mut self, term: String) {
        self.eterms.insert(term);
    }
    pub fn log_fterm(&mut self, doc: DocId, data: (Arc<InvertedIndexReader>, u64)) {
        self.fterms.entry(doc).or_default().push(data);
    }
    pub fn get_fterms(&self, doc: DocId) -> Vec<String> {
        let mut terms = Vec::new();
        for (index, term) in self.fterms.get(&doc).iter().flat_map(|v| v.iter()).cloned() {
            let term_dict = index.terms();
            let mut term_s = vec![];
            let found = term_dict.ord_to_term(term, &mut term_s).unwrap_or(false);
            let elem = if found {
                term_s
            } else {
                vec![]
            };
            match String::from_utf8(elem).ok() {
                Some(v) if v.len() > 2 => terms.push(v),
                _ => (),
            }
        }
        terms
    }
}

#[derive(Default, Clone)]
pub struct SharedTermC(Arc<Mutex<TermCollector>>);
impl SharedTermC {
    pub fn from(termc: TermCollector) -> SharedTermC {
        SharedTermC(Arc::new(Mutex::new(termc)))
    }
    pub fn get_termc(&self) -> TermCollector {
        std::mem::take(&mut self.0.lock().unwrap())
    }
    pub fn set_termc(&self, termc: TermCollector) {
        *self.0.lock().unwrap() = termc;
    }
}

fn term_to_fuzzy(query: Box<dyn Query>, distance: u8, termc: SharedTermC, as_prefix: bool) -> Box<dyn Query> {
    let term_query: &TermQuery = query.downcast_ref().unwrap();
    let term = term_query.term().clone();
    let value = term.value();
    let term_as_str = value.as_str();
    let should_be_prefixed = term_as_str.map(|s| as_prefix && s.len() > 3).unwrap_or_default();
    if should_be_prefixed {
        Box::new(FuzzyTermQuery::new_prefix(term, distance, true, termc))
    } else {
        Box::new(FuzzyTermQuery::new(term, distance, true, termc))
    }
}

fn queryp_map(queries: Vec<QueryP>, distance: u8, as_prefix: Option<usize>, termc: SharedTermC) -> Vec<QueryP> {
    queries
        .into_iter()
        .enumerate()
        .map(|(id, (_, query))| {
            let query = if query.is::<TermQuery>() {
                term_to_fuzzy(query, distance, termc.clone(), as_prefix.map_or(false, |v| id == v))
            } else {
                query
            };
            (Occur::Must, query)
        })
        .collect()
}

fn flat_bool_query(query: BooleanQuery, collector: (usize, Vec<QueryP>)) -> (usize, Vec<QueryP>) {
    query.clauses().iter().map(|(occur, subq)| (*occur, subq.box_clone())).fold(
        collector,
        |(mut id, mut c), (occur, subq)| {
            if subq.is::<BooleanQuery>() {
                let subq: Box<BooleanQuery> = subq.downcast().unwrap();
                flat_bool_query(*subq, (id, c))
            } else if subq.is::<TermQuery>() {
                id = c.len();
                c.push((occur, subq));
                (id, c)
            } else {
                c.push((occur, subq));
                (id, c)
            }
        },
    )
}

fn flat_and_adapt(query: Box<dyn Query>, prefixed: bool, distance: u8, termc: SharedTermC) -> Vec<QueryP> {
    let (queries, as_prefix) = if query.is::<BooleanQuery>() {
        let query: Box<BooleanQuery> = query.downcast().unwrap();
        let (as_prefix, queries) = flat_bool_query(*query, (usize::MAX, vec![]));
        (queries, as_prefix)
    } else if query.is::<TermQuery>() {
        let queries = vec![(Occur::Must, query)];
        let as_prefix = 0;
        (queries, as_prefix)
    } else {
        let queries = vec![(Occur::Must, query)];
        let as_prefix = 1;
        (queries, as_prefix)
    };
    queryp_map(
        queries,
        distance,
        if prefixed {
            Some(as_prefix)
        } else {
            None
        },
        termc,
    )
}

fn fuzzied_queries(query: Box<dyn Query>, prefixed: bool, distance: u8, termc: SharedTermC) -> Vec<QueryP> {
    if query.is::<AllQuery>() {
        vec![]
    } else {
        flat_and_adapt(query, prefixed, distance, termc)
    }
}

fn parse_query(parser: &QueryParser, text: &str) -> Box<dyn Query> {
    if text.is_empty() {
        Box::new(AllQuery) as Box<dyn Query>
    } else {
        parser.parse_query(text).ok().unwrap_or_else(|| Box::new(AllQuery))
    }
}

/// Removes all fuzzy terms identified as stop word.
///
/// A stop word is any fuzzy term that match the following criterias:
/// - Presents in the given list of stop words
/// - Is **NOT** the last term in the query
///
/// The last term of the query is a prefix fuzzy term and must be preserved.
fn remove_stop_words(query: &str) -> Cow<'_, str> {
    match query.rsplit_once(' ') {
        Some((query, last_term)) => {
            query.split(' ').filter(|term| !is_stop_word(&term.to_lowercase())).chain([last_term]).join(" ").into()
        }
        None => query.into(),
    }
}

#[derive(Debug)]
struct ProcessedQuery {
    fuzzy_query: String,
    regular_query: String,
}
fn preprocess_raw_query(query: &str, tc: &mut TermCollector) -> ProcessedQuery {
    let mut regular_query = String::new();
    let mut fuzzy_query = String::new();
    let mut quote_starts = vec![];
    let mut quote_ends = vec![];
    let mut start = 0;
    for (i, d) in query.match_indices('\"').enumerate() {
        if i % 2 == 0 {
            quote_starts.push(d.0)
        } else {
            quote_ends.push(d.0)
        }
    }
    for (qstart, qend) in quote_starts.into_iter().zip(quote_ends.into_iter()) {
        let quote = query[(qstart + 1)..qend].trim();
        let unquote = query[start..qstart].trim();
        let unquote = remove_stop_words(unquote);

        unquote.split(' ').filter(|s| !s.is_empty()).for_each(|t| tc.log_eterm(t.to_string()));
        tc.log_eterm(quote.to_string());

        if !regular_query.is_empty() {
            regular_query.push(' ');
        }
        regular_query.push_str(&unquote);
        regular_query.push(' ');
        regular_query.push('"');
        regular_query.push_str(quote);
        regular_query.push('"');

        if !fuzzy_query.is_empty() {
            fuzzy_query.push(' ');
        }
        fuzzy_query.push_str(&unquote);

        start = qend + 1;
    }
    if start < query.len() {
        let tail = query[start..].trim();
        let tail = remove_stop_words(tail);

        tail.split(' ').filter(|s| !s.is_empty()).for_each(|t| tc.log_eterm(t.to_string()));

        if !regular_query.is_empty() {
            regular_query.push(' ');
        }
        regular_query.push_str(&tail);

        if !fuzzy_query.is_empty() {
            fuzzy_query.push(' ');
        }
        fuzzy_query.push_str(&tail);
    }
    ProcessedQuery {
        regular_query,
        fuzzy_query,
    }
}

fn filter_query(
    schema: &ParagraphSchema,
    prefilter: &PrefilterResult,
    paragraph_formula: &Option<BooleanExpression>,
    filter_or: bool,
) -> Option<Box<dyn Query>> {
    let mut filter_terms = vec![];
    let operator = if filter_or {
        Occur::Should
    } else {
        Occur::Must
    };

    // Paragraph filter
    if let Some(formula) = &paragraph_formula {
        let query = translate_expression(formula, schema);
        filter_terms.push((operator, query));
    }

    // Prefilter
    if let PrefilterResult::Some(field_keys) = prefilter {
        let set_query = Box::new(SetQuery::new(
            schema.field_uuid,
            field_keys.iter().map(|x| format!("{}{}", x.resource_id.simple(), x.field_id)),
        ));
        filter_terms.push((operator, set_query));
    }

    if !filter_terms.is_empty() {
        Some(Box::new(BooleanQuery::new(filter_terms)))
    } else {
        None
    }
}

pub fn suggest_query(
    parser: &QueryParser,
    text: &str,
    request: &ParagraphSuggestRequest,
    prefilter: &PrefilterResult,
    schema: &ParagraphSchema,
    distance: u8,
) -> (Box<dyn Query>, SharedTermC, Box<dyn Query>) {
    let mut term_collector = TermCollector::default();
    let processed = preprocess_raw_query(text, &mut term_collector);
    let query = parse_query(parser, &processed.regular_query);
    let fuzzy_query = parse_query(parser, &processed.fuzzy_query);
    let termc = SharedTermC::from(term_collector);
    let mut fuzzies = fuzzied_queries(fuzzy_query, true, distance, termc.clone());
    let mut originals = vec![(Occur::Must, query)];
    let term = Term::from_field_u64(schema.repeated_in_field, 0);
    let term_query = TermQuery::new(term, IndexRecordOption::Basic);
    fuzzies.push((Occur::Must, Box::new(term_query.clone())));
    originals.push((Occur::Must, Box::new(term_query)));

    let filter_query = filter_query(schema, prefilter, &request.filtering_formula, request.filter_or);
    if let Some(query) = filter_query {
        originals.push((Occur::Must, query.box_clone()));
        fuzzies.push((Occur::Must, query));
    }

    if originals.len() == 1 && originals[0].1.is::<AllQuery>() {
        let original = originals.pop().unwrap().1;
        let fuzzy = Box::new(BooleanQuery::new(vec![]));
        (original, termc, fuzzy)
    } else {
        if processed.fuzzy_query.is_empty() {
            fuzzies.clear();
        }
        let original = Box::new(BooleanQuery::new(originals));
        let fuzzied = Box::new(BoostQuery::new(Box::new(BooleanQuery::new(fuzzies)), 0.5));
        (original, termc, fuzzied)
    }
}

pub fn search_query(
    parser: &QueryParser,
    text: &str,
    search: &ParagraphSearchRequest,
    prefilter: &PrefilterResult,
    schema: &ParagraphSchema,
    distance: u8,
    with_advance: Option<Box<dyn Query>>,
) -> (Box<dyn Query>, SharedTermC, Box<dyn Query>) {
    let mut term_collector = TermCollector::default();
    let processed = preprocess_raw_query(text, &mut term_collector);
    let query = parse_query(parser, &processed.regular_query);
    let fuzzy_query = parse_query(parser, &processed.fuzzy_query);
    let termc = SharedTermC::from(term_collector);
    let mut fuzzies = fuzzied_queries(fuzzy_query, false, distance, termc.clone());
    let mut originals = vec![(Occur::Must, query)];
    if let Some(advance) = with_advance {
        originals.push((Occur::Must, advance.box_clone()));
        fuzzies.push((Occur::Must, advance));
    }

    let filter_query = filter_query(schema, prefilter, &search.filtering_formula, search.filter_or);
    if let Some(query) = filter_query {
        originals.push((Occur::Must, query.box_clone()));
        fuzzies.push((Occur::Must, query));
    }

    if !search.with_duplicates {
        let term = Term::from_field_u64(schema.repeated_in_field, 0);
        let term_query = TermQuery::new(term, IndexRecordOption::Basic);
        fuzzies.push((Occur::Must, Box::new(term_query.clone())));
        originals.push((Occur::Must, Box::new(term_query)))
    }

    if originals.len() == 1 && originals[0].1.is::<AllQuery>() {
        let original = originals.pop().unwrap().1;
        let fuzzy = Box::new(BooleanQuery::new(vec![]));
        (original, termc, fuzzy)
    } else {
        if processed.fuzzy_query.is_empty() {
            fuzzies.clear();
        }
        let original = Box::new(BooleanQuery::new(originals));
        let fuzzied = Box::new(BoostQuery::new(Box::new(BooleanQuery::new(fuzzies)), 0.5));
        (original, termc, fuzzied)
    }
}

pub fn streaming_query(schema: &ParagraphSchema, request: &StreamRequest) -> Box<dyn Query> {
    let mut queries: Vec<(Occur, Box<dyn Query>)> = vec![];
    queries.push((Occur::Must, Box::new(AllQuery)));
    request
        .filter
        .iter()
        .flat_map(|f| f.labels.iter())
        .flat_map(|facet_key| Facet::from_text(facet_key).ok().into_iter())
        .for_each(|facet| {
            let facet_term = Term::from_facet(schema.facets, &facet);
            let facet_term_query = TermQuery::new(facet_term, IndexRecordOption::Basic);
            queries.push((Occur::Must, Box::new(facet_term_query)));
        });
    Box::new(BooleanQuery::new(queries))
}

#[cfg(test)]
mod tests {
    use tantivy::schema::Field;

    use super::*;
    fn dummy_term_query() -> Box<dyn Query> {
        let field = Field::from_field_id(0);
        let term = Term::from_field_u64(field, 0);
        Box::new(TermQuery::new(term, IndexRecordOption::Basic))
    }

    #[test]
    fn test_preprocessor() {
        let text = "own test \"This is great\"";
        let mut term_collector = TermCollector::default();
        let _ = preprocess_raw_query(text, &mut term_collector);
        let terms: HashSet<_> = term_collector.eterms.iter().map(|s| s.as_str()).collect();
        let expect = HashSet::from(["This is great", "test"]);
        assert_eq!(terms, expect);

        let text = "The test \"is correct\" always";
        let mut term_collector = TermCollector::default();
        let processed = preprocess_raw_query(text, &mut term_collector);
        let terms: HashSet<_> = term_collector.eterms.iter().map(|s| s.as_str()).collect();
        let expect = HashSet::from(["test", "always", "is correct"]);
        assert_eq!(terms, expect);
        assert_eq!(processed.regular_query, "test \"is correct\" always");
        assert_eq!(processed.fuzzy_query, "test always");
    }

    #[test]
    fn test() {
        let subqueries0: Vec<_> = vec![dummy_term_query; 12].into_iter().map(|f| (Occur::Must, f())).collect();
        let subqueries1: Vec<_> = vec![dummy_term_query; 12].into_iter().map(|f| (Occur::Must, f())).collect();
        let boolean0: Box<dyn Query> = Box::new(BooleanQuery::new(subqueries0));
        let boolean1: Box<dyn Query> = Box::new(BooleanQuery::new(subqueries1));
        let nested = BooleanQuery::new(vec![(Occur::Should, boolean0), (Occur::Should, boolean1)]);
        let adapted = flat_and_adapt(Box::new(nested), true, 2, SharedTermC::default());
        assert_eq!(adapted.len(), 24);
        assert!(adapted.iter().all(|(occur, _)| *occur == Occur::Must));
        assert!(adapted.iter().all(|(_, query)| query.is::<FuzzyTermQuery>()));
    }

    #[test]
    fn it_removes_stop_word_fterms() {
        let tests = [
            ("nuclia is a database for unstructured data", "nuclia database unstructured data"),
            (
                "nuclia is a database for the",
                // keeps last term even if is a stop word
                "nuclia database the",
            ),
            ("is a for and", "and"),
            ("what does stop is?", "stop is?"),
            ("", ""),
            ("comment s'appelle le train à grande vitesse", "comment s'appelle train grande vitesse"),
            ("¿Qué significa la palabra sentence en español?", "¿Qué significa palabra sentence español?"),
            ("Per què les vaques no són de color rosa?", "vaques color rosa?"),
            ("How can I learn to make a flat white?", "learn make flat white?"),
            ("Qué es escalada en bloque?", "escalada bloque?"),
            (
                "Wer hat gesagt: 'Kaffeetrinken ist integraler Bestandteil des Kletterns'?",
                "Wer gesagt: 'Kaffeetrinken integraler Bestandteil Kletterns'?",
            ),
            ("i pistacchi siciliani sono i migliori al mondo?", "pistacchi siciliani migliori mondo?"),
        ];

        for (query, expected_fuzzy_query) in tests {
            let fuzzy_query = remove_stop_words(query);
            assert_eq!(fuzzy_query, expected_fuzzy_query);
        }
    }
}
