search_hub

at 8f8b2d8 Raw

use crate::config::{EngineConfig, Shortcut};
use crate::search_engines::utils::urlencode;
use crate::search_engines::{EngineError, ResultEntry, SearchEngine};
use crate::storage;
use actix_web::{get, http::header::LOCATION, web, HttpRequest, HttpResponse, Responder};
use serde::Serialize;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Instant;
use tera::Tera;
use tracing::{error, info};

use super::{
    interleave, ApiBookmark, ApiExternal, DbPool, SearchApiResponse, SearchApiResult, SearchQuery,
    ServerConfig, ShortcutHelp, USER_AGENT, VERSION,
};

enum ShortcutKind {
    Bang(String),
    At(usize),
}

/// Parse a !bang or @at shortcut from the start of a query.
fn parse_shortcut<'a>(
    q: &'a str,
    shortcuts: &HashMap<String, Shortcut>,
) -> Option<(ShortcutKind, &'a str)> {
    if q.is_empty() {
        return None;
    }
    // Split off the prefix (! or @) and extract the trigger word up to the first space.
    // e.g. "!w Rust" -> prefix='!', trigger="w", rest="Rust"
    let (prefix, body) = if let Some(body) = q.strip_prefix('!') {
        ('!', body)
    } else if let Some(body) = q.strip_prefix('@') {
        ('@', body)
    } else {
        return None;
    };
    let (trigger, rest) = if let Some(space) = body.find(' ') {
        (&body[..space], &body[space + 1..])
    } else {
        (body, "")  // bare shortcut, no query text
    };
    if trigger.is_empty() {
        return None;
    }
    let sc = shortcuts.get(trigger)?;
    match prefix {
        '!' => Some((ShortcutKind::Bang(sc.bang_url.clone()), rest)),
        '@' => Some((ShortcutKind::At(sc.engine_index?), rest)),  // `@` requires a real engine
        _ => unreachable!(),
    }
}

/// Build the sorted list of shortcuts for the help overlay in the web UI.
/// Entries with `engine_index: None` (standalone [[bangs]]) get `has_at: false`.
fn build_shortcuts_help(shortcuts: &HashMap<String, Shortcut>) -> Vec<ShortcutHelp> {
    let mut help: Vec<ShortcutHelp> = shortcuts
        .values()
        .map(|sc| ShortcutHelp {
            trigger: sc.trigger.clone(),
            name: sc.name.clone(),
            has_at: sc.engine_index.is_some(),
        })
        .collect();
    help.sort_by(|a, b| a.trigger.cmp(&b.trigger));
    help
}

fn run_external_engines(
    q: &str,
    user_agent: &str,
    engines: &[EngineConfig],
    engine_filter: Option<usize>,
) -> (Vec<Vec<ResultEntry>>, usize) {
    let mut all_external: Vec<Vec<ResultEntry>> = Vec::new();
    let mut provider_count: usize = 0;

    let client = match reqwest::Client::builder()
        .user_agent(user_agent)
        .build()
        .ok()
    {
        Some(c) => c,
        None => return (all_external, 0),
    };

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .expect("failed to build tokio runtime");

    let snapshot = engines.to_vec();
    let mut handles = Vec::new();

    for i in 0..snapshot.len() {
        if let Some(filter) = engine_filter {
            if i != filter {
                continue;
            }
        }
        let engine_name = snapshot[i].name().to_string();
        let q_owned = q.to_string();
        let client = client.clone();
        let engines = snapshot.clone();

        handles.push(rt.spawn(async move {
            let t0 = Instant::now();
            let timeout_dur = engines[i].timeout();
            let result = tokio::time::timeout(
                timeout_dur,
                engines[i].fetch_results(&q_owned, &client),
            )
            .await;
            let elapsed = t0.elapsed();
            match result {
                Ok(Ok(results)) => (engine_name, Ok(results), elapsed),
                Ok(Err(e)) => (engine_name, Err(e), elapsed),
                Err(_) => (
                    engine_name,
                    Err(EngineError(format!("timed out after {timeout_dur:?}"))),
                    elapsed,
                ),
            }
        }));
    }

    for handle in handles {
        if let Ok((name, result, elapsed)) = rt.block_on(handle) {
            provider_count += 1;
            match result {
                Ok(results) => {
                    info!("external {} ({} results) [{:.2?}]", name, results.len(), elapsed);
                    all_external.push(results);
                }
                Err(e) => {
                    info!("external {} (error) [{:.2?}]: {}", name, elapsed, e);
                }
            }
        }
    }

    (all_external, provider_count)
}

#[get("/")]
async fn index(templates: web::Data<Tera>, cfg: web::Data<ServerConfig>) -> impl Responder {
    info!("serving index page");
    let mut ctx = tera::Context::new();
    ctx.insert("version", VERSION);
    ctx.insert("port", &cfg.port);
    match templates.render("index.html", &ctx) {
        Ok(rendered) => HttpResponse::Ok().content_type("text/html").body(rendered),
        Err(e) => {
            error!("Template error: {}", e);
            HttpResponse::InternalServerError().finish()
        }
    }
}

#[get("/opensearch.xml")]
async fn opensearch(templates: web::Data<Tera>, cfg: web::Data<ServerConfig>) -> impl Responder {
    let mut ctx = tera::Context::new();
    ctx.insert("port", &cfg.port);
    match templates.render("opensearch.xml", &ctx) {
        Ok(xml) => HttpResponse::Ok().content_type("application/opensearchdescription+xml").body(xml),
        Err(e) => {
            error!("Template error: {}", e);
            HttpResponse::InternalServerError().finish()
        }
    }
}

#[get("/search")]
async fn search(
    req: HttpRequest,
    query: web::Query<SearchQuery>,
    templates: web::Data<Tera>,
    db_pool: web::Data<DbPool>,
    engines: web::Data<Arc<RwLock<Vec<EngineConfig>>>>,
    cfg: web::Data<ServerConfig>,
    shortcuts: web::Data<HashMap<String, Shortcut>>,
) -> impl Responder {
    let start = Instant::now();
    let q_orig = query.q.as_deref().unwrap_or("");
    let page = query.page.unwrap_or(1).max(1);
    let page_size = cfg.page_size;

    // Check for !bang / @at prefix before doing any other work.
    // '!' triggers an immediate HTTP redirect; '@' filters external search
    // to a single engine but still shows local bookmark results.
    let (q, engine_filter) =
        match parse_shortcut(q_orig, &shortcuts) {
            Some((ShortcutKind::Bang(url), rest)) => {
                let redirect_url = url.replace("{}", &urlencode(rest));
                info!("bang redirect: {} -> {}", q_orig, redirect_url);
                return HttpResponse::Found()
                    .append_header((LOCATION, redirect_url))
                    .finish();
            }
            Some((ShortcutKind::At(idx), rest)) => {
                info!("at-filtered search: {} -> engine[{}]", q_orig, idx);
                (rest.to_string(), Some(idx))
            }
            None => (q_orig.to_string(), None),
        };

    let has_query = !q.is_empty();
    info!("search request: query=\"{}\" page={}", q, page);

    let total_results = if has_query {
        storage::count_search_bookmarks(&db_pool.conn(), &q).unwrap_or(0)
    } else {
        storage::count_bookmarks(&db_pool.conn()).unwrap_or(0)
    };
    let total_pages = total_results.div_ceil(page_size);

    let bookmarks = if has_query {
        storage::search_bookmarks(&db_pool.conn(), &q, page, page_size).unwrap_or_default()
    } else {
        storage::list_bookmarks(&db_pool.conn(), page, page_size).unwrap_or_default()
    };

    let user_agent = req
        .headers()
        .get("User-Agent")
        .and_then(|v| v.to_str().ok())
        .unwrap_or(USER_AGENT);

    let snapshot = engines.read().unwrap().clone();
    let (all_external, provider_count) = if has_query {
        run_external_engines(&q, user_agent, &snapshot, engine_filter)
    } else {
        (Vec::new(), 0)
    };

    let external_results = interleave(&all_external);

    let page_elapsed = start.elapsed();
    let page_time_ms = format!("{:.1}", page_elapsed.as_secs_f64() * 1000.0);
    info!(
        "search completed: {} bookmark results, {} external providers [{:.2?}]",
        bookmarks.len(),
        provider_count,
        page_elapsed
    );

    let mut external_engines: Vec<String> = external_results
        .iter()
        .map(|r| r.engine.clone())
        .collect();
    external_engines.sort();
    external_engines.dedup();

    let shortcuts_help = build_shortcuts_help(&shortcuts);

    let mut ctx = tera::Context::new();
    ctx.insert("bookmarks", &bookmarks);
    ctx.insert("query", &q);
    ctx.insert("page", &page);
    ctx.insert("total_pages", &total_pages);
    ctx.insert("total_results", &total_results);
    ctx.insert("version", VERSION);
    ctx.insert("page_time_ms", &page_time_ms);
    ctx.insert("external_results", &external_results);
    ctx.insert("external_engines", &external_engines);
    ctx.insert("shortcuts", &shortcuts_help);

    match templates.render("index.html", &ctx) {
        Ok(rendered) => HttpResponse::Ok().content_type("text/html").body(rendered),
        Err(e) => {
            error!("Template error: {}", e);
            HttpResponse::InternalServerError().finish()
        }
    }
}

#[derive(Serialize)]
struct BangApiResponse {
    #[serde(rename = "type")]
    kind: &'static str,
    trigger: String,
    query: String,
    redirect_url: String,
}

#[get("/api/search")]
async fn api_search(
    req: HttpRequest,
    query: web::Query<SearchQuery>,
    cfg: web::Data<ServerConfig>,
    db_pool: web::Data<DbPool>,
    engines: web::Data<Arc<RwLock<Vec<EngineConfig>>>>,
    shortcuts: web::Data<HashMap<String, Shortcut>>,
) -> impl Responder {
    let start = Instant::now();
    let q_orig = query.q.as_deref().unwrap_or("");
    let page = query.page.unwrap_or(1).max(1);
    let page_size = query.page_size.unwrap_or(cfg.page_size).min(100);

    // Same shortcut logic as /search, but the API JSON response for bangs
    // includes the redirect URL instead of performing the redirect server-side.
    let (q, engine_filter) =
        match parse_shortcut(q_orig, &shortcuts) {
            Some((ShortcutKind::Bang(url), rest)) => {
                let redirect_url = url.replace("{}", &urlencode(rest));
                info!("api bang redirect: {} -> {}", q_orig, redirect_url);
                return HttpResponse::Ok().json(BangApiResponse {
                    kind: "bang",
                    trigger: q_orig.to_string(),
                    query: rest.to_string(),
                    redirect_url,
                });
            }
            Some((ShortcutKind::At(idx), rest)) => {
                info!("api at-filtered search: {} -> engine[{}]", q_orig, idx);
                (rest.to_string(), Some(idx))
            }
            None => (q_orig.to_string(), None),
        };

    let has_query = !q.is_empty();

    let total_results = if has_query {
        storage::count_search_bookmarks(&db_pool.conn(), &q).unwrap_or(0)
    } else {
        storage::count_bookmarks(&db_pool.conn()).unwrap_or(0)
    };
    let total_pages = total_results.div_ceil(page_size);

    let bookmarks = if has_query {
        storage::search_bookmarks(&db_pool.conn(), &q, page, page_size).unwrap_or_default()
    } else {
        storage::list_bookmarks(&db_pool.conn(), page, page_size).unwrap_or_default()
    };

    let user_agent = req
        .headers()
        .get("User-Agent")
        .and_then(|v| v.to_str().ok())
        .unwrap_or(USER_AGENT);

    let snapshot = engines.read().unwrap().clone();
    let external_results = if has_query {
        let (all_external, _) = run_external_engines(&q, user_agent, &snapshot, engine_filter);
        interleave(&all_external)
    } else {
        Vec::new()
    };

    let elapsed = start.elapsed();

    let mut results: Vec<SearchApiResult> = Vec::with_capacity(bookmarks.len() + external_results.len());
    for bm in bookmarks {
        results.push(SearchApiResult::Bookmark(ApiBookmark {
            id: bm.id,
            title: bm.title,
            url: bm.url,
            description: bm.description,
            source: bm.source,
            tags: bm.tags,
            created_at: bm.created_at,
        }));
    }
    for ext in external_results {
        results.push(SearchApiResult::External(ApiExternal {
            title: ext.title,
            url: ext.url,
            description: ext.description,
            engine: ext.engine,
        }));
    }

    let response = SearchApiResponse {
        query: q.to_string(),
        page,
        page_size,
        total_results,
        total_pages,
        page_time_ms: elapsed.as_secs_f64() * 1000.0,
        results,
    };

    HttpResponse::Ok().json(response)
}