use common::search::{
    SchemaFields, json_path_term, normalize_column_name, normalize_exact, tokenize_ngram,
    tokenize_word,
};
use common::proto::komp_ac::search::SearchVersionScope;
use tantivy::query::{
    BooleanQuery, BoostQuery, EmptyQuery, FuzzyTermQuery, Occur, PhraseQuery, Query, QueryParser,
    TermQuery,
};
use tantivy::schema::{IndexRecordOption, Term};
use tantivy::Index;
use tonic::Status;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ConstraintMode {
    Fuzzy,
    Exact,
}

#[derive(Clone, Debug)]
pub struct SearchConstraint {
    pub targets: Vec<SearchConstraintTarget>,
    pub mode: ConstraintMode,
}

#[derive(Clone, Debug)]
pub struct SearchConstraintTarget {
    pub table_name: Option<String>,
    pub column: String,
    pub query: String,
}

pub fn build_master_query(
    index: &Index,
    fields: &SchemaFields,
    free_query: &str,
    must: &[SearchConstraint],
    table_filter: Option<&str>,
    version_scope: SearchVersionScope,
) -> Result<Box<dyn Query>, Status> {
    let mut clauses: Vec<(Occur, Box<dyn Query>)> = Vec::new();
    let mut has_search_clause = false;

    for constraint in must {
        let predicate = constraint_predicate(fields, constraint)?;
        clauses.push((Occur::Must, predicate));
        has_search_clause = true;
    }

    let free_words = tokenize_word(free_query);
    if !free_words.is_empty() {
        let predicate = fuzzy_predicate_unscoped(index, fields, &free_words)?;
        clauses.push((Occur::Must, predicate));
        has_search_clause = true;
    }

    if let Some(table_name) = table_filter {
        let term = Term::from_field_text(fields.table_name, table_name);
        clauses.push((
            Occur::Must,
            Box::new(TermQuery::new(term, IndexRecordOption::Basic)),
        ));
    }

    let archived_rows = |archived| {
        Box::new(TermQuery::new(
            Term::from_field_u64(fields.is_archived, u64::from(archived)),
            IndexRecordOption::Basic,
        )) as Box<dyn Query>
    };
    match version_scope {
        SearchVersionScope::Current => clauses.push((Occur::Must, archived_rows(false))),
        SearchVersionScope::Archived => clauses.push((Occur::Must, archived_rows(true))),
        SearchVersionScope::All => {}
    }

    if !has_search_clause {
        return Ok(Box::new(EmptyQuery));
    }

    Ok(Box::new(BooleanQuery::new(clauses)))
}

fn constraint_predicate(
    fields: &SchemaFields,
    constraint: &SearchConstraint,
) -> Result<Box<dyn Query>, Status> {
    let mut alternatives = Vec::new();

    for target in &constraint.targets {
        let column_predicate = match constraint.mode {
            ConstraintMode::Exact => exact_predicate(fields, &target.column, &target.query)?,
            ConstraintMode::Fuzzy => {
                fuzzy_predicate_scoped(fields, &target.column, &target.query)?
            }
        };

        if let Some(table_name) = &target.table_name {
            let table_term = Term::from_field_text(fields.table_name, table_name);
            alternatives.push((
                Occur::Should,
                Box::new(BooleanQuery::new(vec![
                    (
                        Occur::Must,
                        Box::new(TermQuery::new(table_term, IndexRecordOption::Basic))
                            as Box<dyn Query>,
                    ),
                    (Occur::Must, column_predicate),
                ])) as Box<dyn Query>,
            ));
        } else {
            alternatives.push((Occur::Should, column_predicate));
        }
    }

    if alternatives.is_empty() {
        return Err(Status::invalid_argument(
            "constraint has no searchable column targets",
        ));
    }

    Ok(Box::new(BooleanQuery::new(alternatives)))
}

fn exact_predicate(
    fields: &SchemaFields,
    column: &str,
    query: &str,
) -> Result<Box<dyn Query>, Status> {
    let normalized_value = normalize_exact(query);
    if normalized_value.is_empty() {
        return Err(Status::invalid_argument(
            "exact query is empty after normalization",
        ));
    }

    let column = normalize_column_name(column);
    let term = json_path_term(fields.data_exact, &column, &normalized_value);
    Ok(Box::new(TermQuery::new(term, IndexRecordOption::Basic)))
}

fn fuzzy_predicate_scoped(
    fields: &SchemaFields,
    column: &str,
    query: &str,
) -> Result<Box<dyn Query>, Status> {
    let words = tokenize_word(query);
    if words.is_empty() {
        return Err(Status::invalid_argument(
            "fuzzy query has no searchable tokens",
        ));
    }

    let column = normalize_column_name(column);

    let mut layers: Vec<(Occur, Box<dyn Query>)> = Vec::new();

    let mut per_word_clauses: Vec<(Occur, Box<dyn Query>)> = Vec::new();
    for word in &words {
        let term = json_path_term(fields.data_word, &column, word);
        let mut alternates: Vec<(Occur, Box<dyn Query>)> = Vec::new();

        alternates.push((
            Occur::Should,
            Box::new(BoostQuery::new(
                Box::new(TermQuery::new(term.clone(), IndexRecordOption::WithFreqs)),
                4.0,
            )),
        ));

        alternates.push((
            Occur::Should,
            Box::new(BoostQuery::new(
                Box::new(FuzzyTermQuery::new_prefix(term.clone(), 0, false)),
                3.0,
            )),
        ));

        if let Some(distance) = fuzzy_distance(word.chars().count()) {
            alternates.push((
                Occur::Should,
                Box::new(BoostQuery::new(
                    Box::new(FuzzyTermQuery::new(term.clone(), distance, true)),
                    2.0,
                )),
            ));
        }

        per_word_clauses.push((Occur::Must, Box::new(BooleanQuery::new(alternates))));
    }
    layers.push((Occur::Should, Box::new(BooleanQuery::new(per_word_clauses))));

    if words.len() > 1 {
        let phrase_terms: Vec<(usize, Term)> = words
            .iter()
            .enumerate()
            .map(|(offset, word)| (offset, json_path_term(fields.data_word, &column, word)))
            .collect();
        let phrase = PhraseQuery::new_with_offset_and_slop(phrase_terms, 3);
        layers.push((
            Occur::Should,
            Box::new(BoostQuery::new(Box::new(phrase), 2.0)),
        ));
    }

    let ngrams = tokenize_ngram(query);
    if !ngrams.is_empty() {
        let ngram_clauses: Vec<(Occur, Box<dyn Query>)> = ngrams
            .into_iter()
            .map(|gram| {
                let term = json_path_term(fields.data_ngram, &column, &gram);
                (
                    Occur::Must,
                    Box::new(TermQuery::new(term, IndexRecordOption::Basic)) as Box<dyn Query>,
                )
            })
            .collect();
        layers.push((
            Occur::Should,
            Box::new(BoostQuery::new(
                Box::new(BooleanQuery::new(ngram_clauses)),
                1.0,
            )),
        ));
    }

    Ok(Box::new(BooleanQuery::new(layers)))
}

fn fuzzy_predicate_unscoped(
    index: &Index,
    fields: &SchemaFields,
    words: &[String],
) -> Result<Box<dyn Query>, Status> {
    let mut layers: Vec<(Occur, Box<dyn Query>)> = Vec::new();

    let mut per_word_clauses: Vec<(Occur, Box<dyn Query>)> = Vec::new();
    for word in words {
        let term = Term::from_field_text(fields.all_text, word);
        let mut alternates: Vec<(Occur, Box<dyn Query>)> = Vec::new();

        alternates.push((
            Occur::Should,
            Box::new(BoostQuery::new(
                Box::new(TermQuery::new(term.clone(), IndexRecordOption::WithFreqs)),
                4.0,
            )),
        ));

        alternates.push((
            Occur::Should,
            Box::new(BoostQuery::new(
                Box::new(FuzzyTermQuery::new_prefix(term.clone(), 0, false)),
                3.0,
            )),
        ));

        if let Some(distance) = fuzzy_distance(word.chars().count()) {
            alternates.push((
                Occur::Should,
                Box::new(BoostQuery::new(
                    Box::new(FuzzyTermQuery::new(term, distance, true)),
                    2.0,
                )),
            ));
        }

        per_word_clauses.push((Occur::Must, Box::new(BooleanQuery::new(alternates))));
    }
    layers.push((Occur::Should, Box::new(BooleanQuery::new(per_word_clauses))));

    if words.len() > 1 {
        let parser = QueryParser::for_index(index, vec![fields.all_text]);
        let query_string = format!("\"{}\"~3", words.join(" "));
        if let Ok(query) = parser.parse_query(&query_string) {
            layers.push((Occur::Should, Box::new(BoostQuery::new(query, 2.0))));
        }
    }

    {
        let parser = QueryParser::for_index(index, vec![fields.all_text]);
        let query_string = words
            .iter()
            .map(|word| format!("+{}*", word))
            .collect::<Vec<_>>()
            .join(" ");
        if let Ok(query) = parser.parse_query(&query_string) {
            layers.push((Occur::Should, Box::new(BoostQuery::new(query, 1.0))));
        }
    }

    if layers.is_empty() {
        return Ok(Box::new(EmptyQuery));
    }

    Ok(Box::new(BooleanQuery::new(layers)))
}

fn fuzzy_distance(word_len: usize) -> Option<u8> {
    match word_len {
        0..=3 => None,
        4..=6 => Some(1),
        _ => Some(2),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use common::search::{archived_search_row_key, create_search_schema, register_tokenizers};
    use tantivy::collector::Count;
    use tantivy::schema::OwnedValue;
    use tantivy::TantivyDocument;

    struct TestDocument<'a> {
        row_key: String,
        row_id: u64,
        table_name: &'a str,
        values: &'a [(&'a str, &'a str)],
    }

    fn index_documents(documents: &[TestDocument<'_>]) -> (Index, SchemaFields) {
        let schema = create_search_schema();
        let index = Index::create_in_ram(schema.clone());
        register_tokenizers(&index).expect("tokenizers should register");
        let fields = SchemaFields::from(&schema).expect("schema should match");
        let mut writer = index.writer(50_000_000).expect("writer should open");

        for source in documents {
            let mut document = TantivyDocument::default();
            document.add_u64(fields.pg_id, source.row_id);
            document.add_u64(fields.row_revision, 1);
            let archive = common::search::parse_archived_search_row_key(&source.row_key);
            document.add_u64(fields.is_archived, u64::from(archive.is_some()));
            document.add_u64(
                fields.table_definition_id,
                archive.map(|item| item.0 as u64).unwrap_or_default(),
            );
            document.add_u64(
                fields.version_number,
                archive.map(|item| item.2 as u64).unwrap_or_default(),
            );
            document.add_text(fields.table_name, source.table_name);
            document.add_text(fields.row_key, &source.row_key);
            let mut object = std::collections::BTreeMap::new();
            for (column, value) in source.values {
                document.add_text(fields.all_text, value);
                object.insert((*column).to_string(), OwnedValue::from(*value));
            }
            document.add_object(fields.data_word, object.clone());
            document.add_object(fields.data_ngram, object.clone());
            document.add_object(fields.data_exact, object);
            writer.add_document(document).expect("document should index");
        }
        writer.commit().expect("documents should commit");
        (index, fields)
    }

    fn index_versions() -> (Index, SchemaFields) {
        index_documents(&[
            TestDocument {
                row_key: "customers:7".to_string(),
                row_id: 7,
                table_name: "customers",
                values: &[("1", "new")],
            },
            TestDocument {
                row_key: archived_search_row_key(40, 7, 1),
                row_id: 7,
                table_name: "customers",
                values: &[("1", "old")],
            },
        ])
    }

    fn count(index: &Index, query: &dyn Query) -> usize {
        index
            .reader()
            .expect("reader should open")
            .searcher()
            .search(query, &Count)
            .expect("query should run")
    }

    #[test]
    fn version_scope_separates_current_and_archived_documents() {
        let (index, fields) = index_versions();
        let constraint = SearchConstraint {
            targets: vec![SearchConstraintTarget {
                table_name: None,
                column: "1".to_string(),
                query: "old".to_string(),
            }],
            mode: ConstraintMode::Exact,
        };

        let current = build_master_query(
            &index,
            &fields,
            "",
            std::slice::from_ref(&constraint),
            Some("customers"),
            SearchVersionScope::Current,
        )
        .expect("current query should build");
        let archived = build_master_query(
            &index,
            &fields,
            "",
            &[constraint],
            Some("customers"),
            SearchVersionScope::Archived,
        )
        .expect("archive query should build");

        assert_eq!(count(&index, &*current), 0);
        assert_eq!(count(&index, &*archived), 1);
    }

    #[test]
    fn physical_constraint_survives_any_display_alias_rename() {
        let (index, fields) = index_versions();
        for public_alias in ["adresar", "customer", "renamed_again"] {
            let resolved_physical_column = if !public_alias.is_empty() { "1" } else { unreachable!() };
            let query = build_master_query(
                &index,
                &fields,
                "",
                &[SearchConstraint {
                    targets: vec![SearchConstraintTarget {
                        table_name: None,
                        column: resolved_physical_column.to_string(),
                        query: "new".to_string(),
                    }],
                    mode: ConstraintMode::Exact,
                }],
                Some("customers"),
                SearchVersionScope::Current,
            )
            .expect("renamed alias should resolve to the same physical query");
            assert_eq!(count(&index, &*query), 1, "alias {public_alias}");
        }
    }

    #[test]
    fn exact_constraints_are_anded_and_table_scoped() {
        let (index, fields) = index_documents(&[
            TestDocument {
                row_key: "customers:1".to_string(),
                row_id: 1,
                table_name: "customers",
                values: &[("1", "Alice Example"), ("2", "Bratislava")],
            },
            TestDocument {
                row_key: "customers:2".to_string(),
                row_id: 2,
                table_name: "customers",
                values: &[("1", "Alice Example"), ("2", "Kosice")],
            },
            TestDocument {
                row_key: "suppliers:3".to_string(),
                row_id: 3,
                table_name: "suppliers",
                values: &[("1", "Alice Example"), ("2", "Bratislava")],
            },
        ]);
        let constraints = [
            SearchConstraint {
                targets: vec![SearchConstraintTarget {
                    table_name: None,
                    column: "1".to_string(),
                    query: "Alice Example".to_string(),
                }],
                mode: ConstraintMode::Exact,
            },
            SearchConstraint {
                targets: vec![SearchConstraintTarget {
                    table_name: None,
                    column: "2".to_string(),
                    query: "Bratislava".to_string(),
                }],
                mode: ConstraintMode::Exact,
            },
        ];
        let query = build_master_query(
            &index,
            &fields,
            "",
            &constraints,
            Some("customers"),
            SearchVersionScope::Current,
        )
        .expect("exact query should build");

        assert_eq!(count(&index, &*query), 1);
    }

    #[test]
    fn fuzzy_constraints_stay_within_the_resolved_physical_column() {
        let (index, fields) = index_documents(&[
            TestDocument {
                row_key: "customers:1".to_string(),
                row_id: 1,
                table_name: "customers",
                values: &[("1", "Alice"), ("2", "Bratislava")],
            },
            TestDocument {
                row_key: "customers:2".to_string(),
                row_id: 2,
                table_name: "customers",
                values: &[("1", "Bratislava"), ("2", "Kosice")],
            },
        ]);
        let constraint = |column: &str| SearchConstraint {
            targets: vec![SearchConstraintTarget {
                table_name: None,
                column: column.to_string(),
                query: "Bratislva".to_string(),
            }],
            mode: ConstraintMode::Fuzzy,
        };
        let address_query = build_master_query(
            &index,
            &fields,
            "",
            &[constraint("2")],
            Some("customers"),
            SearchVersionScope::Current,
        )
        .expect("fuzzy address query should build");
        let name_query = build_master_query(
            &index,
            &fields,
            "",
            &[constraint("1")],
            Some("customers"),
            SearchVersionScope::Current,
        )
        .expect("fuzzy name query should build");

        assert_eq!(count(&index, &*address_query), 1);
        assert_eq!(count(&index, &*name_query), 1);
    }

    #[test]
    fn free_text_searches_all_public_values_but_respects_table_filter() {
        let (index, fields) = index_documents(&[
            TestDocument {
                row_key: "customers:1".to_string(),
                row_id: 1,
                table_name: "customers",
                values: &[("1", "Alice Example")],
            },
            TestDocument {
                row_key: "suppliers:2".to_string(),
                row_id: 2,
                table_name: "suppliers",
                values: &[("8", "Alice Example")],
            },
        ]);
        let query = build_master_query(
            &index,
            &fields,
            "Alice Example",
            &[],
            Some("customers"),
            SearchVersionScope::Current,
        )
        .expect("free-text query should build");

        assert_eq!(count(&index, &*query), 1);
    }
}