momentry_core/src/api/search.rs

//! Smart Search API
//! Implements the 5W1H search capability using semantic vectors.

use axum::{extract::State, http::StatusCode, response::Json, routing::post, Router};
use serde::{Deserialize, Serialize};
use serde_json;

use crate::core::embedding::Embedder;

// --- Request / Response Structures ---

#[derive(Debug, Deserialize)]
pub struct SmartSearchRequest {
    pub uuid: String,
    pub query: String,
    pub page: Option<usize>,
    pub page_size: Option<usize>,
    pub limit: Option<usize>,
}

#[derive(Debug, Serialize)]
pub struct SearchResult {
    pub id: i32,
    pub parent_id: i32,
    pub scene_order: Option<i32>,

    // Primary: frame-accurate position (authoritative unit)
    pub start_frame: i64,
    pub end_frame: i64,
    pub fps: f64,

    // Reference: time derived from frames (subject to FPS variation, not precise)
    pub start_time: f64,
    pub end_time: f64,

    pub raw_text: Option<String>, // Text content of the child chunk
    pub summary: Option<String>,  // Summary from parent context
    pub metadata: Option<serde_json::Value>,
    pub similarity: Option<f64>,
}

#[derive(Debug, Serialize)]
pub struct SmartSearchResponse {
    pub query: String,
    pub results: Vec<SearchResult>,
    pub page: usize,
    pub page_size: usize,
    pub strategy: String,
}

// --- API Handler ---

pub async fn smart_search(
    State(state): State<crate::api::server::AppState>,
    Json(req): Json<SmartSearchRequest>,
) -> Result<Json<SmartSearchResponse>, (StatusCode, Json<serde_json::Value>)> {
    let db = &state.db;
    let page = req.page.unwrap_or(1).max(1);
    // Backward compat: if old `limit` sent without `page_size`, use limit as page_size
    let page_size = if req.page_size.is_some() {
        req.page_size.unwrap()
    } else if req.limit.is_some() && req.page.is_none() {
        req.limit.unwrap()
    } else {
        5
    }
    .max(1);
    let hard_limit = req.limit.unwrap_or(usize::MAX);
    let limit = hard_limit.min(page_size);

    // 1. Generate Embedding using EmbeddingGemma via MOMENTRY_EMBED_URL
    let embedder = Embedder::new("embeddinggemma-300m".to_string());
    let embedding = embedder.embed_query(&req.query).await.map_err(
        |e| -> (StatusCode, Json<serde_json::Value>) {
            tracing::error!("Embedding failed: {}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": e.to_string() })),
            )
        },
    )?;

    // 2. Search Database (Drill-Down: Find Parents First)
    let db_parents: Vec<crate::core::db::postgres_db::SemanticSearchResult> = db
        .search_parent_chunks_semantic(&req.uuid, &embedding, limit)
        .await
        .map_err(
            |e: anyhow::Error| -> (StatusCode, Json<serde_json::Value>) {
                tracing::error!("DB search failed: {}", e);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(serde_json::json!({ "error": e.to_string() })),
                )
            },
        )?;

    if db_parents.is_empty() {
        return Ok(Json(SmartSearchResponse {
            query: req.query,
            results: vec![],
            page,
            page_size,
            strategy: "semantic_vector_search".to_string(),
        }));
    }

    // Collect Parent IDs
    let parent_ids: Vec<i32> = db_parents.iter().map(|p| p.id).collect();

    // 3. Fetch Children for these Parents (Drill Down)
    // We fetch all children for these parents (limit can be adjusted)
    let children: Vec<crate::core::db::postgres_db::ChildChunkResult> = db
        .get_children_for_parents(&parent_ids, 10) // Fetch top 10 children per parent
        .await
        .map_err(
            |e: anyhow::Error| -> (StatusCode, Json<serde_json::Value>) {
                tracing::error!("Fetching children failed: {}", e);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(serde_json::json!({ "error": e.to_string() })),
                )
            },
        )?;

    // 4. Map Parents to a lookup table
    let parent_map: std::collections::HashMap<
        i32,
        &crate::core::db::postgres_db::SemanticSearchResult,
    > = db_parents.iter().map(|p| (p.id, p)).collect();

    // Map Children to API response struct
    let results: Vec<SearchResult> = children
        .into_iter()
        .map(|c| {
            let parent = parent_map.get(&c.parent_id);
            SearchResult {
                id: c.id,
                parent_id: c.parent_id,
                scene_order: parent.map(|p| p.scene_order),

                start_frame: c.start_frame,
                end_frame: c.end_frame,
                fps: c.fps,

                start_time: c.start_time,
                end_time: c.end_time,
                raw_text: Some(c.raw_text),
                summary: parent.map(|p| p.summary.clone()),
                metadata: parent.map(|p| p.metadata.clone()),
                similarity: parent.and_then(|p| p.similarity),
            }
        })
        .collect();

    // 6. Sort results by similarity (descending)
    // Since all children of a parent have the same parent similarity, this groups relevant chunks together
    let mut results = results;
    results.sort_by(|a, b| {
        b.similarity
            .partial_cmp(&a.similarity)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // 7. Limit the final results (optional, but good for API consistency)
    let truncate_limit = hard_limit.min(page_size * 5); // Allow more children per parent context
    results.truncate(truncate_limit);

    // 8. Format Response
    let response = SmartSearchResponse {
        query: req.query,
        results,
        page,
        page_size,
        strategy: "drill_down_semantic_search".to_string(),
    };

    Ok(Json(response))
}

// --- Router Setup ---

pub fn search_routes() -> Router<crate::api::server::AppState> {
    Router::new().route("/api/v1/search/smart", post(smart_search))
}