use axum::{ extract::{Query, State}, http::StatusCode, response::Json, routing::{get, post}, Router, }; use serde::{Deserialize, Serialize}; use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; use crate::core::db::{Database, PostgresDb, QdrantDb, RedisClient, VideoRecord}; use crate::{Embedder, FileManager}; #[derive(Clone)] struct AppState { embedder: Arc, #[allow(dead_code)] embedder_model: String, } #[derive(Debug, Deserialize)] struct RegisterRequest { path: String, } #[derive(Debug, Serialize)] struct RegisterResponse { uuid: String, video_id: i64, file_name: String, duration: f64, width: u32, height: u32, } #[derive(Debug, Deserialize)] struct SearchRequest { query: String, limit: Option, uuid: Option, } #[derive(Debug, Serialize)] struct SearchResult { uuid: String, chunk_id: String, chunk_type: String, start_time: f64, end_time: f64, text: String, score: f32, } #[derive(Debug, Serialize)] struct SearchResponse { results: Vec, query: String, } #[derive(Debug, Deserialize)] struct LookupQuery { path: Option, uuid: Option, } #[derive(Debug, Serialize)] struct LookupResponse { uuid: String, file_path: Option, file_name: Option, duration: Option, } #[derive(Debug, Serialize)] struct VideoInfoResponse { uuid: String, file_path: String, file_name: String, duration: f64, width: u32, height: u32, } #[derive(Debug, Serialize)] struct VideosResponse { videos: Vec, } async fn register( State(_state): State, Json(req): Json, ) -> Result, StatusCode> { let path = req.path; let uuid = crate::uuid::compute_uuid_from_path(&path); let probe_result = crate::core::probe::probe_video(&path).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let duration = probe_result .format .duration .as_ref() .and_then(|s: &String| s.parse::().ok()) .unwrap_or(0.0); let mut width = 0u32; let mut height = 0u32; for stream in &probe_result.streams { if stream.codec_type.as_deref() == Some("video") { width = stream.width.unwrap_or(0); height = stream.height.unwrap_or(0); } } let file_manager = FileManager::new(std::path::PathBuf::from(".")); let json_str = serde_json::to_string(&probe_result).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let _json_path = file_manager .save_json(&uuid, "probe", &json_str) .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let db = PostgresDb::init() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let file_path = Path::new(&path) .canonicalize() .map(|p| p.to_string_lossy().to_string()) .unwrap_or_else(|_| path.clone()); let file_name = Path::new(&path) .file_name() .map(|n| n.to_string_lossy().to_string()) .unwrap_or_default(); let record = VideoRecord { id: 0, uuid: uuid.clone(), file_path, file_name: file_name.clone(), duration, width, height, fps: 0.0, probe_json: Some(json_str), storage: Default::default(), created_at: String::new(), }; let video_id = db .register_video(&record) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(Json(RegisterResponse { uuid, video_id, file_name, duration, width, height, })) } async fn search( State(state): State, Json(req): Json, ) -> Result, StatusCode> { let limit = req.limit.unwrap_or(10); let query_vector = state.embedder.embed_query(&req.query).await.map_err(|e| { tracing::error!("Failed to embed query: {}", e); StatusCode::INTERNAL_SERVER_ERROR })?; let qdrant: QdrantDb = QdrantDb::init() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let pg: PostgresDb = PostgresDb::init() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let search_results = if let Some(uuid) = &req.uuid { let query_f64: Vec = query_vector.iter().map(|&x| x as f64).collect(); qdrant .search_in_uuid(&query_f64, uuid, limit) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? } else { qdrant .search(&query_vector, limit) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? }; let mut results = Vec::new(); for r in search_results { let chunks = pg.get_chunks_by_uuid(&r.chunk_id).await.unwrap_or_default(); for chunk in chunks { let text = chunk .content .get("text") .and_then(|v: &serde_json::Value| v.as_str()) .unwrap_or("") .to_string(); results.push(SearchResult { uuid: chunk.uuid.clone(), chunk_id: chunk.chunk_id.clone(), chunk_type: chunk.chunk_type.as_str().to_string(), start_time: chunk.start_time, end_time: chunk.end_time, text, score: r.score, }); } } Ok(Json(SearchResponse { results, query: req.query, })) } async fn lookup(Query(query): Query) -> Result, StatusCode> { let db: PostgresDb = PostgresDb::init() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; if let Some(path) = query.path { let uuid = crate::uuid::compute_uuid_from_path(&path); return Ok(Json(LookupResponse { uuid, file_path: None, file_name: None, duration: None, })); } if let Some(uuid) = query.uuid { let video = db .get_video_by_uuid(&uuid) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; if let Some(v) = video { return Ok(Json(LookupResponse { uuid: v.uuid, file_path: Some(v.file_path), file_name: Some(v.file_name), duration: Some(v.duration), })); } } Err(StatusCode::NOT_FOUND) } pub async fn start_server(host: &str, port: u16) -> anyhow::Result<()> { let embedder = Arc::new(Embedder::new("nomic-embed-text:v1.5".to_string())); let state = AppState { embedder, embedder_model: "nomic-embed-text:v1.5".to_string(), }; let app = Router::new() .route("/api/v1/register", post(register)) .route("/api/v1/search", post(search)) .route("/api/v1/lookup", get(lookup)) .route("/api/v1/videos", get(list_videos)) .route("/api/v1/progress/:uuid", get(get_progress)) .with_state(state); let addr = SocketAddr::new(host.parse().unwrap(), port); tracing::info!("Starting API server at http://{}", addr); let listener = tokio::net::TcpListener::bind(addr).await?; axum::serve(listener, app).await?; Ok(()) } async fn list_videos() -> Result, StatusCode> { let db: PostgresDb = PostgresDb::init() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let videos = db .list_videos() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let video_infos: Vec = videos .into_iter() .map(|v| VideoInfoResponse { uuid: v.uuid, file_path: v.file_path, file_name: v.file_name, duration: v.duration, width: v.width, height: v.height, }) .collect(); Ok(Json(VideosResponse { videos: video_infos, })) } #[derive(Debug, Serialize)] struct ProgressResponse { uuid: String, processors: Vec, } #[derive(Debug, Serialize)] struct ProcessorProgressInfo { name: String, status: String, current: u32, total: u32, message: String, } async fn get_progress( axum::extract::Path(uuid): axum::extract::Path, ) -> Result, StatusCode> { let redis = RedisClient::new().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let mut conn = redis .get_conn() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let processor_names = ["asr", "cut", "asrx", "yolo", "ocr", "face", "pose"]; let mut processors = Vec::new(); for name in processor_names { let key = format!("momentry:job:{}:processor:{}", uuid, name); let status: String = redis::cmd("HGET") .arg(&key) .arg("status") .query_async(&mut conn) .await .unwrap_or_else(|_| "pending".to_string()); let current: u32 = redis::cmd("HGET") .arg(&key) .arg("current") .query_async(&mut conn) .await .unwrap_or_else(|_| "0".to_string()) .parse() .unwrap_or(0); let total: u32 = redis::cmd("HGET") .arg(&key) .arg("total") .query_async(&mut conn) .await .unwrap_or_else(|_| "0".to_string()) .parse() .unwrap_or(0); let message: String = redis::cmd("HGET") .arg(&key) .arg("message") .query_async(&mut conn) .await .unwrap_or_else(|_| "".to_string()); processors.push(ProcessorProgressInfo { name: name.to_string(), status, current, total, message, }); } Ok(Json(ProgressResponse { uuid, processors })) }