use axum::{ body::Body, extract::{Path, Query, State}, http::{header, StatusCode}, response::{IntoResponse, Json}, routing::{get, post}, Router, }; use serde::{Deserialize, Serialize}; use std::process::Command; use crate::core::db::{Database, PostgresDb}; #[derive(Debug, Deserialize)] pub struct CreateIdentityRequest { pub face_json_path: String, pub identity_name: String, pub schema: Option, } #[derive(Debug, Serialize)] pub struct CreateIdentityResponse { pub success: bool, pub message: String, pub identity_uuid: Option, pub identity_name: String, pub total_vectors: Option, pub angle_coverage: Option>, pub quality_avg: Option, } pub fn identity_routes() -> Router { Router::new() .route("/api/v1/identities", get(list_identities)) .route("/api/v1/identity", post(create_identity)) .route("/api/v1/faces/candidates", get(list_face_candidates)) .route("/api/v1/traces/unassigned", get(list_unassigned_traces)) } /// Register a Global Identity from face.json with multi-angle reference vectors. /// Calls select_face_reference_vectors_v2.py for automatic reference selection. async fn create_identity( State(_state): State, Json(req): Json, ) -> Result, (StatusCode, String)> { let schema = req.schema.unwrap_or("dev".to_string()); let python_path = std::env::var("MOMENTRY_PYTHON_PATH").unwrap_or("/opt/homebrew/bin/python3.11".to_string()); let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR").unwrap_or_else(|_| { let mut path = std::env::current_dir().unwrap_or_default(); path.push("scripts"); path.to_string_lossy().to_string() }); let script_path = format!("{}/select_face_reference_vectors_v2.py", scripts_dir); tracing::info!( "Registering identity '{}' from face.json: {}", req.identity_name, req.face_json_path ); let output = Command::new(&python_path) .arg(&script_path) .arg("--face-json") .arg(&req.face_json_path) .arg("--identity-name") .arg(&req.identity_name) .arg("--register") .arg("--schema") .arg(&schema) .output() .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to execute script: {}", e), ) })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Script failed: {}", stderr), )); } let db = PostgresDb::init().await.map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {}", e), ) })?; let id_table = crate::core::db::schema::table_name("identities"); let query = format!( "SELECT uuid, reference_data->'total_references' as total, reference_data->'angles_covered' as angles, reference_data->'quality_avg' as quality FROM {} WHERE name = $1 ORDER BY created_at DESC LIMIT 1", id_table ); let row: Option<(String, Option, Option>, Option)> = sqlx::query_as(&query) .bind(&req.identity_name) .fetch_optional(db.pool()) .await .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), ) })?; match row { Some((uuid, total, angles, quality)) => Ok(Json(CreateIdentityResponse { success: true, message: format!( "Successfully registered identity '{}' with {} reference vectors", req.identity_name, total.unwrap_or(0) ), identity_uuid: Some(uuid), identity_name: req.identity_name, total_vectors: total, angle_coverage: angles, quality_avg: quality, })), None => Ok(Json(CreateIdentityResponse { success: true, message: format!( "Identity '{}' registered, but details not found", req.identity_name ), identity_uuid: None, identity_name: req.identity_name, total_vectors: None, angle_coverage: None, quality_avg: None, })), } } /// List all global identities async fn list_identities( State(_state): State, Query(query): Query, ) -> Result, (StatusCode, String)> { let db = match PostgresDb::init().await { Ok(db) => db, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {}", e), )) } }; let page = query.page.unwrap_or(1); let page_size = query.page_size.unwrap_or(20); let offset = ((page - 1) as i64) * (page_size as i64); let id_table = crate::core::db::schema::table_name("identities"); let total: i64 = sqlx::query_scalar(&format!( "SELECT COUNT(*) FROM {} WHERE status IS NULL OR status != 'merged'", id_table )) .fetch_one(db.pool()) .await .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e), ) })?; let sql = format!( r#"SELECT i.id::int, i.uuid, i.name, i.metadata, i.status, i.starred, COALESCE( jsonb_agg(jsonb_build_object( 'file_uuid', fi.file_uuid, 'confidence', fi.confidence, 'source', fi.metadata->>'source' ) ORDER BY fi.created_at DESC), '[]'::jsonb ) as file_bindings FROM {} i LEFT JOIN {} fi ON i.id = fi.identity_id WHERE i.status IS NULL OR i.status != 'merged' GROUP BY i.id, i.uuid, i.name, i.metadata, i.status, i.starred ORDER BY i.id DESC LIMIT $1 OFFSET $2"#, id_table, crate::core::db::schema::table_name("file_identities") ); let rows: Vec<( i32, uuid::Uuid, String, Option, Option, Option, serde_json::Value, )> = match sqlx::query_as(&sql) .bind(page_size as i64) .bind(offset) .fetch_all(db.pool()) .await { Ok(rows) => rows, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } }; let identities: Vec = rows .into_iter() .map(|r| { let file_bindings: Vec = r.6.as_array() .map(|arr| { arr.iter() .filter_map(|v| serde_json::from_value(v.clone()).ok()) .collect() }) .unwrap_or_default(); let file_uuids: Vec = file_bindings .iter() .map(|fb| fb.file_uuid.clone()) .collect(); IdentityResponse { id: r.0, identity_uuid: r.1.to_string().replace('-', ""), name: r.2, metadata: r.3, status: r.4, starred: r.5.unwrap_or(false), file_uuids, file_bindings: Some(file_bindings), } }) .collect(); let identities_table = crate::core::db::schema::table_name("identities"); let total_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}", identities_table)) .fetch_one(db.pool()) .await .unwrap_or(0); let tmdb_identities: i64 = sqlx::query_scalar(&format!( "SELECT COUNT(*) FROM {} WHERE source = 'tmdb'", identities_table )) .fetch_one(db.pool()) .await .unwrap_or(0); let auto_identities: i64 = sqlx::query_scalar(&format!( "SELECT COUNT(*) FROM {} WHERE file_uuid IS NOT NULL", crate::core::db::schema::table_name("strangers") )) .fetch_one(db.pool()) .await .unwrap_or(0); Ok(Json(IdentityListResponse { identities, count: total, page, page_size, total_identities, tmdb_identities, auto_identities, })) } #[derive(Debug, Deserialize)] pub struct ListIdentitiesQuery { pub page: Option, pub page_size: Option, } #[derive(Debug, Deserialize)] pub struct FaceCandidatesQuery { pub file_uuid: Option, pub min_confidence: Option, pub page: Option, pub page_size: Option, pub limit: Option, } #[derive(Debug, Serialize)] pub struct FaceCandidate { pub id: i32, pub face_id: Option, pub file_uuid: String, pub frame_number: i64, pub confidence: f32, pub bbox: Option, pub attributes: Option, } #[derive(Debug, Serialize)] pub struct FaceCandidatesResponse { pub candidates: Vec, pub total: i64, pub page: usize, pub page_size: usize, } #[derive(Debug, Serialize, Deserialize)] pub struct FileBinding { pub file_uuid: String, pub confidence: f64, pub source: Option, } #[derive(Debug, Serialize)] pub struct IdentityResponse { pub id: i32, pub identity_uuid: String, pub name: String, pub metadata: Option, pub status: Option, pub starred: bool, pub file_uuids: Vec, pub file_bindings: Option>, } #[derive(Debug, Serialize)] pub struct IdentityListResponse { pub identities: Vec, pub count: i64, pub page: usize, pub page_size: usize, pub total_identities: i64, pub tmdb_identities: i64, pub auto_identities: i64, } async fn list_face_candidates( Query(query): Query, ) -> Result, (StatusCode, String)> { let page = query.page.unwrap_or(1); let page_size = std::cmp::min(query.page_size.unwrap_or(15), 100); let offset = (page - 1) * page_size; let min_confidence = query.min_confidence.unwrap_or(0.5); // Query Qdrant _faces for unbound faces (identity_id IS NULL) let qdrant = crate::core::db::qdrant_db::QdrantDb::new(); let mut filter_must = vec![ serde_json::json!({"is_null": {"key": "identity_id"}}), serde_json::json!({"key": "confidence", "range": {"gte": min_confidence}}), ]; if let Some(ref file_uuid) = query.file_uuid { filter_must.push(serde_json::json!({"key": "file_uuid", "match": {"value": file_uuid}})); } let scroll_filter = serde_json::json!({"must": filter_must}); let all_points = qdrant .scroll_all_points("_faces", scroll_filter, 1000) .await .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Qdrant scroll failed: {}", e), ) })?; let total = all_points.len() as i64; // Sort by confidence DESC then paginate let mut sorted: Vec<&serde_json::Value> = all_points.iter().collect(); sorted.sort_by(|a, b| { let ca = a["payload"]["confidence"].as_f64().unwrap_or(0.0); let cb = b["payload"]["confidence"].as_f64().unwrap_or(0.0); cb.partial_cmp(&ca).unwrap_or(std::cmp::Ordering::Equal) }); let paginated: Vec<&&serde_json::Value> = sorted.iter().skip(offset).take(page_size).collect(); let candidates: Vec = paginated .into_iter() .map(|p| { let payload = &p["payload"]; let point_id = p["id"].as_u64().unwrap_or(0); FaceCandidate { id: point_id as i32, face_id: Some(format!("{:x}", point_id)), file_uuid: payload["file_uuid"].as_str().unwrap_or("").to_string(), frame_number: payload["frame"].as_i64().unwrap_or(0), confidence: payload["confidence"].as_f64().unwrap_or(0.0) as f32, bbox: payload.get("bbox").cloned(), attributes: None, } }) .collect(); Ok(Json(FaceCandidatesResponse { candidates, total, page, page_size, })) } #[derive(Debug, Deserialize)] pub struct UnassignedTracesQuery { pub file_uuid: Option, pub page: Option, pub page_size: Option, } #[derive(Debug, Serialize)] pub struct UnassignedTrace { pub trace_id: i32, pub file_uuid: String, pub frame_count: i64, pub start_frame: i64, pub end_frame: i64, pub best_face_id: i32, pub best_face_frame: i64, pub best_face_confidence: f64, pub best_face_bbox: Option, } #[derive(Debug, Serialize)] pub struct UnassignedTracesResponse { pub traces: Vec, pub total: i64, pub page: usize, pub page_size: usize, } /// List unassigned traces (identity_id IS NULL, grouped by trace_id) async fn list_unassigned_traces( Query(query): Query, ) -> Result, (StatusCode, String)> { let page = query.page.unwrap_or(1); let page_size = std::cmp::min(query.page_size.unwrap_or(20), 100); let offset = (page - 1) * page_size; // Query Qdrant _faces for unbound traces (identity_id IS NULL, trace_id > 0) let qdrant = crate::core::db::qdrant_db::QdrantDb::new(); let mut filter_must: Vec = vec![ serde_json::json!({"is_null": {"key": "identity_id"}}), serde_json::json!({"key": "trace_id", "range": {"gt": 0}}), ]; if let Some(ref file_uuid) = query.file_uuid { filter_must.push(serde_json::json!({"key": "file_uuid", "match": {"value": file_uuid}})); } let scroll_filter = serde_json::json!({"must": filter_must}); let all_points = qdrant .scroll_all_points("_faces", scroll_filter, 1000) .await .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Qdrant scroll failed: {}", e), ) })?; // Group by (file_uuid, trace_id) and aggregate use std::collections::BTreeMap; #[derive(Default)] struct TraceAgg { frame_count: i64, start_frame: i64, end_frame: i64, best_confidence: f64, best_point_id: i64, best_frame: i64, best_bbox: Option, } let mut trace_map: BTreeMap<(String, i32), TraceAgg> = BTreeMap::new(); for point in &all_points { let payload = &point["payload"]; let file_uuid = match payload["file_uuid"].as_str() { Some(f) => f.to_string(), None => continue, }; let trace_id = payload["trace_id"].as_i64().unwrap_or(0) as i32; if trace_id <= 0 { continue; } let frame = payload["frame"].as_i64().unwrap_or(0); let confidence = payload["confidence"].as_f64().unwrap_or(0.0); let point_id = point["id"].as_i64().unwrap_or(0); let entry = trace_map.entry((file_uuid, trace_id)).or_default(); entry.frame_count += 1; if frame < entry.start_frame || entry.start_frame == 0 { entry.start_frame = frame; } if frame > entry.end_frame { entry.end_frame = frame; } if confidence > entry.best_confidence { entry.best_confidence = confidence; entry.best_point_id = point_id; entry.best_frame = frame; entry.best_bbox = payload.get("bbox").cloned(); } } let total = trace_map.len() as i64; // Sort by frame_count DESC, paginate let mut sorted_traces: Vec<((String, i32), TraceAgg)> = trace_map.into_iter().collect(); sorted_traces.sort_by(|a, b| b.1.frame_count.cmp(&a.1.frame_count)); let paginated: Vec<_> = sorted_traces .into_iter() .skip(offset) .take(page_size) .collect(); let traces: Vec = paginated .into_iter() .map(|((file_uuid, trace_id), agg)| UnassignedTrace { trace_id, file_uuid, frame_count: agg.frame_count, start_frame: agg.start_frame, end_frame: agg.end_frame, best_face_id: agg.best_point_id as i32, best_face_frame: agg.best_frame, best_face_confidence: agg.best_confidence, best_face_bbox: agg.best_bbox, }) .collect(); Ok(Json(UnassignedTracesResponse { traces, total, page, page_size, })) }