phrona-api 0.1.1

use std::sync::Arc;

use axum::extract::State;
use axum::response::IntoResponse;
use axum::{Json, http::HeaderMap};
use serde::{Deserialize, Serialize};

use phrona::SearchOptions;
use phrona::models::{Category, ResultItem, TimeRange};

use crate::{AppError, AppResult, AppState, JsonBody};

/// Tavily-compatible request body.
///
/// The Tavily API (https://docs.tavily.com) is the de-facto standard for
/// AI search. Clients such as `tavily-python` can target this server by
/// setting `base_url` to it and calling `/search`.
#[derive(Deserialize)]
pub struct TavilyRequest {
    pub query: String,
    #[serde(default)]
    pub api_key: Option<String>,
    #[serde(default)]
    pub search_depth: Option<String>,
    #[serde(default)]
    pub topic: Option<String>,
    #[serde(default)]
    pub days: Option<u32>,
    #[serde(default)]
    pub max_results: Option<usize>,
    #[serde(default)]
    pub include_images: bool,
    #[serde(default)]
    pub include_answer: bool,
    #[serde(default)]
    pub include_raw_content: bool,
    /// Accepted for compatibility with Tavily clients; only meaningful for
    /// Tavily's image-search endpoint, which this server does not expose.
    #[serde(default)]
    pub include_image_descriptions: bool,
    #[serde(default)]
    pub include_domains: Option<Vec<String>>,
    #[serde(default)]
    pub exclude_domains: Option<Vec<String>>,
}

#[derive(Serialize)]
pub struct TavilyResponse {
    pub query: String,
    pub follow_up_questions: Vec<String>,
    pub response_time: f64,
    pub answer: Option<String>,
    pub images: Option<Vec<String>>,
    pub results: Vec<TavilyResult>,
}

#[derive(Serialize)]
pub struct TavilyResult {
    pub title: String,
    pub url: String,
    pub content: String,
    pub score: f64,
    pub raw_content: Option<String>,
}

fn days_to_range(days: u32) -> Option<TimeRange> {
    Some(match days {
        0..=1 => TimeRange::Day,
        2..=7 => TimeRange::Week,
        8..=30 => TimeRange::Month,
        _ => TimeRange::Year,
    })
}

fn apply_domains(query: &mut String, include: &[String], exclude: &[String]) {
    if !include.is_empty() {
        let sites: Vec<String> = include.iter().map(|d| format!("site:{d}")).collect();
        query.push_str(&format!(" ({})", sites.join(" OR ")));
    }
    for d in exclude {
        query.push_str(&format!(" -site:{d}"));
    }
}

fn to_tavily_result(r: &ResultItem, pos: usize) -> (String, String, String, f64) {
    let score = (1.0 - pos as f64 * 0.05).max(0.05);
    match r {
        ResultItem::Web(w) => (w.title.clone(), w.url.clone(), w.description.clone(), score),
        ResultItem::News(n) => (n.title.clone(), n.url.clone(), n.description.clone(), score),
        ResultItem::Video(v) => (v.title.clone(), v.url.clone(), v.description.clone(), score),
        ResultItem::Image(i) => (i.title.clone(), i.url.clone(), i.source.clone(), score),
        ResultItem::Book(b) => (b.title.clone(), b.url.clone(), b.info.clone(), score),
    }
}

pub async fn search(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    JsonBody(req): JsonBody<TavilyRequest>,
) -> AppResult<impl IntoResponse> {
    // Tavily SDKs pass credentials as `api_key` in the JSON body
    // (langchain-tavily, llama-index) or as headers; both are honored.
    let key = crate::auth_key(&headers, req.api_key.as_deref());
    if !state.authorized(key.as_deref()) {
        return Err(AppError::unauthorized());
    }

    let mut opts = SearchOptions::new(req.query.clone());
    let topic_is_news = matches!(req.topic.as_deref(), Some("news"));
    opts.category = if topic_is_news {
        Category::News
    } else {
        Category::Web
    };
    let depth = req.search_depth.as_deref().unwrap_or("basic");
    // "advanced" is honored by querying every engine in the category;
    // anything else is rejected loudly instead of silently coerced.
    if !matches!(depth, "basic" | "advanced") {
        return Err(AppError::bad_request(format!(
            "invalid search_depth '{depth}', expected 'basic' or 'advanced'"
        )));
    }
    if depth == "basic" {
        let mut engines = match opts.category {
            Category::News => vec!["bing_news".into(), "duckduckgo_news".into()],
            _ => vec!["bing".into(), "duckduckgo".into()],
        };
        if req.include_answer {
            engines.push("grokipedia".into());
        }
        opts.engines = engines;
    }
    if let Some(days) = req.days {
        opts.time_range = days_to_range(days);
    } else if topic_is_news {
        // news topic without an explicit window: last week, like Tavily
        opts.time_range = Some(TimeRange::Week);
    }
    opts.max_results = req.max_results.unwrap_or(5).clamp(1, 20);
    if let Some(include) = &req.include_domains {
        apply_domains(&mut opts.query, include, &[]);
    }
    if let Some(exclude) = &req.exclude_domains {
        apply_domains(&mut opts.query, &[], exclude);
    }

    let started = std::time::Instant::now();
    let resp = state.client.search(opts).await?;
    let response_time = started.elapsed().as_secs_f64();

    let limit = resp.total.min(req.max_results.unwrap_or(5).clamp(1, 20));
    let mut results: Vec<TavilyResult> = Vec::with_capacity(limit);
    let mut images: Vec<String> = Vec::new();
    for (i, r) in resp.results.iter().take(limit).enumerate() {
        let (title, url, content, score) = to_tavily_result(r, i);
        if let ResultItem::Image(img) = r
            && !img.image_url.is_empty()
        {
            images.push(img.image_url.clone());
        }
        results.push(TavilyResult {
            title,
            url,
            content,
            score,
            raw_content: None,
        });
    }

    if req.include_images {
        // Tavily's `images` field lists image results alongside the web
        // hits; run a dedicated image search to populate it honestly.
        let mut img_opts = SearchOptions::new(req.query.clone());
        img_opts.category = Category::Images;
        img_opts.max_results = limit.clamp(1, 8);
        if let Ok(img_resp) = state.client.search(img_opts).await {
            for r in img_resp.results.iter().take(8) {
                if let ResultItem::Image(img) = r
                    && !img.image_url.is_empty()
                {
                    images.push(img.image_url.clone());
                }
            }
        }
    }

    if req.include_raw_content {
        let client = state.client.http();
        let urls: Vec<String> = results.iter().map(|r| r.url.clone()).collect();
        let pages = phrona::extract_many(client, &urls, 8000, Some(&req.query)).await;
        for (r, page) in results.iter_mut().zip(pages) {
            r.raw_content = Some(match page {
                Ok(p) => p.text,
                Err(e) => format!("extract failed: {e}"),
            });
        }
    }

    Ok(Json(TavilyResponse {
        query: req.query.clone(),
        follow_up_questions: Vec::new(),
        response_time,
        answer: req.include_answer.then_some(resp.answer.clone()).flatten(),
        images: (req.include_images && !images.is_empty()).then_some(images),
        results,
    }))
}