search_hub

at 6806854 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::web::Bytes;
use actix_web::{get, http::header::LOCATION, web, HttpRequest, HttpResponse, Responder};
use futures::Stream;
use serde::Serialize;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};
use std::time::Instant;
use tera::Tera;
use tokio::sync::mpsc;
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
}

async 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 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(tokio::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)) = handle.await {
            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);
    ctx.insert("tagging_feature", &cfg!(feature = "tagging"));
    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, cfg.tag_weight).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 js_mode = query.js.unwrap_or(false);

    let (all_external, provider_count) = if has_query && !js_mode {
        let snapshot = engines.read().unwrap().clone();
        run_external_engines(&q, user_agent, &snapshot, engine_filter).await
    } 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);
    ctx.insert("tagging_feature", &cfg!(feature = "tagging"));

    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, cfg.tag_weight).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).await;
        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)
}

#[derive(Serialize)]
struct SseBookmarkData {
    results: Vec<SearchApiResult>,
    total: usize,
    page: usize,
    page_size: usize,
    total_pages: usize,
}

#[derive(Serialize)]
struct SseExternalData {
    engine: String,
    results: Vec<SearchApiResult>,
}

fn sse_event(event_type: &str, data: &str) -> Bytes {
    Bytes::from(format!("event: {event_type}\ndata: {data}\n\n"))
}

fn sse_json_event<T: serde::Serialize>(event_type: &str, payload: &T) -> Bytes {
    let data = serde_json::to_string(payload).unwrap_or_else(|_| "{}".into());
    sse_event(event_type, &data)
}

fn stream_from_rx(
    rx: mpsc::Receiver<Result<Bytes, std::io::Error>>,
) -> impl Stream<Item = Result<Bytes, std::io::Error>> {
    futures::stream::unfold(rx, |mut rx| async {
        rx.recv().await.map(|item| (item, rx))
    })
}

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

    let (q, engine_filter) =
        match parse_shortcut(&q_orig, &shortcuts) {
            Some((ShortcutKind::Bang(url), rest)) => {
                let redirect_url = url.replace("{}", &urlencode(rest));
                info!("sse bang redirect: {} -> {}", q_orig, redirect_url);
                let payload = serde_json::json!({
                    "redirect_url": redirect_url,
                    "query": rest,
                    "trigger": q_orig,
                });
                return HttpResponse::Ok()
                    .content_type("text/event-stream")
                    .streaming(futures::stream::once(async move {
                        Ok::<_, std::io::Error>(sse_json_event("bang", &payload))
                    }));
            }
            Some((ShortcutKind::At(idx), rest)) => {
                info!("sse at-filtered search: {} -> engine[{}]", q_orig, idx);
                (rest.to_string(), Some(idx))
            }
            None => (q_orig.clone(), None),
        };

    let has_query = !q.is_empty();

    if !has_query {
        return HttpResponse::Ok()
            .content_type("text/event-stream")
            .streaming(futures::stream::once(async {
                Ok::<_, std::io::Error>(sse_event("done", "{}"))
            }));
    }

    let (tx, rx) = mpsc::channel::<Result<Bytes, std::io::Error>>(64);

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

    let snapshot = engines.read().unwrap().clone();

    tokio::spawn(async move {
        // 1. Bookmark results
        let total_results = storage::count_search_bookmarks(&db_pool.conn(), &q).unwrap_or(0);
        let total_pages = total_results.div_ceil(page_size);

        let bookmarks =
            storage::search_bookmarks(&db_pool.conn(), &q, page, page_size, cfg.tag_weight).unwrap_or_default();

        let api_bookmarks: Vec<SearchApiResult> = bookmarks
            .into_iter()
            .map(|bm| {
                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,
                })
            })
            .collect();

        let bookmark_data = SseBookmarkData {
            results: api_bookmarks,
            total: total_results,
            page,
            page_size,
            total_pages,
        };

        if tx
            .send(Ok(sse_json_event("bookmarks", &bookmark_data)))
            .await
            .is_err()
        {
            return;
        }

        // 2. External engines
        let mut spawn_count = 0;
        for i in 0..snapshot.len() {
            if let Some(filter) = engine_filter {
                if i != filter {
                    continue;
                }
            }
            spawn_count += 1;
        }

        if spawn_count > 0 {
            let client = match reqwest::Client::builder()
                .user_agent(&user_agent)
                .build()
            {
                Ok(c) => c,
                Err(_) => {
                    let _ = tx.send(Ok(sse_event("done", "{}"))).await;
                    return;
                }
            };

            let remaining = Arc::new(AtomicUsize::new(spawn_count));

            for i in 0..snapshot.len() {
                if let Some(filter) = engine_filter {
                    if i != filter {
                        continue;
                    }
                }

                let tx = tx.clone();
                let remaining = remaining.clone();
                let client = client.clone();
                let engine = snapshot[i].clone();
                let q = q.clone();

                tokio::spawn(async move {
                    let t0 = Instant::now();
                    let timeout_dur = engine.timeout();
                    let result = tokio::time::timeout(
                        timeout_dur,
                        engine.fetch_results(&q, &client),
                    )
                    .await;
                    let elapsed = t0.elapsed();

                    match result {
                        Ok(Ok(results)) => {
                            info!(
                                "external {} ({} results) [{:.2?}]",
                                engine.name(),
                                results.len(),
                                elapsed
                            );
                            if !results.is_empty() {
                                let api_results: Vec<SearchApiResult> = results
                                    .into_iter()
                                    .map(|r| {
                                        SearchApiResult::External(ApiExternal {
                                            title: r.title,
                                            url: r.url,
                                            description: r.description,
                                            engine: r.engine,
                                        })
                                    })
                                    .collect();

                                let ext_data = SseExternalData {
                                    engine: engine.name().to_string(),
                                    results: api_results,
                                };
                                let _ = tx
                                    .send(Ok(sse_json_event("external", &ext_data)))
                                    .await;
                            }
                        }
                        Ok(Err(e)) => {
                            info!("external {} (error) [{:.2?}]: {}", engine.name(), elapsed, e);
                        }
                        Err(_) => {
                            info!("external {} (timeout) [{:.2?}]", engine.name(), elapsed);
                        }
                    }

                    if remaining.fetch_sub(1, Ordering::SeqCst) == 1 {
                        let _ = tx.send(Ok(sse_event("done", "{}"))).await;
                    }
                });
            }
        } else {
            let _ = tx.send(Ok(sse_event("done", "{}"))).await;
        }
    });

    let stream = stream_from_rx(rx);

    HttpResponse::Ok()
        .content_type("text/event-stream")
        .streaming(stream)
}