RustBox/src/routes/search.rs

use std::sync::Arc;

use axum::extract::{Path, Query};
use axum::response::{IntoResponse, Response};
use axum::{Extension, Json};
use serde::{Deserialize, Serialize};

use crate::auth::session::SessionState;
use crate::config::AppConfig;
use crate::db;
use crate::error::AppError;
use crate::search::engine::{SearchEngine, SearchQuery};

struct ParsedQuery {
    text: String,
    subject_only: Option<String>,
    folder: Option<String>,
    from: Option<String>,
    to: Option<String>,
    date_from: Option<i64>,
    date_to: Option<i64>,
    has_attachment: Option<bool>,
    unread: Option<bool>,
    flagged: Option<bool>,
}

fn parse_search_query(input: &str) -> ParsedQuery {
    let mut text_parts: Vec<String> = Vec::new();
    let mut subject_only: Option<String> = None;
    let mut folder: Option<String> = None;
    let mut from: Option<String> = None;
    let mut to: Option<String> = None;
    let mut date_from: Option<i64> = None;
    let mut date_to: Option<i64> = None;
    let mut has_attachment: Option<bool> = None;
    let mut unread: Option<bool> = None;
    let mut flagged: Option<bool> = None;

    let tokens = tokenize_query(input);

    for token in tokens {
        if let Some((operator, value)) = split_operator(&token) {
            let op = operator.to_lowercase();
            match op.as_str() {
                "from" => from = Some(value),
                "to" => to = Some(value),
                "subject" => subject_only = Some(value),
                "in" | "folder" => folder = Some(value),
                "date" => {
                    if let Some((start, end)) = parse_date_range(&value) {
                        date_from = Some(start);
                        date_to = Some(end);
                    }
                }
                "after" => {
                    if let Some(epoch) = parse_date_start(&value) {
                        date_from = Some(epoch);
                    }
                }
                "before" => {
                    if let Some(epoch) = parse_date_start(&value) {
                        date_to = Some(epoch);
                    }
                }
                "has" if value.eq_ignore_ascii_case("attachment") => {
                    has_attachment = Some(true);
                }
                "is" if value.eq_ignore_ascii_case("unread") => unread = Some(true),
                "is" if value.eq_ignore_ascii_case("read") => unread = Some(false),
                "is" if value.eq_ignore_ascii_case("flagged")
                    || value.eq_ignore_ascii_case("starred") =>
                {
                    flagged = Some(true)
                }
                "is" if value.eq_ignore_ascii_case("unflagged")
                    || value.eq_ignore_ascii_case("unstarred") =>
                {
                    flagged = Some(false)
                }
                _ => {
                    text_parts.push(token);
                }
            }
        } else {
            text_parts.push(token);
        }
    }

    ParsedQuery {
        text: text_parts.join(" "),
        subject_only,
        folder,
        from,
        to,
        date_from,
        date_to,
        has_attachment,
        unread,
        flagged,
    }
}

fn tokenize_query(input: &str) -> Vec<String> {
    let mut tokens = Vec::new();
    let mut chars = input.chars().peekable();
    let mut current = String::new();

    while let Some(&ch) = chars.peek() {
        if ch.is_whitespace() {
            if !current.is_empty() {
                tokens.push(std::mem::take(&mut current));
            }
            chars.next();
        } else if ch == '"' {
            chars.next();
            while let Some(&qch) = chars.peek() {
                if qch == '"' {
                    chars.next();
                    break;
                }
                current.push(qch);
                chars.next();
            }
        } else {
            current.push(ch);
            chars.next();
        }
    }

    if !current.is_empty() {
        tokens.push(current);
    }

    tokens
}

fn split_operator(token: &str) -> Option<(String, String)> {
    if let Some(colon_pos) = token.find(':') {
        let op = &token[..colon_pos];
        let val = &token[colon_pos + 1..];

        if !op.is_empty() && op.chars().all(|c| c.is_ascii_alphabetic()) && !val.is_empty() {
            return Some((op.to_string(), val.to_string()));
        }
    }
    None
}

fn parse_date_range(date_str: &str) -> Option<(i64, i64)> {
    let start = parse_date_start(date_str)?;
    let end = start + 86_400 - 1;
    Some((start, end))
}

fn parse_date_start(date_str: &str) -> Option<i64> {
    let parts: Vec<&str> = date_str.split('-').collect();
    if parts.len() != 3 {
        return None;
    }
    let year: i64 = parts[0].parse().ok()?;
    let month: i64 = parts[1].parse().ok()?;
    let day: i64 = parts[2].parse().ok()?;

    if !(1..=12).contains(&month) || !(1..=31).contains(&day) {
        return None;
    }

    let mut total_days: i64 = 0;

    for y in 1970..year {
        total_days += if is_leap_year(y) { 366 } else { 365 };
    }

    let days_in_months = [
        31,
        28 + i64::from(is_leap_year(year)),
        31,
        30,
        31,
        30,
        31,
        31,
        30,
        31,
        30,
        31,
    ];
    for days in days_in_months.iter().take((month - 1) as usize) {
        total_days += days;
    }

    total_days += day - 1;

    Some(total_days * 86_400)
}

fn is_leap_year(y: i64) -> bool {
    (y % 4 == 0 && y % 100 != 0) || y % 400 == 0
}

#[derive(Deserialize)]
pub struct SearchParams {
    pub q: Option<String>,

    pub folder: Option<String>,

    pub subject: Option<String>,

    pub from: Option<String>,

    pub to: Option<String>,

    pub date_from: Option<i64>,

    pub date_to: Option<i64>,

    pub has_attachment: Option<bool>,

    pub unread: Option<bool>,

    pub flagged: Option<bool>,

    #[serde(default = "default_limit")]
    pub limit: usize,

    #[serde(default)]
    pub offset: usize,

    #[serde(default = "default_sort")]
    pub sort: String,

    #[serde(default)]
    pub all: bool,
}

fn default_sort() -> String {
    "date_desc".to_string()
}

fn default_limit() -> usize {
    10_000
}

#[derive(Serialize)]
pub struct SearchResponse {
    pub results: Vec<SearchResultItem>,
    pub total_count: usize,
    pub query: String,
}

#[derive(Serialize)]
pub struct SearchResultItem {
    pub uid: u32,
    pub folder: String,
    pub score: f32,
    pub subject: String,
    pub from_address: String,
    pub from_name: String,
    pub to_addresses: String,
    pub date: String,
    pub flags: String,
    pub has_attachments: bool,
    pub snippet: String,
}

pub async fn search_messages(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Extension(search_engine): Extension<Arc<SearchEngine>>,
    Query(params): Query<SearchParams>,
) -> Result<Response, AppError> {
    let query_text = params.q.as_deref().unwrap_or("").trim().to_string();

    let has_direct_filters = params.folder.is_some()
        || params.subject.is_some()
        || params.from.is_some()
        || params.to.is_some()
        || params.date_from.is_some()
        || params.date_to.is_some()
        || params.has_attachment.is_some()
        || params.unread.is_some()
        || params.flagged.is_some();
    if query_text.is_empty() && !has_direct_filters {
        return Err(AppError::BadRequest(
            "A search query or at least one filter is required".to_string(),
        ));
    }

    let sort_order = params.sort.clone();
    let select_all = params.all;
    let requested_limit = params.limit.clamp(1, 10_000);
    let search_offset = params.offset;

    let parsed = parse_search_query(&query_text);

    let folder = params.folder.or(parsed.folder);
    let subject = params.subject.or(parsed.subject_only.clone());
    let from = params.from.or(parsed.from);
    let to = params.to.or(parsed.to);
    let date_from = params.date_from.or(parsed.date_from);
    let date_to = params.date_to.or(parsed.date_to);
    let has_attachment = params.has_attachment.or(parsed.has_attachment);
    let unread = params.unread.or(parsed.unread);
    let flagged = params.flagged.or(parsed.flagged);

    let data_dir = config.data_dir.clone();
    let user_hash = session.user_hash.clone();
    let text_clone = parsed.text.clone();
    let parsed_subject_clone = subject.clone();
    let folder_clone = folder.clone();
    let from_clone = from.clone();
    let to_clone = to.clone();
    let (sqlite_results, limit) = tokio::task::spawn_blocking(move || {
        let conn = db::pool::open_user_db(&data_dir, &user_hash)
            .map_err(|e| format!("Database error: {e}"))?;
        let effective_limit = if select_all {
            db::messages::count_all_messages(&conn)?.max(1)
        } else {
            requested_limit
        };
        let results = db::messages::search_messages_sqlite(
            &conn,
            &text_clone,
            parsed_subject_clone.as_deref(),
            folder_clone.as_deref(),
            from_clone.as_deref(),
            to_clone.as_deref(),
            date_from,
            date_to,
            has_attachment,
            unread,
            flagged,
            effective_limit,
        )?;
        Ok::<_, String>((results, effective_limit))
    })
    .await
    .map_err(|e| AppError::InternalError(format!("Task error: {e}")))?
    .map_err(|e| AppError::InternalError(format!("Search error: {e}")))?;

    let mut seen = std::collections::HashSet::new();
    let mut results: Vec<SearchResultItem> = sqlite_results
        .into_iter()
        .map(|msg| {
            seen.insert((msg.folder.clone(), msg.uid));
            SearchResultItem {
                uid: msg.uid,
                folder: msg.folder,
                score: 1.0,
                subject: msg.subject,
                from_address: msg.from_address,
                from_name: msg.from_name,
                to_addresses: msg.to_addresses,
                date: msg.date,
                flags: msg.flags,
                has_attachments: msg.has_attachments,
                snippet: msg.snippet,
            }
        })
        .collect();

    if !parsed.text.is_empty()
        && results.len() < limit
        && let Ok(user_index) = search_engine.open_user_index(&session.user_hash)
        && let Ok((tantivy_results, _)) = user_index
            .search(&SearchQuery {
                text: parsed.text,
                subject_only: subject,
                folder,
                from,
                to,
                date_from,
                date_to,
                has_attachment,
                limit,
                offset: if select_all { 0 } else { search_offset },
            })
            .await
    {
        let data_dir2 = config.data_dir.clone();
        let user_hash2 = session.user_hash.clone();
        let remaining = limit.saturating_sub(results.len());
        let tantivy_enriched = tokio::task::spawn_blocking(move || {
            let conn = db::pool::open_user_db(&data_dir2, &user_hash2)
                .map_err(|e| format!("Database error: {e}"))?;
            let mut items = Vec::new();
            for sr in &tantivy_results {
                if items.len() >= remaining {
                    break;
                }
                if let Ok(Some(msg)) = db::messages::get_single_message(&conn, &sr.folder, sr.uid) {
                    let is_unread = !msg.flags.contains("\\Seen");
                    let is_flagged = msg.flags.contains("\\Flagged");
                    if unread.is_some_and(|wanted| wanted != is_unread) {
                        continue;
                    }
                    if flagged.is_some_and(|wanted| wanted != is_flagged) {
                        continue;
                    }
                    items.push((sr.score, sr.snippet.clone(), msg));
                }
            }
            Ok::<_, String>(items)
        })
        .await
        .map_err(|e| AppError::InternalError(format!("Task error: {e}")))?
        .map_err(|e| AppError::InternalError(format!("Database error: {e}")))?;

        for (score, snippet, msg) in tantivy_enriched {
            if seen.contains(&(msg.folder.clone(), msg.uid)) {
                continue;
            }
            seen.insert((msg.folder.clone(), msg.uid));
            results.push(SearchResultItem {
                uid: msg.uid,
                folder: msg.folder,
                score,
                subject: msg.subject,
                from_address: msg.from_address,
                from_name: msg.from_name,
                to_addresses: msg.to_addresses,
                date: msg.date,
                flags: msg.flags,
                has_attachments: msg.has_attachments,
                snippet: if snippet.is_empty() {
                    msg.snippet
                } else {
                    snippet
                },
            });
        }
    }

    let sort_asc = sort_order == "date_asc";
    results.sort_by(|a, b| {
        let da = crate::db::messages::parse_date_to_epoch_public(&a.date);
        let db_val = crate::db::messages::parse_date_to_epoch_public(&b.date);
        if sort_asc {
            da.cmp(&db_val)
        } else {
            db_val.cmp(&da)
        }
    });

    results.truncate(limit);

    Ok(Json(SearchResponse {
        total_count: results.len(),
        results,
        query: query_text,
    })
    .into_response())
}

#[derive(Debug, Deserialize)]
pub struct SavedSearchRequest {
    pub name: String,
    pub query_string: String,
}

pub async fn list_saved_searches(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
) -> Result<Response, AppError> {
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;
    let searches = db::saved_searches::list(&conn).map_err(AppError::InternalError)?;
    Ok(Json(serde_json::json!({ "searches": searches })).into_response())
}

pub async fn create_saved_search(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Json(data): Json<SavedSearchRequest>,
) -> Result<Response, AppError> {
    if data.name.trim().is_empty() || data.query_string.trim().is_empty() {
        return Err(AppError::BadRequest(
            "Saved search name and query are required".to_string(),
        ));
    }
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;
    let saved = db::saved_searches::create(&conn, &data.name, &data.query_string)
        .map_err(AppError::BadRequest)?;
    Ok((axum::http::StatusCode::CREATED, Json(saved)).into_response())
}

pub async fn update_saved_search(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Path(id): Path<i64>,
    Json(data): Json<SavedSearchRequest>,
) -> Result<Response, AppError> {
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;
    match db::saved_searches::update(&conn, id, &data.name, &data.query_string)
        .map_err(AppError::BadRequest)?
    {
        Some(saved) => Ok(Json(saved).into_response()),
        None => Err(AppError::NotFound("Saved search not found".to_string())),
    }
}

pub async fn delete_saved_search(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Path(id): Path<i64>,
) -> Result<Response, AppError> {
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;
    if db::saved_searches::delete(&conn, id).map_err(AppError::InternalError)? {
        Ok(Json(serde_json::json!({ "status": "deleted" })).into_response())
    } else {
        Err(AppError::NotFound("Saved search not found".to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_plain_text() {
        let p = parse_search_query("hello world");
        assert_eq!(p.text, "hello world");
        assert!(p.from.is_none());
        assert!(p.to.is_none());
        assert!(p.subject_only.is_none());
        assert!(p.folder.is_none());
    }

    #[test]
    fn parse_from_operator() {
        let p = parse_search_query("from:alice@example.com hello");
        assert_eq!(p.text, "hello");
        assert_eq!(p.from.as_deref(), Some("alice@example.com"));
    }

    #[test]
    fn parse_to_operator() {
        let p = parse_search_query("to:bob@example.com updates");
        assert_eq!(p.text, "updates");
        assert_eq!(p.to.as_deref(), Some("bob@example.com"));
    }

    #[test]
    fn parse_subject_operator() {
        let p = parse_search_query("subject:meeting notes");
        assert_eq!(p.subject_only.as_deref(), Some("meeting"));
        assert_eq!(p.text, "notes");
    }

    #[test]
    fn parse_subject_quoted() {
        let p = parse_search_query("subject:\"meeting notes\" extra");
        assert_eq!(p.subject_only.as_deref(), Some("meeting notes"));
        assert_eq!(p.text, "extra");
    }

    #[test]
    fn parse_folder_operators() {
        let p = parse_search_query("in:Sent hello");
        assert_eq!(p.folder.as_deref(), Some("Sent"));
        assert_eq!(p.text, "hello");

        let p2 = parse_search_query("folder:INBOX test");
        assert_eq!(p2.folder.as_deref(), Some("INBOX"));
    }

    #[test]
    fn parse_date_operator() {
        let p = parse_search_query("date:2024-01-15 news");
        assert_eq!(p.text, "news");
        assert!(p.date_from.is_some());
        assert!(p.date_to.is_some());

        assert_eq!(p.date_from, Some(1_705_276_800));
        assert_eq!(p.date_to, Some(1_705_276_800 + 86_399));
    }

    #[test]
    fn parse_after_before() {
        let p = parse_search_query("after:2024-01-01 before:2024-06-01 hello");
        assert_eq!(p.text, "hello");

        assert_eq!(p.date_from, Some(1_704_067_200));

        assert_eq!(p.date_to, Some(1_717_200_000));
    }

    #[test]
    fn parse_has_attachment() {
        let p = parse_search_query("has:attachment report");
        assert_eq!(p.text, "report");
        assert_eq!(p.has_attachment, Some(true));
    }

    #[test]
    fn parse_state_operators() {
        let p = parse_search_query("is:unread is:flagged report");
        assert_eq!(p.text, "report");
        assert_eq!(p.unread, Some(true));
        assert_eq!(p.flagged, Some(true));

        let p2 = parse_search_query("is:read is:unflagged");
        assert_eq!(p2.unread, Some(false));
        assert_eq!(p2.flagged, Some(false));
    }
    #[test]
    fn parse_multiple_operators() {
        let p = parse_search_query(
            "from:alice@test.com to:bob@test.com in:INBOX has:attachment meeting",
        );
        assert_eq!(p.from.as_deref(), Some("alice@test.com"));
        assert_eq!(p.to.as_deref(), Some("bob@test.com"));
        assert_eq!(p.folder.as_deref(), Some("INBOX"));
        assert_eq!(p.has_attachment, Some(true));
        assert_eq!(p.text, "meeting");
    }

    #[test]
    fn parse_quoted_from() {
        let p = parse_search_query("from:\"John Doe\" hello");
        assert_eq!(p.from.as_deref(), Some("John Doe"));
        assert_eq!(p.text, "hello");
    }

    #[test]
    fn parse_unknown_operator_kept_as_text() {
        let p = parse_search_query("foo:bar hello");
        assert_eq!(p.text, "foo:bar hello");
    }

    #[test]
    fn parse_empty_value_not_operator() {
        let p = parse_search_query("hello from:");
        assert_eq!(p.text, "hello from:");
        assert!(p.from.is_none());
    }

    #[test]
    fn date_parsing_basic() {
        assert_eq!(parse_date_start("2024-01-01"), Some(1_704_067_200));

        assert_eq!(parse_date_start("not-a-date"), None);
        assert_eq!(parse_date_start("2024-13-01"), None);
    }
}