//! Query processing and document retrieval.

use snafu::{ResultExt, Snafu};
use tokio::task::JoinError;

use crate::{
    storage::{
        self,
        queries::{DocumentMatch, Queries},
    },
    text_encoder::{self, TextEncoder},
    tokenize::{self, Tokenizer},
};

/// Errors that occur during query processing.
#[derive(Debug, Snafu)]
pub enum AskError {
    #[snafu(display("Failed to encode query."))]
    Encode { source: tokenize::EncodeError },
    #[snafu(display("Failed to embed query."))]
    Embed { source: text_encoder::EmbedError },
    #[snafu(display("Embedding task failed to execute."))]
    EmbedJoin { source: JoinError },
    #[snafu(display("Failed to retrieve similar documents."))]
    Query {
        source: storage::queries::QueryError,
    },
    #[snafu(display("Failed to rerank documents."))]
    Rerank { source: text_encoder::RerankError },
    #[snafu(display("Reranking task failed to execute."))]
    RerankJoin { source: JoinError },
}

/// Process a user query and return ranked document matches.
pub async fn ask(
    query: &str,
    db: &Queries,
    tokenizer: &Tokenizer,
    embedder: &TextEncoder,
    reranker: &TextEncoder,
    chunk_size: usize,
    limit: usize,
) -> Result<Vec<DocumentMatch>, AskError> {
    let encodings = tokenizer.encode(query, chunk_size).context(EncodeSnafu)?;

    let embedder = embedder.clone();
    let encoding = encodings[0].clone();
    let embeddings = tokio::task::spawn_blocking(move || embedder.embed(encoding))
        .await
        .context(EmbedJoinSnafu)?
        .context(EmbedSnafu)?;

    let documents = db
        .query(embeddings, (limit * 10) as i32)
        .await
        .context(QuerySnafu)?;

    let reranker = reranker.clone();
    let tokenizer = tokenizer.clone();
    let query = query.to_string();
    let reranked_docs = tokio::task::spawn_blocking(move || {
        reranker.rerank(&query, documents, &tokenizer, limit)
    })
    .await
    .context(RerankJoinSnafu)?
    .context(RerankSnafu)?;

    Ok(reranked_docs)
}